From 41bb770808dd84750cda6ce8a01197d207565ae2 Mon Sep 17 00:00:00 2001 From: Codex Date: Sat, 26 Sep 2026 03:35:29 +0800 Subject: [PATCH 1/2] fix(openai): bind credentials and WS continuations to owners --- .../handler/openai_gateway_handler.go | 32 ++ .../handler/openai_gateway_handler_test.go | 2 + .../handler/openai_ws_response_owner_test.go | 267 +++++++++++++ backend/internal/repository/account_repo.go | 25 +- .../account_repo_openai_identity_test.go | 44 +++ backend/internal/service/oauth_refresh_api.go | 56 ++- .../service/oauth_refresh_api_test.go | 25 ++ .../service/openai_identity_mutation_test.go | 373 ++++++++++++++++++ .../internal/service/openai_token_provider.go | 86 +++- .../internal/service/openai_v183_compat.go | 10 + 10 files changed, 903 insertions(+), 17 deletions(-) create mode 100644 backend/internal/handler/openai_ws_response_owner_test.go create mode 100644 backend/internal/repository/account_repo_openai_identity_test.go create mode 100644 backend/internal/service/openai_identity_mutation_test.go diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 25fd898bd827..9b5f6b2857c8 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -2647,6 +2647,30 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, service.OpenAIContinuationAnchorValidationMessage) return } + responseOwnerGroupID := int64(0) + if apiKey.GroupID != nil { + responseOwnerGroupID = *apiKey.GroupID + } + const responseOwnerRejection = "previous_response_id is not available for this user" + // Upstream affinity is not downstream authorization. Share the HTTP + // ownership contract across first frames and every later response.create. + ownsPreviousResponse := func(anchor string) bool { + if anchor == "" { + return true + } + owned, ownerErr := h.gatewayService.ValidateOpenAIHTTPResponseOwner( + ctx, responseOwnerGroupID, anchor, subject.UserID, apiKey.ID, + ) + if ownerErr != nil { + reqLog.Warn("openai.websocket_previous_response_owner_lookup_failed") + } + return ownerErr == nil && owned + } + if !ownsPreviousResponse(previousResponseID) { + closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, responseOwnerRejection) + return + } + service.SetOpenAIHTTPResponseOwner(c, subject.UserID, apiKey.ID) previousResponseIDKind := service.ClassifyOpenAIPreviousResponseIDKind(previousResponseID) reqLog = reqLog.With( zap.Bool("ws_ingress", true), @@ -3108,6 +3132,13 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { if !gjson.ValidBytes(payload) { return newOpenAIWSLocalTurnCloseError(coderws.StatusPolicyViolation, "invalid websocket request payload", errors.New("invalid json")) } + anchor, anchorErr := service.ParseOpenAIContinuationAnchor(payload) + if anchorErr != nil { + return newOpenAIWSLocalTurnCloseError(coderws.StatusPolicyViolation, service.OpenAIContinuationAnchorValidationMessage, anchorErr) + } + if !ownsPreviousResponse(anchor) { + return newOpenAIWSLocalTurnCloseError(coderws.StatusPolicyViolation, responseOwnerRejection, nil) + } model := strings.TrimSpace(originalModel) if model == "" { model = strings.TrimSpace(gjson.GetBytes(payload, "model").String()) @@ -3262,6 +3293,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { if result == nil { return } + h.gatewayService.BindOpenAIWSResponseOwner(ctx, c, account, result) result.BillingModel = openAIWSTurnBillingModel(result, turnMapping, turnRequestedModel, turnUpstreamModel) reqLog.Debug("openai.websocket_turn_billing", zap.Int("turn", turn), diff --git a/backend/internal/handler/openai_gateway_handler_test.go b/backend/internal/handler/openai_gateway_handler_test.go index bb6f51e39c01..4e35b84784cc 100644 --- a/backend/internal/handler/openai_gateway_handler_test.go +++ b/backend/internal/handler/openai_gateway_handler_test.go @@ -1276,6 +1276,8 @@ func TestOpenAIResponsesWebSocket_PreviousResponseIDKindLoggedBeforeAcquireFailu }, } h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache) + // This case exercises user-slot failure after an authorized continuation. + require.NoError(t, h.gatewayService.BindOpenAIHTTPResponseOwner(context.Background(), 2, "resp_prev_123", 1, 101)) wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1}) defer wsServer.Close() diff --git a/backend/internal/handler/openai_ws_response_owner_test.go b/backend/internal/handler/openai_ws_response_owner_test.go new file mode 100644 index 000000000000..15545cab1c1e --- /dev/null +++ b/backend/internal/handler/openai_ws_response_owner_test.go @@ -0,0 +1,267 @@ +package handler + +import ( + "context" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "sync" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" + "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + coderws "github.com/coder/websocket" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +const wsResponseOwnerFixtureGroupID int64 = 61201 + +// Every upstream response is synthetic. Native modes use a loopback WS fake; +// HTTP bridge mode terminates in Do without opening an upstream connection. +type wsResponseOwnerUpstream struct { + service.HTTPUpstream + mu sync.Mutex + requests [][]byte +} + +func (u *wsResponseOwnerUpstream) complete(body []byte) string { + u.mu.Lock() + defer u.mu.Unlock() + u.requests = append(u.requests, append([]byte(nil), body...)) + return fmt.Sprintf(`{"type":"response.completed","response":{"id":"resp_ws_owner_%d","object":"response","status":"completed","model":"gpt-5.2","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"local owner fixture"}]}],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}`, len(u.requests)) +} + +func (u *wsResponseOwnerUpstream) Do(request *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + body, err := io.ReadAll(request.Body) + if err != nil { + return nil, err + } + event := u.complete(body) + return &http.Response{StatusCode: http.StatusOK, Header: http.Header{"Content-Type": {"text/event-stream"}}, Body: io.NopCloser(strings.NewReader("data: " + event + "\n\n"))}, nil +} + +func (u *wsResponseOwnerUpstream) DoWithTLS(request *http.Request, proxy string, accountID int64, concurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) { + return u.Do(request, proxy, accountID, concurrency) +} + +func (u *wsResponseOwnerUpstream) snapshot() [][]byte { + u.mu.Lock() + defer u.mu.Unlock() + return append([][]byte(nil), u.requests...) +} + +type wsResponseOwnerFixture struct { + handler *OpenAIGatewayHandler + upstream *wsResponseOwnerUpstream + server *httptest.Server + finished chan struct{} +} + +func newWSResponseOwnerFixture(t *testing.T, mode string) *wsResponseOwnerFixture { + t.Helper() + gin.SetMode(gin.TestMode) + f := &wsResponseOwnerFixture{upstream: &wsResponseOwnerUpstream{}, finished: make(chan struct{}, 8)} + baseURL := "https://upstream.invalid" + if mode != service.OpenAIWSIngressModeHTTPBridge { + upstreamServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover}) + if err != nil { + return + } + defer func() { _ = conn.CloseNow() }() + for { + _, body, readErr := conn.Read(r.Context()) + if readErr != nil { + return + } + if err := conn.Write(r.Context(), coderws.MessageText, []byte(f.upstream.complete(body))); err != nil { + return + } + } + })) + t.Cleanup(upstreamServer.Close) + baseURL = upstreamServer.URL + } + cfg := &config.Config{RunMode: config.RunModeSimple} + cfg.Default.RateMultiplier = 1 + cfg.Security.URLAllowlist.Enabled = false + cfg.Security.URLAllowlist.AllowInsecureHTTP = true + cfg.Gateway.OpenAIWS.Enabled = true + cfg.Gateway.OpenAIWS.APIKeyEnabled = true + cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true + cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true + cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2 + cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2 + cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3 + cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3 + repo := &httpResponseOwnerAccountRepo{&openAIWSFailoverHandlerAccountRepoStub{accounts: []service.Account{{ + ID: 61202, Name: "local-ws-owner-fixture", Platform: service.PlatformOpenAI, + Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Concurrency: 2, + Credentials: map[string]any{"api_key": "synthetic-key", "base_url": baseURL}, + Extra: map[string]any{ + "openai_passthrough": true, + "openai_apikey_responses_websockets_v2_enabled": true, + "openai_apikey_responses_websockets_v2_mode": mode, + }, + }}}} + billingCache := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil) + t.Cleanup(billingCache.Stop) + concurrency := service.NewConcurrencyService(&concurrencyCacheMock{ + acquireUserSlotFn: func(context.Context, int64, int, string) (bool, error) { return true, nil }, + acquireAccountSlotFn: func(context.Context, int64, int, string) (bool, error) { return true, nil }, + }) + gateway := service.NewOpenAIGatewayService( + repo, &openAIWSUsageHandlerUsageLogRepoStub{}, nil, nil, nil, nil, nil, cfg, nil, concurrency, + service.NewBillingService(cfg, nil), nil, billingCache, f.upstream, &service.DeferredService{}, + nil, nil, nil, nil, nil, nil, nil, + ) + t.Cleanup(gateway.CloseOpenAIWSPool) + f.handler = &OpenAIGatewayHandler{ + cfg: cfg, gatewayService: gateway, billingCacheService: billingCache, + apiKeyService: &service.APIKeyService{}, + concurrencyHelper: NewConcurrencyHelper(concurrency, SSEPingFormatNone, time.Second), + } + group := &service.Group{ID: wsResponseOwnerFixtureGroupID, Platform: service.PlatformOpenAI, Status: service.StatusActive} + wireOpenAIWSFixtureGroupReader(t, f.handler, cfg, group) + router := gin.New() + router.GET("/openai/v1/responses", func(c *gin.Context) { + // Fixture-only middleware identities are independent from all payload fields. + userID, _ := strconv.ParseInt(c.Query("user"), 10, 64) + keyID, _ := strconv.ParseInt(c.Query("key"), 10, 64) + groupID := wsResponseOwnerFixtureGroupID + c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{ + ID: keyID, UserID: userID, GroupID: &groupID, Status: service.StatusActive, + User: &service.User{ID: userID, Status: service.StatusActive}, Group: group, + }) + c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: userID, Concurrency: 1}) + defer func() { f.finished <- struct{}{} }() + f.handler.ResponsesWebSocket(c) + }) + f.server = httptest.NewServer(router) + t.Cleanup(f.server.Close) + return f +} + +func (f *wsResponseOwnerFixture) connect(t *testing.T, userID, keyID int64) *coderws.Conn { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + endpoint := fmt.Sprintf("ws%s/openai/v1/responses?user=%d&key=%d", strings.TrimPrefix(f.server.URL, "http"), userID, keyID) + conn, _, err := coderws.Dial(ctx, endpoint, nil) + require.NoError(t, err) + t.Cleanup(func() { _ = conn.CloseNow() }) + return conn +} + +func wsResponseOwnerSend(t *testing.T, conn *coderws.Conn, anchor string) { + t.Helper() + body := `{"type":"response.create","model":"gpt-5.2","input":"local fixture","stream":true,"store":true,"user_id":101,"api_key_id":1001}` + if anchor != "" { + body = fmt.Sprintf(`{"type":"response.create","model":"gpt-5.2","input":"local fixture","stream":true,"store":true,"previous_response_id":%q,"user_id":101,"api_key_id":1001}`, anchor) + } + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + require.NoError(t, conn.Write(ctx, coderws.MessageText, []byte(body))) +} + +func wsResponseOwnerReadCompleted(t *testing.T, conn *coderws.Conn) string { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + for { + _, event, err := conn.Read(ctx) + require.NoError(t, err) + if gjson.GetBytes(event, "type").String() == "response.completed" { + return gjson.GetBytes(event, "response.id").String() + } + } +} + +func (f *wsResponseOwnerFixture) close(t *testing.T, conn *coderws.Conn) { + t.Helper() + _ = conn.Close(coderws.StatusNormalClosure, "fixture done") + select { + case <-f.finished: + case <-time.After(3 * time.Second): + t.Fatal("owner fixture handler did not finish") + } +} + +func wsResponseOwnerAssertRejected(t *testing.T, conn *coderws.Conn) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + _, event, err := conn.Read(ctx) + require.Error(t, err, "foreign anchor reached a synthetic upstream and produced %s", event) + var closeErr coderws.CloseError + require.ErrorAs(t, err, &closeErr) + require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code) + require.Equal(t, "previous_response_id is not available for this user", closeErr.Reason) +} + +func TestOpenAIResponsesWSOwner_RejectsForeignFirstAndSubsequentFrames(t *testing.T) { + for _, mode := range []string{service.OpenAIWSIngressModeHTTPBridge, service.OpenAIWSIngressModeCtxPool, service.OpenAIWSIngressModePassthrough} { + for _, subsequent := range []bool{false, true} { + t.Run(fmt.Sprintf("%s/subsequent=%t", mode, subsequent), func(t *testing.T) { + f := newWSResponseOwnerFixture(t, mode) + owner := f.connect(t, 101, 1001) + wsResponseOwnerSend(t, owner, "") + privateID := wsResponseOwnerReadCompleted(t, owner) + require.Equal(t, "resp_ws_owner_1", privateID) + f.close(t, owner) + + other := f.connect(t, 202, 2002) + if subsequent { + wsResponseOwnerSend(t, other, "") + require.Equal(t, "resp_ws_owner_2", wsResponseOwnerReadCompleted(t, other)) + } + before := len(f.upstream.snapshot()) + wsResponseOwnerSend(t, other, privateID) + wsResponseOwnerAssertRejected(t, other) + require.Len(t, f.upstream.snapshot(), before, "foreign continuation must produce no additional upstream IO") + owned, err := f.handler.gatewayService.ValidateOpenAIHTTPResponseOwner(context.Background(), wsResponseOwnerFixtureGroupID, privateID, 101, 1001) + require.NoError(t, err) + require.True(t, owned, "a rejected frame must preserve the original response owner") + foreignOwned, err := f.handler.gatewayService.ValidateOpenAIHTTPResponseOwner(context.Background(), wsResponseOwnerFixtureGroupID, privateID, 202, 2002) + require.NoError(t, err) + require.False(t, foreignOwned, "a rejected anchor must not be rebound to the caller") + }) + } + } +} + +func TestOpenAIResponsesWSOwner_RegistersAndContinuesSameUser(t *testing.T) { + for _, mode := range []string{service.OpenAIWSIngressModeHTTPBridge, service.OpenAIWSIngressModeCtxPool, service.OpenAIWSIngressModePassthrough} { + t.Run(mode, func(t *testing.T) { + f := newWSResponseOwnerFixture(t, mode) + first := f.connect(t, 202, 2002) + wsResponseOwnerSend(t, first, "") + firstID := wsResponseOwnerReadCompleted(t, first) + f.close(t, first) + owned, err := f.handler.gatewayService.ValidateOpenAIHTTPResponseOwner(context.Background(), wsResponseOwnerFixtureGroupID, firstID, 202, 2002) + require.NoError(t, err) + require.True(t, owned, "real WS completion must register authenticated user, not body user_id") + + // Existing HTTP contract permits another key of the same user. + next := f.connect(t, 202, 2003) + wsResponseOwnerSend(t, next, firstID) + nextID := wsResponseOwnerReadCompleted(t, next) + wsResponseOwnerSend(t, next, nextID) + require.Equal(t, "resp_ws_owner_3", wsResponseOwnerReadCompleted(t, next)) + f.close(t, next) + requests := f.upstream.snapshot() + require.Len(t, requests, 3) + require.Equal(t, firstID, gjson.GetBytes(requests[1], "previous_response_id").String()) + require.Equal(t, nextID, gjson.GetBytes(requests[2], "previous_response_id").String()) + }) + } +} diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index 410e9f47244a..572c82591d07 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -1679,6 +1679,29 @@ func (r *accountRepository) UpdateGrokOAuthCredentialsIfUnchanged( expectedCredentials map[string]any, expectedProxyID *int64, credentials map[string]any, +) (bool, error) { + return r.updateOAuthCredentialsIfUnchanged(ctx, id, service.PlatformGrok, expectedCredentials, expectedProxyID, credentials) +} + +// UpdateOpenAIOAuthCredentialsIfUnchanged uses the same atomic credential and +// scheduler-outbox boundary, scoped specifically to an OpenAI OAuth row. +func (r *accountRepository) UpdateOpenAIOAuthCredentialsIfUnchanged( + ctx context.Context, + id int64, + expectedCredentials map[string]any, + expectedProxyID *int64, + credentials map[string]any, +) (bool, error) { + return r.updateOAuthCredentialsIfUnchanged(ctx, id, service.PlatformOpenAI, expectedCredentials, expectedProxyID, credentials) +} + +func (r *accountRepository) updateOAuthCredentialsIfUnchanged( + ctx context.Context, + id int64, + platform string, + expectedCredentials map[string]any, + expectedProxyID *int64, + credentials map[string]any, ) (bool, error) { if r == nil || r.sql == nil { return false, errors.New("account repository SQL executor is not configured") @@ -1709,7 +1732,7 @@ func (r *accountRepository) UpdateGrokOAuthCredentialsIfUnchanged( `, string(credentialsJSON), id, - service.PlatformGrok, + platform, service.AccountTypeOAuth, string(expectedJSON), expectedProxyID, diff --git a/backend/internal/repository/account_repo_openai_identity_test.go b/backend/internal/repository/account_repo_openai_identity_test.go new file mode 100644 index 000000000000..014bec7ae364 --- /dev/null +++ b/backend/internal/repository/account_repo_openai_identity_test.go @@ -0,0 +1,44 @@ +package repository + +import ( + "context" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestAccountRepository_OAuthIdentityCASPlatformAndAtomicOutbox(t *testing.T) { + for _, platform := range []string{service.PlatformOpenAI, service.PlatformGrok} { + t.Run(platform, func(t *testing.T) { + exec := &recordingSQLExecutor{result: rowsAffectedResult(1)} + repo := newAccountRepositoryWithSQL(nil, exec, nil) + proxyID := int64(29) + update := repo.UpdateOpenAIOAuthCredentialsIfUnchanged + if platform == service.PlatformGrok { + update = repo.UpdateGrokOAuthCredentialsIfUnchanged + } + applied, err := update(context.Background(), 42, + map[string]any{"refresh_token": "synthetic-old", "_token_version": int64(9)}, + &proxyID, + map[string]any{"refresh_token": "synthetic-new", "_token_version": int64(10)}, + ) + require.NoError(t, err) + require.True(t, applied) + require.Len(t, exec.execQueries, 1) + query := normalizeSQLWhitespace(exec.execQueries[0]) + require.Contains(t, query, "WITH updated AS") + require.Contains(t, query, "a.platform = $3") + require.Contains(t, query, "a.type = $4") + require.Contains(t, query, "a.credentials = $5::jsonb") + require.Contains(t, query, "a.proxy_id IS NOT DISTINCT FROM $6") + require.Contains(t, query, "INSERT INTO scheduler_outbox") + require.NotContains(t, query, "SET status") + require.Len(t, exec.execArgs[0], 7) + require.Equal(t, platform, exec.execArgs[0][2]) + require.Equal(t, service.AccountTypeOAuth, exec.execArgs[0][3]) + require.Equal(t, &proxyID, exec.execArgs[0][5]) + require.Equal(t, service.SchedulerOutboxEventAccountChanged, exec.execArgs[0][6]) + }) + } +} diff --git a/backend/internal/service/oauth_refresh_api.go b/backend/internal/service/oauth_refresh_api.go index 4bff2dad9108..54d16f23604c 100644 --- a/backend/internal/service/oauth_refresh_api.go +++ b/backend/internal/service/oauth_refresh_api.go @@ -34,6 +34,18 @@ type GrokOAuthRefreshSuccessRepository interface { ) (bool, error) } +// OpenAIOAuthRefreshSuccessRepository keeps a completed token rotation from +// overwriting credentials or a proxy selected by a concurrent admin edit. +type OpenAIOAuthRefreshSuccessRepository interface { + UpdateOpenAIOAuthCredentialsIfUnchanged( + ctx context.Context, + id int64, + expectedCredentials map[string]any, + expectedProxyID *int64, + credentials map[string]any, + ) (bool, error) +} + const ( defaultRefreshLockTTL = 60 * time.Second defaultRefreshLockReleaseTimeout = 2 * time.Second @@ -337,7 +349,7 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded( ) return &OAuthRefreshResult{Account: currentAccount}, nil } - durableAccount, readErr := api.loadGrokDurableAccountAfterPersist(ctx, cacheKey, freshAccount.ID) + durableAccount, readErr := api.loadDurableAccountAfterPersist(ctx, cacheKey, freshAccount.ID) if readErr != nil || durableAccount == nil { if readErr == nil { readErr = fmt.Errorf("account not found after Grok OAuth success CAS") @@ -351,6 +363,46 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded( // while the provider call was in flight. Return the durable row so // post-refresh cache publication cannot restore that stale snapshot. freshAccount = durableAccount + } else if freshAccount.IsOpenAIOAuth() { + conditionalRepo, ok := api.accountRepo.(OpenAIOAuthRefreshSuccessRepository) + if !ok { + return nil, &providerConfigurationRefreshError{ + err: fmt.Errorf("OpenAI OAuth refresh success CAS repository is not configured"), + } + } + applied, updateErr := conditionalRepo.UpdateOpenAIOAuthCredentialsIfUnchanged( + ctx, freshAccount.ID, attemptedAccount.Credentials, attemptedAccount.ProxyID, newCredentials, + ) + if updateErr != nil { + return nil, &providerCycleContainmentRefreshError{ + err: fmt.Errorf("OpenAI OAuth refresh succeeded but credential persistence failed: %w", updateErr), + } + } + if !applied { + currentAccount, readErr := api.accountRepo.GetByID(ctx, freshAccount.ID) + if readErr != nil || currentAccount == nil { + return nil, &providerCycleContainmentRefreshError{ + err: fmt.Errorf("OpenAI OAuth success CAS lost and current state is unavailable"), + } + } + return &OAuthRefreshResult{Account: currentAccount}, nil + } + durableAccount, readErr := api.loadDurableAccountAfterPersist(ctx, cacheKey, freshAccount.ID) + if readErr != nil || durableAccount == nil { + return nil, &providerCycleContainmentRefreshError{ + err: fmt.Errorf("OpenAI OAuth success persisted but durable account state is unavailable"), + } + } + persistedIdentity := *attemptedAccount + persistedIdentity.Credentials = newCredentials + if validateOpenAITokenIdentity(&persistedIdentity, durableAccount) != nil || + persistedIdentity.GetOpenAIAccessToken() != durableAccount.GetOpenAIAccessToken() || + persistedIdentity.GetCredentialAsInt64("_token_version") != durableAccount.GetCredentialAsInt64("_token_version") { + // An admin edit can also win after CAS but before the durable read. + // Do not apply the old refresh's post-success actions to that edit. + return &OAuthRefreshResult{Account: durableAccount}, nil + } + freshAccount = durableAccount } else if updateErr := persistAccountCredentials(ctx, api.accountRepo, freshAccount, newCredentials); updateErr != nil { slog.Error("oauth_refresh_update_failed", "account_id", freshAccount.ID, @@ -385,7 +437,7 @@ func (api *OAuthRefreshAPI) releaseRefreshLock(parent context.Context, cacheKey } } -func (api *OAuthRefreshAPI) loadGrokDurableAccountAfterPersist(parent context.Context, cacheKey string, accountID int64) (*Account, error) { +func (api *OAuthRefreshAPI) loadDurableAccountAfterPersist(parent context.Context, cacheKey string, accountID int64) (*Account, error) { cleanupParent := context.Background() if parent != nil { cleanupParent = context.WithoutCancel(parent) diff --git a/backend/internal/service/oauth_refresh_api_test.go b/backend/internal/service/oauth_refresh_api_test.go index 3baebd4349ba..4bb5587bfba4 100644 --- a/backend/internal/service/oauth_refresh_api_test.go +++ b/backend/internal/service/oauth_refresh_api_test.go @@ -103,6 +103,31 @@ func (r *refreshAPIAccountRepo) UpdateGrokOAuthCredentialsIfUnchanged( return true, nil } +func (r *refreshAPIAccountRepo) UpdateOpenAIOAuthCredentialsIfUnchanged( + _ context.Context, + id int64, + expectedCredentials map[string]any, + expectedProxyID *int64, + credentials map[string]any, +) (bool, error) { + r.successCASCalls++ + if r.beforeSuccessCAS != nil { + r.beforeSuccessCAS(r) + } + if r.updateErr != nil { + return false, r.updateErr + } + if r.account == nil || r.account.ID != id || !r.account.IsOpenAIOAuth() || + !reflect.DeepEqual(r.account.Credentials, expectedCredentials) || + !reflect.DeepEqual(r.account.ProxyID, expectedProxyID) { + return false, nil + } + r.updateCalls++ + r.updateCredentialsCalls++ + r.account.Credentials = shallowCopyMap(credentials) + return true, nil +} + // refreshAPIExecutorStub implements OAuthRefreshExecutor for tests. type refreshAPIExecutorStub struct { needsRefresh bool diff --git a/backend/internal/service/openai_identity_mutation_test.go b/backend/internal/service/openai_identity_mutation_test.go new file mode 100644 index 000000000000..789e00fa91ea --- /dev/null +++ b/backend/internal/service/openai_identity_mutation_test.go @@ -0,0 +1,373 @@ +//go:build unit + +package service + +import ( + "bytes" + "context" + "errors" + "net/http" + "net/http/httptest" + "reflect" + "sync" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// Return independent snapshots, as the real repository does. Sharing the same +// Account pointer would hide a refresh overwriting a concurrent admin edit. +type openAIIdentityMutationRepo struct { + AccountRepository + mu sync.Mutex + account *Account + getErr error + casErr error + getCalls int + afterCAS func() +} + +func copyOpenAIIdentityMutationAccount(account *Account) *Account { + if account == nil { + return nil + } + copy := *account + copy.Credentials = shallowCopyMap(account.Credentials) + copy.Extra = shallowCopyMap(account.Extra) + return © +} + +func (r *openAIIdentityMutationRepo) GetByID(context.Context, int64) (*Account, error) { + r.mu.Lock() + defer r.mu.Unlock() + r.getCalls++ + if r.getErr != nil { + return nil, r.getErr + } + return copyOpenAIIdentityMutationAccount(r.account), nil +} + +func (r *openAIIdentityMutationRepo) Update(_ context.Context, account *Account) error { + r.mu.Lock() + defer r.mu.Unlock() + r.account = copyOpenAIIdentityMutationAccount(account) + return nil +} + +func (r *openAIIdentityMutationRepo) UpdateCredentials(_ context.Context, _ int64, credentials map[string]any) error { + r.mu.Lock() + defer r.mu.Unlock() + r.account.Credentials = shallowCopyMap(credentials) + return nil +} + +func (r *openAIIdentityMutationRepo) UpdateOpenAIOAuthCredentialsIfUnchanged( + _ context.Context, + id int64, + expectedCredentials map[string]any, + expectedProxyID *int64, + credentials map[string]any, +) (bool, error) { + r.mu.Lock() + if r.casErr != nil { + err := r.casErr + r.mu.Unlock() + return false, err + } + if r.account.ID != id || !r.account.IsOpenAIOAuth() || + !reflect.DeepEqual(r.account.Credentials, expectedCredentials) || + !reflect.DeepEqual(r.account.ProxyID, expectedProxyID) { + r.mu.Unlock() + return false, nil + } + r.account.Credentials = shallowCopyMap(credentials) + afterCAS := r.afterCAS + r.afterCAS = nil + r.mu.Unlock() + if afterCAS != nil { + afterCAS() + } + return true, nil +} + +type openAIIdentityMutationCache struct { + *openAITokenCacheStub + getCalls int + beforeGet func(int) +} + +func (c *openAIIdentityMutationCache) GetAccessToken(ctx context.Context, key string) (string, error) { + c.getCalls++ + if c.beforeGet != nil { + c.beforeGet(c.getCalls) + } + return c.openAITokenCacheStub.GetAccessToken(ctx, key) +} + +type openAIIdentityMutationExecutor struct { + *refreshAPIExecutorStub +} + +func (e *openAIIdentityMutationExecutor) CacheKey(account *Account) string { + return OpenAITokenCacheKey(account) +} + +func openAIIdentityMutationCredentials(identity string) map[string]any { + return map[string]any{ + "access_token": "synthetic-token-" + identity, + "refresh_token": "synthetic-refresh-" + identity, + "chatgpt_account_id": "synthetic-workspace-" + identity, + "chatgpt_user_id": "synthetic-user-" + identity, + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + } +} + +func newOpenAIIdentityMutationRepo() *openAIIdentityMutationRepo { + return &openAIIdentityMutationRepo{account: &Account{ + ID: 62001, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Credentials: openAIIdentityMutationCredentials("a"), + }} +} + +func buildOpenAIIdentityMutationRequest(t *testing.T, gateway *OpenAIGatewayService, account *Account) *http.Request { + t.Helper() + ctx := context.Background() + token, _, err := gateway.GetAccessToken(ctx, account) + require.NoError(t, err) + body := []byte(`{"model":"gpt-5.5","input":"synthetic identity fixture"}`) + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + request, err := gateway.buildUpstreamRequest(ctx, c, account, body, token, true, "", true) + require.NoError(t, err) + return request +} + +func TestOpenAIIdentity_AdminCredentialUpdateUsesCurrentTokenAndWorkspace(t *testing.T) { + for _, cacheState := range []string{"already_cached", "old_snapshot_fills_cache_after_update"} { + t.Run(cacheState, func(t *testing.T) { + ctx := context.Background() + repo := newOpenAIIdentityMutationRepo() + cache := newOpenAITokenCacheStub() + provider := NewOpenAITokenProvider(repo, cache, nil) + selectedBeforeEdit, err := repo.GetByID(ctx, repo.account.ID) + require.NoError(t, err) + if cacheState == "already_cached" { + token, err := provider.GetAccessToken(ctx, selectedBeforeEdit) + require.NoError(t, err) + require.Equal(t, "synthetic-token-a", token) + } + + admin := &adminServiceImpl{accountRepo: repo} + updated, err := admin.UpdateAccount(ctx, selectedBeforeEdit.ID, &UpdateAccountInput{ + Credentials: openAIIdentityMutationCredentials("b"), + }) + require.NoError(t, err) + require.Equal(t, "synthetic-token-b", updated.GetOpenAIAccessToken()) + if cacheState == "old_snapshot_fills_cache_after_update" { + // A request selected before the edit may finish or fail, but it must + // not populate an A token that a later B request can consume. + _, _ = provider.GetAccessToken(ctx, selectedBeforeEdit) + } + + gateway := &OpenAIGatewayService{accountRepo: repo, openAITokenProvider: provider} + request := buildOpenAIIdentityMutationRequest(t, gateway, updated) + require.Equal(t, chatgptCodexURL, request.URL.String()) + require.Equal(t, "synthetic-workspace-b", request.Header.Get("Chatgpt-Account-Id")) + require.Equal(t, "Bearer synthetic-token-b", request.Header.Get("Authorization"), + "the real Responses builder must pair the current credential with its workspace") + }) + } + + t.Run("matching_cache_hit_does_not_read_repository", func(t *testing.T) { + ctx := context.Background() + repo := newOpenAIIdentityMutationRepo() + account, err := repo.GetByID(ctx, repo.account.ID) + require.NoError(t, err) + cache := newOpenAITokenCacheStub() + require.NoError(t, cache.SetAccessToken(ctx, OpenAITokenCacheKey(account), account.GetOpenAIAccessToken(), time.Minute)) + provider := NewOpenAITokenProvider(repo, cache, nil) + reads := repo.getCalls + request := buildOpenAIIdentityMutationRequest(t, &OpenAIGatewayService{accountRepo: repo, openAITokenProvider: provider}, account) + require.Equal(t, "Bearer synthetic-token-a", request.Header.Get("Authorization")) + require.Equal(t, reads, repo.getCalls) + }) + + t.Run("unverified_different_cache_token_is_ignored", func(t *testing.T) { + ctx := context.Background() + repo := newOpenAIIdentityMutationRepo() + admin := &adminServiceImpl{accountRepo: repo} + updated, err := admin.UpdateAccount(ctx, repo.account.ID, &UpdateAccountInput{Credentials: openAIIdentityMutationCredentials("b")}) + require.NoError(t, err) + cache := newOpenAITokenCacheStub() + require.NoError(t, cache.SetAccessToken(ctx, OpenAITokenCacheKey(updated), "synthetic-token-a", time.Minute)) + repo.getErr = errors.New("synthetic repository unavailable") + provider := NewOpenAITokenProvider(repo, cache, nil) + request := buildOpenAIIdentityMutationRequest(t, &OpenAIGatewayService{accountRepo: repo, openAITokenProvider: provider}, updated) + require.Equal(t, "Bearer synthetic-token-b", request.Header.Get("Authorization")) + require.Equal(t, "synthetic-workspace-b", request.Header.Get("Chatgpt-Account-Id")) + }) + + t.Run("version_check_cannot_adopt_another_workspace", func(t *testing.T) { + ctx := context.Background() + repo := newOpenAIIdentityMutationRepo() + selected, err := repo.GetByID(ctx, repo.account.ID) + require.NoError(t, err) + credentials := openAIIdentityMutationCredentials("b") + credentials["_token_version"] = int64(2) + admin := &adminServiceImpl{accountRepo: repo} + _, err = admin.UpdateAccount(ctx, selected.ID, &UpdateAccountInput{Credentials: credentials}) + require.NoError(t, err) + provider := NewOpenAITokenProvider(repo, newOpenAITokenCacheStub(), nil) + token, err := provider.GetAccessToken(ctx, selected) + require.ErrorIs(t, err, errOpenAITokenIdentityChanged) + require.Empty(t, token) + }) + + t.Run("lock_wait_cannot_adopt_another_workspace", func(t *testing.T) { + ctx := context.Background() + repo := newOpenAIIdentityMutationRepo() + repo.account.Credentials["expires_at"] = time.Now().Add(time.Minute).UTC().Format(time.RFC3339) + selected, err := repo.GetByID(ctx, repo.account.ID) + require.NoError(t, err) + cache := &openAIIdentityMutationCache{openAITokenCacheStub: newOpenAITokenCacheStub()} + cache.lockAcquired = false + cache.beforeGet = func(call int) { + if call != 2 { + return + } + admin := &adminServiceImpl{accountRepo: repo} + _, err := admin.UpdateAccount(ctx, selected.ID, &UpdateAccountInput{Credentials: openAIIdentityMutationCredentials("b")}) + require.NoError(t, err) + require.NoError(t, cache.SetAccessToken(ctx, OpenAITokenCacheKey(selected), "synthetic-token-b", time.Minute)) + } + provider := NewOpenAITokenProvider(repo, cache, nil) + token, err := provider.GetAccessToken(ctx, selected) + require.ErrorIs(t, err, errOpenAITokenIdentityChanged) + require.Empty(t, token) + }) +} + +func TestOpenAIIdentity_LateRefreshCannotOverwriteAdminCredentials(t *testing.T) { + t.Run("admin_reauthorization_before_cas", func(t *testing.T) { + ctx := context.Background() + repo := newOpenAIIdentityMutationRepo() + account, err := repo.GetByID(ctx, repo.account.ID) + require.NoError(t, err) + admin := &adminServiceImpl{accountRepo: repo} + lateCredentials := openAIIdentityMutationCredentials("a") + lateCredentials["access_token"] = "synthetic-late-refresh-a" + executor := &refreshAPIExecutorStub{ + needsRefresh: true, + credentials: lateCredentials, + onRefresh: func() { + // Deterministically perform the admin commit after RefreshIfNeeded + // reads A and before the provider returns A's replacement token. + updated, err := admin.UpdateAccount(ctx, account.ID, &UpdateAccountInput{ + Credentials: openAIIdentityMutationCredentials("b"), + }) + require.NoError(t, err) + require.Equal(t, "synthetic-token-b", updated.GetOpenAIAccessToken()) + }, + } + api := NewOAuthRefreshAPI(repo, newOpenAITokenCacheStub()) + result, err := api.RefreshIfNeeded(ctx, account, executor, 3*time.Minute) + require.NoError(t, err) + durable, err := repo.GetByID(ctx, account.ID) + require.NoError(t, err) + require.Equal(t, "synthetic-token-b", durable.GetOpenAIAccessToken(), + "a successful admin edit must survive an older OAuth refresh finishing later") + require.Equal(t, "synthetic-workspace-b", durable.GetChatGPTAccountID()) + require.Equal(t, "synthetic-token-b", result.Account.GetOpenAIAccessToken()) + require.False(t, result.Refreshed, "a discarded stale refresh must not be published as committed") + }) + + t.Run("normal_same_workspace_rotation", func(t *testing.T) { + ctx := context.Background() + repo := newOpenAIIdentityMutationRepo() + repo.account.Credentials["expires_at"] = time.Now().Add(time.Minute).UTC().Format(time.RFC3339) + selected, err := repo.GetByID(ctx, repo.account.ID) + require.NoError(t, err) + credentials := openAIIdentityMutationCredentials("a") + credentials["access_token"] = "synthetic-rotated-token-a" + cache := newOpenAITokenCacheStub() + executor := &openAIIdentityMutationExecutor{&refreshAPIExecutorStub{needsRefresh: true, credentials: credentials}} + provider := NewOpenAITokenProvider(repo, cache, nil) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), executor) + gateway := &OpenAIGatewayService{accountRepo: repo, openAITokenProvider: provider} + request := buildOpenAIIdentityMutationRequest(t, gateway, selected) + require.Equal(t, "Bearer synthetic-rotated-token-a", request.Header.Get("Authorization")) + require.Equal(t, "synthetic-workspace-a", request.Header.Get("Chatgpt-Account-Id")) + // A second old scheduling snapshot may consume the verified background + // rotation from the cache without mixing workspaces or refreshing twice. + request = buildOpenAIIdentityMutationRequest(t, gateway, selected) + require.Equal(t, "Bearer synthetic-rotated-token-a", request.Header.Get("Authorization")) + require.Equal(t, 1, executor.refreshCalls) + }) + + t.Run("refresh_result_cannot_change_the_callers_workspace", func(t *testing.T) { + ctx := context.Background() + repo := newOpenAIIdentityMutationRepo() + repo.account.Credentials["expires_at"] = time.Now().Add(time.Minute).UTC().Format(time.RFC3339) + selected, err := repo.GetByID(ctx, repo.account.ID) + require.NoError(t, err) + admin := &adminServiceImpl{accountRepo: repo} + executor := &openAIIdentityMutationExecutor{&refreshAPIExecutorStub{ + needsRefresh: true, + credentials: openAIIdentityMutationCredentials("a"), + onRefresh: func() { + _, err := admin.UpdateAccount(ctx, selected.ID, &UpdateAccountInput{Credentials: openAIIdentityMutationCredentials("b")}) + require.NoError(t, err) + }, + }} + cache := newOpenAITokenCacheStub() + provider := NewOpenAITokenProvider(repo, cache, nil) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), executor) + token, err := provider.GetAccessToken(ctx, selected) + require.ErrorIs(t, err, errOpenAITokenIdentityChanged) + require.Empty(t, token) + require.Equal(t, "synthetic-token-b", repo.account.GetOpenAIAccessToken()) + }) + + t.Run("admin_reauthorization_after_cas", func(t *testing.T) { + ctx := context.Background() + repo := newOpenAIIdentityMutationRepo() + selected, err := repo.GetByID(ctx, repo.account.ID) + require.NoError(t, err) + admin := &adminServiceImpl{accountRepo: repo} + repo.afterCAS = func() { + _, err := admin.UpdateAccount(ctx, selected.ID, &UpdateAccountInput{Credentials: openAIIdentityMutationCredentials("b")}) + require.NoError(t, err) + } + api := NewOAuthRefreshAPI(repo, nil) + result, err := api.RefreshIfNeeded(ctx, selected, &refreshAPIExecutorStub{needsRefresh: true, credentials: openAIIdentityMutationCredentials("a")}, time.Minute) + require.NoError(t, err) + require.False(t, result.Refreshed) + require.Nil(t, result.NewCredentials) + require.Equal(t, "synthetic-token-b", result.Account.GetOpenAIAccessToken()) + }) + + t.Run("cas_failure_does_not_publish_or_overwrite", func(t *testing.T) { + ctx := context.Background() + repo := newOpenAIIdentityMutationRepo() + repo.account.Credentials["expires_at"] = time.Now().Add(time.Minute).UTC().Format(time.RFC3339) + selected, err := repo.GetByID(ctx, repo.account.ID) + require.NoError(t, err) + repo.casErr = errors.New("synthetic CAS write failure") + cache := newOpenAITokenCacheStub() + provider := NewOpenAITokenProvider(repo, cache, nil) + provider.SetRefreshAPI(NewOAuthRefreshAPI(repo, cache), &openAIIdentityMutationExecutor{&refreshAPIExecutorStub{ + needsRefresh: true, credentials: openAIIdentityMutationCredentials("a"), + }}) + token, err := provider.GetAccessToken(ctx, selected) + var containment *providerCycleContainmentRefreshError + require.ErrorAs(t, err, &containment) + require.Empty(t, token, "an ambiguous rotating-token persistence result must not fall back to the old bearer") + require.Equal(t, "synthetic-token-a", repo.account.GetOpenAIAccessToken()) + }) +} diff --git a/backend/internal/service/openai_token_provider.go b/backend/internal/service/openai_token_provider.go index 09997616231f..7cceced33867 100644 --- a/backend/internal/service/openai_token_provider.go +++ b/backend/internal/service/openai_token_provider.go @@ -130,6 +130,41 @@ func (p *OpenAITokenProvider) ensureMetrics() { } } +var errOpenAITokenIdentityChanged = errors.New("openai OAuth account identity changed; select the account again") + +// A provider returns only a token. Its caller still owns the selected account +// and may already have derived request/session headers from that snapshot. +func validateOpenAITokenIdentity(selected, latest *Account) error { + if selected == nil || latest == nil || selected.ID != latest.ID || + !latest.IsOpenAIOAuth() || + selected.GetChatGPTAccountID() != latest.GetChatGPTAccountID() || + codexAccountIdentityNamespace(selected) != codexAccountIdentityNamespace(latest) || + selected.IsChatGPTAccountFedRAMP() != latest.IsChatGPTAccountFedRAMP() || + selected.IsOpenAIAgentIdentity() != latest.IsOpenAIAgentIdentity() || + selected.IsOpenAIPersonalAccessToken() != latest.IsOpenAIPersonalAccessToken() { + return errOpenAITokenIdentityChanged + } + return nil +} + +// Matching snapshots retain the cache-only hot path. A different bearer may +// be a normal background rotation or an old value left by an admin edit; only +// the durable account can distinguish them. Repository-less callers retain +// the existing standalone cache contract; production always supplies a repo. +func (p *OpenAITokenProvider) cachedTokenMatchesAccount(ctx context.Context, account *Account, token string) (bool, error) { + if token == account.GetOpenAIAccessToken() || p.accountRepo == nil { + return true, nil + } + latest, err := p.accountRepo.GetByID(ctx, account.ID) + if err != nil || latest == nil { + return false, nil + } + if err := validateOpenAITokenIdentity(account, latest); err != nil { + return false, err + } + return token == latest.GetOpenAIAccessToken(), nil +} + // GetAccessToken returns a valid access_token. func (p *OpenAITokenProvider) GetAccessToken(ctx context.Context, account *Account) (string, error) { p.ensureMetrics() @@ -140,13 +175,20 @@ func (p *OpenAITokenProvider) GetAccessToken(ctx context.Context, account *Accou return "", errors.New("not an openai oauth account") } + selectedAccount := account cacheKey := OpenAITokenCacheKey(account) // 1) Try cache first. if p.tokenCache != nil { if token, err := p.tokenCache.GetAccessToken(ctx, cacheKey); err == nil && strings.TrimSpace(token) != "" { - slog.Debug("openai_token_cache_hit", "account_id", account.ID) - return token, nil + matches, identityErr := p.cachedTokenMatchesAccount(ctx, selectedAccount, token) + if identityErr != nil { + return "", identityErr + } + if matches { + slog.Debug("openai_token_cache_hit", "account_id", account.ID) + return token, nil + } } else if err != nil { slog.Warn("openai_token_cache_get_failed", "account_id", account.ID, "error", err) } @@ -177,7 +219,10 @@ func (p *OpenAITokenProvider) GetAccessToken(ctx context.Context, account *Accou result, err := p.refreshAPI.RefreshIfNeeded(ctx, account, p.executor, openAITokenRefreshSkew) if err != nil { - if p.refreshPolicy.OnRefreshError == ProviderRefreshErrorReturn { + var containment *providerCycleContainmentRefreshError + var configuration *providerConfigurationRefreshError + if errors.As(err, &containment) || errors.As(err, &configuration) || + p.refreshPolicy.OnRefreshError == ProviderRefreshErrorReturn { return "", err } slog.Warn("openai_token_refresh_failed", "account_id", account.ID, "error", err) @@ -187,7 +232,7 @@ func (p *OpenAITokenProvider) GetAccessToken(ctx context.Context, account *Accou if p.refreshPolicy.OnLockHeld == ProviderLockHeldWaitForCache { p.metrics.lockContention.Add(1) p.metrics.touchNow() - token, waitErr := p.waitForTokenAfterLockRace(ctx, cacheKey) + token, waitErr := p.waitForTokenAfterLockRace(ctx, selectedAccount, cacheKey) if waitErr != nil { return "", waitErr } @@ -196,11 +241,13 @@ func (p *OpenAITokenProvider) GetAccessToken(ctx context.Context, account *Accou return token, nil } } - } else if result.Refreshed { - p.metrics.refreshSuccess.Add(1) - account = result.Account - expiresAt = account.GetCredentialAsTime("expires_at") } else { + if identityErr := validateOpenAITokenIdentity(selectedAccount, result.Account); identityErr != nil { + return "", identityErr + } + if result.Refreshed { + p.metrics.refreshSuccess.Add(1) + } account = result.Account expiresAt = account.GetCredentialAsTime("expires_at") } @@ -218,7 +265,7 @@ func (p *OpenAITokenProvider) GetAccessToken(ctx context.Context, account *Accou } else { p.metrics.lockContention.Add(1) p.metrics.touchNow() - token, waitErr := p.waitForTokenAfterLockRace(ctx, cacheKey) + token, waitErr := p.waitForTokenAfterLockRace(ctx, selectedAccount, cacheKey) if waitErr != nil { return "", waitErr } @@ -237,6 +284,11 @@ func (p *OpenAITokenProvider) GetAccessToken(ctx context.Context, account *Accou // 3) Populate cache with TTL. if p.tokenCache != nil { latestAccount, isStale := CheckTokenVersion(ctx, account, p.accountRepo) + if latestAccount != nil { + if identityErr := validateOpenAITokenIdentity(selectedAccount, latestAccount); identityErr != nil { + return "", identityErr + } + } if isStale && latestAccount != nil { slog.Debug("openai_token_version_stale_use_latest", "account_id", account.ID) accessToken = latestAccount.GetOpenAIAccessToken() @@ -308,7 +360,7 @@ func (p *OpenAITokenProvider) disableAccountMissingRefreshToken(account *Account ) } -func (p *OpenAITokenProvider) waitForTokenAfterLockRace(ctx context.Context, cacheKey string) (string, error) { +func (p *OpenAITokenProvider) waitForTokenAfterLockRace(ctx context.Context, account *Account, cacheKey string) (string, error) { wait := openAILockInitialWait totalWaitMs := int64(0) for i := 0; i < openAILockMaxAttempts; i++ { @@ -337,11 +389,17 @@ func (p *OpenAITokenProvider) waitForTokenAfterLockRace(ctx context.Context, cac token, err := p.tokenCache.GetAccessToken(ctx, cacheKey) if err == nil && strings.TrimSpace(token) != "" { - p.metrics.lockWaitHit.Add(1) - if totalWaitMs >= openAILockWarnThresholdMs { - slog.Warn("openai_token_lock_wait_high", "wait_ms", totalWaitMs, "attempts", i+1) + matches, identityErr := p.cachedTokenMatchesAccount(ctx, account, token) + if identityErr != nil { + return "", identityErr + } + if matches { + p.metrics.lockWaitHit.Add(1) + if totalWaitMs >= openAILockWarnThresholdMs { + slog.Warn("openai_token_lock_wait_high", "wait_ms", totalWaitMs, "attempts", i+1) + } + return token, nil } - return token, nil } if wait < openAILockMaxWait { diff --git a/backend/internal/service/openai_v183_compat.go b/backend/internal/service/openai_v183_compat.go index df02b80e1366..c0a0a4fa7c9d 100644 --- a/backend/internal/service/openai_v183_compat.go +++ b/backend/internal/service/openai_v183_compat.go @@ -338,6 +338,16 @@ func (s *OpenAIGatewayService) BindOpenAIHTTPResponseOwner(ctx context.Context, return s.getOpenAIWSStateStore().BindHTTPResponseOwner(ctx, groupID, responseID, userID, apiKeyID, s.openAIWSResponseStickyTTL()) } +// BindOpenAIWSResponseOwner registers the actual successful upstream response +// using the same authenticated owner and bounded persistence as HTTP. Never +// register a client-supplied continuation anchor or payload identity here. +func (s *OpenAIGatewayService) BindOpenAIWSResponseOwner(ctx context.Context, c *gin.Context, account *Account, result *OpenAIForwardResult) { + if result == nil || result.ClientDisconnect || !isOpenAIWSSuccessTerminalEvent(result.UpstreamTerminalEvent) { + return + } + s.bindHTTPResponseAccount(ctx, c, account, result.RequestID) +} + func (s *OpenAIGatewayService) newOpenAIAccountFailoverError( account *Account, statusCode int, From 7e1a7334d1c8f7cd7ae06e54192a9d18af10253a Mon Sep 17 00:00:00 2001 From: Codex Date: Sat, 26 Sep 2026 04:04:35 +0800 Subject: [PATCH 2/2] fix(openai): preserve admin metadata during token rotation --- backend/internal/repository/account_repo.go | 60 ++++++++- .../account_repo_openai_identity_test.go | 14 ++- ...account_usage_service_spark_shadow_test.go | 1 + backend/internal/service/oauth_refresh_api.go | 5 + .../service/oauth_refresh_api_test.go | 11 +- .../service/openai_identity_mutation_test.go | 119 +++++++++++++++++- .../service/openai_quota_credits_test.go | 2 +- .../openai_quota_reset_credits_test.go | 1 + .../service/openai_quota_spark_window_test.go | 7 +- .../service/openai_referral_service_test.go | 3 +- 10 files changed, 211 insertions(+), 12 deletions(-) diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index 572c82591d07..34abcb9b3b87 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -1683,8 +1683,12 @@ func (r *accountRepository) UpdateGrokOAuthCredentialsIfUnchanged( return r.updateOAuthCredentialsIfUnchanged(ctx, id, service.PlatformGrok, expectedCredentials, expectedProxyID, credentials) } -// UpdateOpenAIOAuthCredentialsIfUnchanged uses the same atomic credential and -// scheduler-outbox boundary, scoped specifically to an OpenAI OAuth row. +const openAIOAuthRefreshAdminMetadataKeysSQL = "ARRAY['model_mapping', 'compact_model_mapping', 'intercept_warmup_requests']::text[]" + +// UpdateOpenAIOAuthCredentialsIfUnchanged compares all identity/credential fields +// and the proxy, while merging the three known admin-owned settings from the +// current row. A metadata edit must neither discard a rotated refresh token nor +// be overwritten by settings captured before the provider call. func (r *accountRepository) UpdateOpenAIOAuthCredentialsIfUnchanged( ctx context.Context, id int64, @@ -1692,7 +1696,57 @@ func (r *accountRepository) UpdateOpenAIOAuthCredentialsIfUnchanged( expectedProxyID *int64, credentials map[string]any, ) (bool, error) { - return r.updateOAuthCredentialsIfUnchanged(ctx, id, service.PlatformOpenAI, expectedCredentials, expectedProxyID, credentials) + if r == nil || r.sql == nil { + return false, errors.New("account repository SQL executor is not configured") + } + expectedJSON, err := json.Marshal(normalizeJSONMap(expectedCredentials)) + if err != nil { + return false, err + } + credentialsJSON, err := json.Marshal(normalizeJSONMap(credentials)) + if err != nil { + return false, err + } + result, err := r.sql.ExecContext(ctx, ` + WITH updated AS ( + UPDATE accounts AS a + SET credentials = ($1::jsonb - `+openAIOAuthRefreshAdminMetadataKeysSQL+`) || COALESCE( + (SELECT jsonb_object_agg(metadata.key, metadata.value) + FROM jsonb_each(a.credentials) AS metadata + WHERE metadata.key = ANY (`+openAIOAuthRefreshAdminMetadataKeysSQL+`)), + '{}'::jsonb), + updated_at = NOW() + WHERE a.id = $2 + AND a.deleted_at IS NULL + AND a.platform = $3 + AND a.type = $4 + AND (a.credentials - `+openAIOAuthRefreshAdminMetadataKeysSQL+`) = ($5::jsonb - `+openAIOAuthRefreshAdminMetadataKeysSQL+`) + AND a.proxy_id IS NOT DISTINCT FROM $6 + RETURNING a.id + ) + INSERT INTO scheduler_outbox (event_type, account_id, group_id, payload) + SELECT $7, updated.id, NULL, NULL FROM updated + `, + string(credentialsJSON), + id, + service.PlatformOpenAI, + service.AccountTypeOAuth, + string(expectedJSON), + expectedProxyID, + service.SchedulerOutboxEventAccountChanged, + ) + if err != nil { + return false, err + } + rowsAffected, err := result.RowsAffected() + if err != nil { + return false, err + } + if rowsAffected == 0 { + return false, nil + } + r.syncSchedulerAccountSnapshotDetached(ctx, id) + return true, nil } func (r *accountRepository) updateOAuthCredentialsIfUnchanged( diff --git a/backend/internal/repository/account_repo_openai_identity_test.go b/backend/internal/repository/account_repo_openai_identity_test.go index 014bec7ae364..9c5fc4fb420c 100644 --- a/backend/internal/repository/account_repo_openai_identity_test.go +++ b/backend/internal/repository/account_repo_openai_identity_test.go @@ -30,7 +30,19 @@ func TestAccountRepository_OAuthIdentityCASPlatformAndAtomicOutbox(t *testing.T) require.Contains(t, query, "WITH updated AS") require.Contains(t, query, "a.platform = $3") require.Contains(t, query, "a.type = $4") - require.Contains(t, query, "a.credentials = $5::jsonb") + if platform == service.PlatformOpenAI { + metadataKeys := "ARRAY['model_mapping', 'compact_model_mapping', 'intercept_warmup_requests']::text[]" + require.Contains(t, query, "(a.credentials - "+metadataKeys+") = ($5::jsonb - "+metadataKeys+")") + require.Contains(t, query, "SET credentials = ($1::jsonb - "+metadataKeys+") || COALESCE(") + require.Contains(t, query, "jsonb_object_agg(metadata.key, metadata.value)") + require.Contains(t, query, "FROM jsonb_each(a.credentials) AS metadata") + require.Contains(t, query, "metadata.key = ANY ("+metadataKeys+")") + require.NotContains(t, query, "jsonb_strip_nulls", "explicit admin null values must survive") + } else { + require.Contains(t, query, "a.credentials = $5::jsonb") + require.NotContains(t, query, "model_mapping", "Grok must retain its full credential comparison") + require.NotContains(t, query, "jsonb_object_agg") + } require.Contains(t, query, "a.proxy_id IS NOT DISTINCT FROM $6") require.Contains(t, query, "INSERT INTO scheduler_outbox") require.NotContains(t, query, "SET status") diff --git a/backend/internal/service/account_usage_service_spark_shadow_test.go b/backend/internal/service/account_usage_service_spark_shadow_test.go index 442ac6b4b0d5..9a15cb6b5f62 100644 --- a/backend/internal/service/account_usage_service_spark_shadow_test.go +++ b/backend/internal/service/account_usage_service_spark_shadow_test.go @@ -69,6 +69,7 @@ func TestGetOpenAIUsage_SparkShadow_WritesExtraAndReturnsNonEmptyWindows(t *test Status: StatusActive, Credentials: map[string]any{ "chatgpt_account_id": "org-spark-parent", + "access_token": "fake-access-token", }, } diff --git a/backend/internal/service/oauth_refresh_api.go b/backend/internal/service/oauth_refresh_api.go index 54d16f23604c..92e2a3576db2 100644 --- a/backend/internal/service/oauth_refresh_api.go +++ b/backend/internal/service/oauth_refresh_api.go @@ -36,6 +36,8 @@ type GrokOAuthRefreshSuccessRepository interface { // OpenAIOAuthRefreshSuccessRepository keeps a completed token rotation from // overwriting credentials or a proxy selected by a concurrent admin edit. +// Known admin-owned model mappings and warmup settings are preserved from the +// current row without making a metadata-only edit discard the rotation. type OpenAIOAuthRefreshSuccessRepository interface { UpdateOpenAIOAuthCredentialsIfUnchanged( ctx context.Context, @@ -402,6 +404,9 @@ func (api *OAuthRefreshAPI) RefreshIfNeeded( // Do not apply the old refresh's post-success actions to that edit. return &OAuthRefreshResult{Account: durableAccount}, nil } + // The atomic update preserves current admin-owned metadata. Return + // that durable document, not the settings captured before refresh. + newCredentials = shallowCopyMap(durableAccount.Credentials) freshAccount = durableAccount } else if updateErr := persistAccountCredentials(ctx, api.accountRepo, freshAccount, newCredentials); updateErr != nil { slog.Error("oauth_refresh_update_failed", diff --git a/backend/internal/service/oauth_refresh_api_test.go b/backend/internal/service/oauth_refresh_api_test.go index 4bb5587bfba4..415ed33d0b38 100644 --- a/backend/internal/service/oauth_refresh_api_test.go +++ b/backend/internal/service/oauth_refresh_api_test.go @@ -118,13 +118,20 @@ func (r *refreshAPIAccountRepo) UpdateOpenAIOAuthCredentialsIfUnchanged( return false, r.updateErr } if r.account == nil || r.account.ID != id || !r.account.IsOpenAIOAuth() || - !reflect.DeepEqual(r.account.Credentials, expectedCredentials) || + !reflect.DeepEqual(openAIIdentityTestComparableCredentials(r.account.Credentials), openAIIdentityTestComparableCredentials(expectedCredentials)) || !reflect.DeepEqual(r.account.ProxyID, expectedProxyID) { return false, nil } r.updateCalls++ r.updateCredentialsCalls++ - r.account.Credentials = shallowCopyMap(credentials) + merged := shallowCopyMap(credentials) + for _, key := range openAIIdentityTestAdminMetadataKeys { + delete(merged, key) + if value, exists := r.account.Credentials[key]; exists { + merged[key] = value + } + } + r.account.Credentials = merged return true, nil } diff --git a/backend/internal/service/openai_identity_mutation_test.go b/backend/internal/service/openai_identity_mutation_test.go index 789e00fa91ea..a5eb65f46435 100644 --- a/backend/internal/service/openai_identity_mutation_test.go +++ b/backend/internal/service/openai_identity_mutation_test.go @@ -63,6 +63,16 @@ func (r *openAIIdentityMutationRepo) UpdateCredentials(_ context.Context, _ int6 return nil } +var openAIIdentityTestAdminMetadataKeys = []string{"model_mapping", "compact_model_mapping", "intercept_warmup_requests"} + +func openAIIdentityTestComparableCredentials(credentials map[string]any) map[string]any { + comparable := shallowCopyMap(credentials) + for _, key := range openAIIdentityTestAdminMetadataKeys { + delete(comparable, key) + } + return comparable +} + func (r *openAIIdentityMutationRepo) UpdateOpenAIOAuthCredentialsIfUnchanged( _ context.Context, id int64, @@ -77,12 +87,19 @@ func (r *openAIIdentityMutationRepo) UpdateOpenAIOAuthCredentialsIfUnchanged( return false, err } if r.account.ID != id || !r.account.IsOpenAIOAuth() || - !reflect.DeepEqual(r.account.Credentials, expectedCredentials) || + !reflect.DeepEqual(openAIIdentityTestComparableCredentials(r.account.Credentials), openAIIdentityTestComparableCredentials(expectedCredentials)) || !reflect.DeepEqual(r.account.ProxyID, expectedProxyID) { r.mu.Unlock() return false, nil } - r.account.Credentials = shallowCopyMap(credentials) + merged := shallowCopyMap(credentials) + for _, key := range openAIIdentityTestAdminMetadataKeys { + delete(merged, key) + if value, exists := r.account.Credentials[key]; exists { + merged[key] = value + } + } + r.account.Credentials = merged afterCAS := r.afterCAS r.afterCAS = nil r.mu.Unlock() @@ -254,6 +271,104 @@ func TestOpenAIIdentity_AdminCredentialUpdateUsesCurrentTokenAndWorkspace(t *tes } func TestOpenAIIdentity_LateRefreshCannotOverwriteAdminCredentials(t *testing.T) { + t.Run("concurrent_non_auth_metadata_preserves_rotation", func(t *testing.T) { + for _, edit := range []string{"add", "replace", "delete", "null"} { + t.Run(edit, func(t *testing.T) { + ctx := context.Background() + repo := newOpenAIIdentityMutationRepo() + repo.account.Credentials["_token_version"] = int64(17) + if edit != "add" { + repo.account.Credentials["model_mapping"] = map[string]any{"old-alias": "gpt-5.5"} + repo.account.Credentials["compact_model_mapping"] = map[string]any{"old-compact": "gpt-5.5"} + repo.account.Credentials["intercept_warmup_requests"] = false + } + selected, err := repo.GetByID(ctx, repo.account.ID) + require.NoError(t, err) + rotated := shallowCopyMap(selected.Credentials) + rotated["access_token"] = "synthetic-rotated-token-a" + rotated["refresh_token"] = "synthetic-rotated-refresh-a" + var adminCredentials map[string]any + admin := &adminServiceImpl{accountRepo: repo} + executor := &refreshAPIExecutorStub{ + needsRefresh: true, + credentials: rotated, + onRefresh: func() { + // The real full-object edit preserves A's credentials and + // changes only admin-owned settings while A's refresh runs. + current, err := repo.GetByID(ctx, selected.ID) + require.NoError(t, err) + edited := shallowCopyMap(current.Credentials) + switch edit { + case "add", "replace": + // Replace whole maps; mutating a shared nested map would + // alter the refresh snapshot and hide this interleaving. + edited["model_mapping"] = map[string]any{"admin-alias": "gpt-5.5"} + edited["compact_model_mapping"] = map[string]any{"admin-compact": "gpt-5.5"} + edited["intercept_warmup_requests"] = true + case "delete": + for _, key := range openAIIdentityTestAdminMetadataKeys { + delete(edited, key) + } + case "null": + for _, key := range openAIIdentityTestAdminMetadataKeys { + edited[key] = nil + } + } + updated, err := admin.UpdateAccount(ctx, selected.ID, &UpdateAccountInput{Credentials: edited}) + require.NoError(t, err) + adminCredentials = shallowCopyMap(updated.Credentials) + }, + } + result, err := NewOAuthRefreshAPI(repo, nil).RefreshIfNeeded(ctx, selected, executor, time.Minute) + require.NoError(t, err) + require.True(t, result.Refreshed, "a metadata edit must not discard a completed token rotation") + durable, err := repo.GetByID(ctx, selected.ID) + require.NoError(t, err) + require.Equal(t, "synthetic-rotated-token-a", durable.GetOpenAIAccessToken()) + require.Equal(t, "synthetic-rotated-refresh-a", durable.GetOpenAIRefreshToken()) + require.Equal(t, "synthetic-workspace-a", durable.GetChatGPTAccountID()) + require.Greater(t, durable.GetCredentialAsInt64("_token_version"), int64(17)) + for _, key := range openAIIdentityTestAdminMetadataKeys { + expectedValue, expectedExists := adminCredentials[key] + for _, actual := range []map[string]any{durable.Credentials, result.Account.Credentials, result.NewCredentials} { + actualValue, actualExists := actual[key] + require.Equal(t, expectedExists, actualExists, "preserve admin key presence: %s", key) + require.Equal(t, expectedValue, actualValue, "preserve admin key value: %s", key) + } + } + }) + } + }) + + t.Run("same_workspace_admin_token_change_still_wins", func(t *testing.T) { + ctx := context.Background() + repo := newOpenAIIdentityMutationRepo() + selected, err := repo.GetByID(ctx, repo.account.ID) + require.NoError(t, err) + rotated := shallowCopyMap(selected.Credentials) + rotated["access_token"] = "synthetic-late-token-a" + rotated["refresh_token"] = "synthetic-late-refresh-a" + admin := &adminServiceImpl{accountRepo: repo} + executor := &refreshAPIExecutorStub{ + needsRefresh: true, + credentials: rotated, + onRefresh: func() { + edited := shallowCopyMap(selected.Credentials) + edited["access_token"] = "synthetic-admin-token-a" + edited["refresh_token"] = "synthetic-admin-refresh-a" + _, err := admin.UpdateAccount(ctx, selected.ID, &UpdateAccountInput{Credentials: edited}) + require.NoError(t, err) + }, + } + result, err := NewOAuthRefreshAPI(repo, nil).RefreshIfNeeded(ctx, selected, executor, time.Minute) + require.NoError(t, err) + require.False(t, result.Refreshed) + require.Nil(t, result.NewCredentials) + require.Equal(t, "synthetic-workspace-a", result.Account.GetChatGPTAccountID()) + require.Equal(t, "synthetic-admin-token-a", result.Account.GetOpenAIAccessToken()) + require.Equal(t, "synthetic-admin-refresh-a", result.Account.GetOpenAIRefreshToken()) + }) + t.Run("admin_reauthorization_before_cas", func(t *testing.T) { ctx := context.Background() repo := newOpenAIIdentityMutationRepo() diff --git a/backend/internal/service/openai_quota_credits_test.go b/backend/internal/service/openai_quota_credits_test.go index 37f2a613b7e3..25bb0a688356 100644 --- a/backend/internal/service/openai_quota_credits_test.go +++ b/backend/internal/service/openai_quota_credits_test.go @@ -26,7 +26,7 @@ func TestQueryUsageCodexCredits(t *testing.T) { } { t.Run(tc.name, func(t *testing.T) { account := &Account{ID: 100, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, - Credentials: map[string]any{"chatgpt_account_id": "test-workspace"}} + Credentials: map[string]any{"chatgpt_account_id": "test-workspace", "access_token": "test-token"}} repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{100: account}} tokens := &stubQuotaTokenCache{tokens: map[string]string{OpenAITokenCacheKey(account): "test-token"}} var paths []string diff --git a/backend/internal/service/openai_quota_reset_credits_test.go b/backend/internal/service/openai_quota_reset_credits_test.go index 02edef3e4bb5..fa33d4b38485 100644 --- a/backend/internal/service/openai_quota_reset_credits_test.go +++ b/backend/internal/service/openai_quota_reset_credits_test.go @@ -152,6 +152,7 @@ func TestQueryUsageResetCreditCountPrecedence(t *testing.T) { Status: StatusActive, Credentials: map[string]any{ "chatgpt_account_id": "org-parent123", + "access_token": "fake-token", }, } repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{100: account}} diff --git a/backend/internal/service/openai_quota_spark_window_test.go b/backend/internal/service/openai_quota_spark_window_test.go index 8bda444bf93d..08fe722875c0 100644 --- a/backend/internal/service/openai_quota_spark_window_test.go +++ b/backend/internal/service/openai_quota_spark_window_test.go @@ -194,7 +194,7 @@ func TestResetCreditShadowRejected(t *testing.T) { func TestResetCreditTargetedSendsStableCreditAndRedeemIDs(t *testing.T) { account := &Account{ ID: 203, Platform: PlatformOpenAI, Type: AccountTypeOAuth, - Credentials: map[string]any{"chatgpt_account_id": "account-targeted"}, + Credentials: map[string]any{"chatgpt_account_id": "account-targeted", "access_token": "fake-token"}, } repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{account.ID: account}} tokenCache := &stubQuotaTokenCache{tokens: map[string]string{OpenAITokenCacheKey(account): "fake-token"}} @@ -363,6 +363,7 @@ func TestPrepareUpstreamCallShadowResolve(t *testing.T) { Status: StatusActive, Credentials: map[string]any{ "chatgpt_account_id": "org-parent123", + "access_token": "fake-access-token", }, } repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{200: shadow, 100: parent}} @@ -536,6 +537,7 @@ func TestQueryUsageIncludesResetCreditExpirations_EndToEnd(t *testing.T) { Status: StatusActive, Credentials: map[string]any{ "chatgpt_account_id": "org-parent123", + "access_token": "fake-token", }, } repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{100: account}} @@ -599,6 +601,7 @@ func TestQueryUsageResetCreditDetails401NonFatal(t *testing.T) { Status: StatusActive, Credentials: map[string]any{ "chatgpt_account_id": "org-parent123", + "access_token": "fake-token", }, } repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{100: account}} @@ -756,7 +759,7 @@ func TestQueryUsageShadowResolve_EndToEnd(t *testing.T) { } parent := &Account{ ID: 100, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, - Credentials: map[string]any{"chatgpt_account_id": "org-e2e-parent"}, + Credentials: map[string]any{"chatgpt_account_id": "org-e2e-parent", "access_token": "fake-token-e2e"}, } repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{200: shadow, 100: parent}} diff --git a/backend/internal/service/openai_referral_service_test.go b/backend/internal/service/openai_referral_service_test.go index 579a01a87b03..e5c2727798d8 100644 --- a/backend/internal/service/openai_referral_service_test.go +++ b/backend/internal/service/openai_referral_service_test.go @@ -32,8 +32,9 @@ func (s *referralClientStub) SendInvite(_ context.Context, call OpenAIReferralCa func referralTestService(t *testing.T, plan string, client OpenAIReferralClient) (*OpenAIQuotaService, *stubQuotaAccountRepo) { t.Helper() + // The cached synthetic bearer must belong to the persisted test account. a := &Account{ID: 100, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, - Credentials: map[string]any{"chatgpt_account_id": "workspace-test", "plan_type": plan}} + Credentials: map[string]any{"chatgpt_account_id": "workspace-test", "plan_type": plan, "access_token": "test-token"}} repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{100: a}} tokens := &stubQuotaTokenCache{tokens: map[string]string{OpenAITokenCacheKey(a): "test-token"}} factory := func(string) (*req.Client, error) {