mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-30 16:23:31 +00:00
fix(openai): harden referral outcomes and isolate HTTP transport
This commit is contained in:
@@ -199,14 +199,14 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
antigravityQuotaFetcher := service.NewAntigravityQuotaFetcher(proxyRepository, configConfig)
|
||||
grokQuotaFetcher := service.NewGrokQuotaFetcher()
|
||||
grokQuotaService := service.ProvideGrokQuotaService(accountRepository, proxyRepository, grokTokenProvider, httpUpstream, configConfig, usageLogRepository, settingService)
|
||||
openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory, openAIGatewayService)
|
||||
openAIReferralClient := repository.NewOpenAIReferralClient(privacyClientFactory)
|
||||
openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory, openAIReferralClient, openAIGatewayService)
|
||||
usageCache := service.NewUsageCache()
|
||||
accountUsageService := service.ProvideAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, grokQuotaService, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService, openAIGatewayService)
|
||||
pluginRepository := repository.NewPluginRepository(db)
|
||||
pluginHostInfo := providePluginHostInfo(buildInfo)
|
||||
pluginKVStore := repository.NewPluginKVStore(redisClient)
|
||||
pluginManager := service.NewPluginManager(pluginRepository, secretEncryptor, configConfig, pluginHostInfo, pluginKVStore)
|
||||
pluginManager.SetAccountDirectory(openAIGatewayService)
|
||||
pluginManager := service.ProvidePluginManager(pluginRepository, secretEncryptor, configConfig, pluginHostInfo, pluginKVStore, openAIGatewayService)
|
||||
accountTestService := service.ProvideAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService, openAIGatewayService, settingService, pluginManager)
|
||||
crsSyncService := service.NewCRSSyncService(accountRepository, proxyRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig)
|
||||
accountHandler := admin.ProvideAccountHandler(configConfig, adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator, grokQuotaService)
|
||||
|
||||
@@ -81,13 +81,18 @@ func (h *OpenAIOAuthHandler) SendReferralInvite(c *gin.Context) {
|
||||
}
|
||||
// The email is already sent. Refresh failure must not turn it into a failed
|
||||
// submission and encourage a duplicate send, even if the browser disconnects.
|
||||
ctx, cancel := context.WithTimeout(context.WithoutCancel(c.Request.Context()), 8*time.Second)
|
||||
defer cancel()
|
||||
eligibility, refreshErr := h.referralService.QueryReferralEligibility(ctx, id)
|
||||
baseCtx := context.WithoutCancel(c.Request.Context())
|
||||
refreshCtx, cancelRefresh := context.WithTimeout(baseCtx, 8*time.Second)
|
||||
eligibility, refreshErr := h.referralService.QueryReferralEligibility(refreshCtx, id)
|
||||
cancelRefresh()
|
||||
if refreshErr != nil {
|
||||
eligibility = nil
|
||||
}
|
||||
cacheErr := h.referralService.CacheReferralSnapshot(ctx, id, eligibility)
|
||||
// Refresh may exhaust its entire deadline. Give persistence (including
|
||||
// invalidating a stale snapshot) a fresh, independent deadline.
|
||||
cacheCtx, cancelCache := context.WithTimeout(baseCtx, 3*time.Second)
|
||||
defer cancelCache()
|
||||
cacheErr := h.referralService.CacheReferralSnapshot(cacheCtx, id, eligibility)
|
||||
response.Success(c, openAIReferralSendResponse{
|
||||
OpenAIReferralSendResult: *result,
|
||||
openAIReferralRefreshResponse: openAIReferralRefreshResponse{Eligibility: eligibility, CachePersisted: cacheErr == nil},
|
||||
|
||||
@@ -10,6 +10,8 @@ import (
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"testing/synctest"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -17,6 +19,8 @@ import (
|
||||
)
|
||||
|
||||
type referralHandlerStub struct {
|
||||
query func(context.Context) (*service.OpenAIReferralEligibility, error)
|
||||
cache func(context.Context, *service.OpenAIReferralEligibility) error
|
||||
eligibility *service.OpenAIReferralEligibility
|
||||
queryErr, cacheErr error
|
||||
input service.OpenAIReferralSendRequest
|
||||
@@ -24,10 +28,16 @@ type referralHandlerStub struct {
|
||||
sends int
|
||||
}
|
||||
|
||||
func (s *referralHandlerStub) QueryReferralEligibility(context.Context, int64) (*service.OpenAIReferralEligibility, error) {
|
||||
func (s *referralHandlerStub) QueryReferralEligibility(ctx context.Context, _ int64) (*service.OpenAIReferralEligibility, error) {
|
||||
if s.query != nil {
|
||||
return s.query(ctx)
|
||||
}
|
||||
return s.eligibility, s.queryErr
|
||||
}
|
||||
func (s *referralHandlerStub) CacheReferralSnapshot(_ context.Context, _ int64, e *service.OpenAIReferralEligibility) error {
|
||||
func (s *referralHandlerStub) CacheReferralSnapshot(ctx context.Context, _ int64, e *service.OpenAIReferralEligibility) error {
|
||||
if s.cache != nil {
|
||||
return s.cache(ctx, e)
|
||||
}
|
||||
s.cached = e
|
||||
return s.cacheErr
|
||||
}
|
||||
@@ -88,3 +98,36 @@ func TestOpenAIReferralHandlerRejectsInvalidInput(t *testing.T) {
|
||||
}
|
||||
require.Zero(t, stub.sends)
|
||||
}
|
||||
|
||||
func TestOpenAIReferralHandlerInvalidatesCacheAfterRefreshDeadline(t *testing.T) {
|
||||
synctest.Test(t, func(t *testing.T) {
|
||||
cached := false
|
||||
stub := &referralHandlerStub{
|
||||
query: func(ctx context.Context) (*service.OpenAIReferralEligibility, error) {
|
||||
<-ctx.Done()
|
||||
require.ErrorIs(t, ctx.Err(), context.DeadlineExceeded)
|
||||
return nil, ctx.Err()
|
||||
},
|
||||
cache: func(ctx context.Context, e *service.OpenAIReferralEligibility) error {
|
||||
require.NoError(t, ctx.Err(), "cache invalidation needs a fresh deadline")
|
||||
deadline, ok := ctx.Deadline()
|
||||
require.True(t, ok)
|
||||
require.Equal(t, 3*time.Second, time.Until(deadline))
|
||||
require.Nil(t, e, "expired count must be cleared")
|
||||
cached = true
|
||||
return nil
|
||||
},
|
||||
}
|
||||
rec := referralHandlerRequest(t, stub, "/100/referrals/invite", `{"email":"friend@example.com","program_id":"codex_referral_consumer","confirmed":true}`)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var body struct {
|
||||
Data openAIReferralSendResponse `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body))
|
||||
require.True(t, cached)
|
||||
require.True(t, body.Data.Sent)
|
||||
require.True(t, body.Data.RefreshFailed)
|
||||
require.True(t, body.Data.CachePersisted)
|
||||
require.Equal(t, 1, stub.sends)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -65,6 +65,8 @@ var schedulerNeutralExtraKeyPrefixes = []string{
|
||||
|
||||
var schedulerNeutralExtraKeys = map[string]struct{}{
|
||||
"codex_usage_updated_at": {},
|
||||
"codex_credits_snapshot": {},
|
||||
"codex_referral_snapshot": {},
|
||||
"grok_billing_snapshot": {},
|
||||
"session_window_utilization": {},
|
||||
}
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"regexp"
|
||||
"testing"
|
||||
|
||||
"entgo.io/ent/dialect"
|
||||
entsql "entgo.io/ent/dialect/sql"
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestUpdateExtraCodexDisplaySnapshotsAvoidSchedulerOutbox(t *testing.T) {
|
||||
for _, key := range []string{"codex_credits_snapshot", "codex_referral_snapshot"} {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
value any
|
||||
schedulingChange bool
|
||||
}{
|
||||
{"refresh", map[string]any{"fetched_at": 1770000000}, false},
|
||||
{"invalidate", nil, false},
|
||||
{"mixed scheduling update", nil, true},
|
||||
} {
|
||||
t.Run(key+"/"+tc.name, func(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
client := dbent.NewClient(dbent.Driver(entsql.OpenDB(dialect.Postgres, db)))
|
||||
t.Cleanup(func() { _ = client.Close() })
|
||||
repo := newAccountRepositoryWithSQL(client, db, nil)
|
||||
updates := map[string]any{key: tc.value}
|
||||
if tc.schedulingChange {
|
||||
updates["openai_compact_enabled"] = true
|
||||
mock.ExpectBegin()
|
||||
}
|
||||
payload, err := json.Marshal(updates)
|
||||
require.NoError(t, err)
|
||||
mock.ExpectExec(regexp.QuoteMeta("UPDATE accounts SET extra = COALESCE(extra, '{}'::jsonb) || $1::jsonb, updated_at = NOW() WHERE id = $2 AND deleted_at IS NULL")).
|
||||
WithArgs(string(payload), int64(27)).WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
if tc.schedulingChange {
|
||||
mock.ExpectExec(regexp.QuoteMeta("INSERT INTO scheduler_outbox")).
|
||||
WithArgs(service.SchedulerOutboxEventAccountChanged, int64(27), nil, nil, sqlmock.AnyArg()).WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectCommit()
|
||||
}
|
||||
// Snapshot-only changes must perform just the UPDATE, without even a
|
||||
// transaction or account_changed event that would rebuild scheduler buckets.
|
||||
require.NoError(t, repo.UpdateExtra(context.Background(), 27, updates))
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/imroc/req/v3"
|
||||
)
|
||||
|
||||
// Desktop 26.908.40834, app-primary-44ec287874b7.js: SIt, EIt, PIt, FIt.
|
||||
// Referrals use /backend-api/referrals, not /backend-api/wham.
|
||||
const openAIReferralURL = "https://chatgpt.com/backend-api/referrals/invite"
|
||||
const openAIReferralEntrypoint = "persistent"
|
||||
|
||||
type openAIReferralClient struct {
|
||||
clientFactory service.PrivacyClientFactory
|
||||
baseURL string
|
||||
}
|
||||
|
||||
func NewOpenAIReferralClient(factory service.PrivacyClientFactory) service.OpenAIReferralClient {
|
||||
return &openAIReferralClient{clientFactory: factory, baseURL: openAIReferralURL}
|
||||
}
|
||||
|
||||
func (c *openAIReferralClient) request(ctx context.Context, call service.OpenAIReferralCall) (*req.Request, error) {
|
||||
client, err := c.clientFactory(call.ProxyURL)
|
||||
if err != nil {
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_CLIENT_ERROR", "failed to build referral client")
|
||||
}
|
||||
// Sending has no documented idempotency key. Override any factory retries.
|
||||
return client.R().SetContext(ctx).SetHeaders(call.Headers).SetRetryCount(0), nil
|
||||
}
|
||||
|
||||
func (c *openAIReferralClient) QueryEligibility(ctx context.Context, call service.OpenAIReferralCall) (*service.OpenAIReferralEligibility, error) {
|
||||
r, err := c.request(ctx, call)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resp, err := r.SetQueryParams(map[string]string{"program_id": call.ProgramID, "entrypoint": openAIReferralEntrypoint}).Get(c.baseURL + "/eligibility")
|
||||
if err != nil {
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_QUERY_FAILED", "failed to query invitation eligibility")
|
||||
}
|
||||
if !resp.IsSuccessState() {
|
||||
return nil, referralHTTPError(resp.StatusCode)
|
||||
}
|
||||
var result *service.OpenAIReferralEligibility
|
||||
if err := json.Unmarshal(resp.Bytes(), &result); err != nil || result == nil {
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_INVALID_RESPONSE", "invalid invitation eligibility response")
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (c *openAIReferralClient) SendInvite(ctx context.Context, call service.OpenAIReferralCall, email string) error {
|
||||
r, err := c.request(ctx, call)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
resp, err := r.SetBody(map[string]any{
|
||||
"program_id": call.ProgramID, "entrypoint": openAIReferralEntrypoint, "emails": []string{email},
|
||||
}).Post(c.baseURL)
|
||||
if err != nil {
|
||||
return referralSendUnknown()
|
||||
}
|
||||
if !resp.IsSuccessState() {
|
||||
return referralHTTPError(resp.StatusCode)
|
||||
}
|
||||
var payload struct {
|
||||
Invites []json.RawMessage `json:"invites"`
|
||||
}
|
||||
if err := json.Unmarshal(resp.Bytes(), &payload); err != nil || len(payload.Invites) != 1 || string(payload.Invites[0]) == "null" {
|
||||
return referralSendUnknown()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func referralSendUnknown() error {
|
||||
return infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_SEND_UNKNOWN", "invitation outcome is unknown; check the invitation status in Codex before sending again")
|
||||
}
|
||||
|
||||
// Upstream bodies can contain personal data or credentials. Return only stable
|
||||
// error codes distinguishing validation failures, duplicates and limits.
|
||||
func referralHTTPError(status int) error {
|
||||
switch status {
|
||||
case http.StatusBadRequest, http.StatusUnprocessableEntity:
|
||||
return infraerrors.New(http.StatusBadRequest, "OPENAI_REFERRAL_REJECTED", "the invitation request was rejected; check the email and eligibility")
|
||||
case http.StatusUnauthorized:
|
||||
return infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_AUTH_ERROR", "upstream authentication failed; refresh the account credentials")
|
||||
case http.StatusForbidden:
|
||||
return infraerrors.New(http.StatusForbidden, "OPENAI_REFERRAL_FORBIDDEN", "this account is not eligible for invitations")
|
||||
case http.StatusConflict:
|
||||
return infraerrors.New(http.StatusConflict, "OPENAI_REFERRAL_ALREADY_EXISTS", "an invitation already exists for this recipient")
|
||||
case http.StatusTooManyRequests:
|
||||
return infraerrors.New(http.StatusTooManyRequests, "OPENAI_REFERRAL_RATE_LIMITED", "invitation limit reached; try again later")
|
||||
default:
|
||||
return infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_UPSTREAM_ERROR", "invitation service is unavailable")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/imroc/req/v3"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func referralTestClient(t *testing.T, handler http.HandlerFunc) *openAIReferralClient {
|
||||
t.Helper()
|
||||
srv := httptest.NewServer(handler)
|
||||
t.Cleanup(srv.Close)
|
||||
adapter := NewOpenAIReferralClient(func(proxy string) (*req.Client, error) {
|
||||
require.Equal(t, "test-proxy", proxy)
|
||||
// A shared factory may have retries; the adapter must override them.
|
||||
return req.C().SetCommonRetryCount(2).SetCommonRetryCondition(func(_ *req.Response, _ error) bool { return true }), nil
|
||||
})
|
||||
client, ok := adapter.(*openAIReferralClient)
|
||||
require.True(t, ok)
|
||||
client.baseURL = srv.URL + "/backend-api/referrals/invite"
|
||||
return client
|
||||
}
|
||||
|
||||
func referralTestCall() service.OpenAIReferralCall {
|
||||
return service.OpenAIReferralCall{ProxyURL: "test-proxy", ProgramID: "codex_referral_consumer",
|
||||
Headers: map[string]string{"Authorization": "Bearer test-token", "ChatGPT-Account-ID": "workspace-test"}}
|
||||
}
|
||||
|
||||
func TestOpenAIReferralClientProtocol(t *testing.T) {
|
||||
var gets, posts atomic.Int32
|
||||
client := referralTestClient(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
require.Equal(t, "Bearer test-token", r.Header.Get("Authorization"))
|
||||
require.Equal(t, "workspace-test", r.Header.Get("ChatGPT-Account-ID"))
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if r.Method == http.MethodGet {
|
||||
gets.Add(1)
|
||||
require.Equal(t, "/backend-api/referrals/invite/eligibility", r.URL.Path)
|
||||
require.Equal(t, "codex_referral_consumer", r.URL.Query().Get("program_id"))
|
||||
require.Equal(t, "persistent", r.URL.Query().Get("entrypoint"))
|
||||
_, _ = w.Write([]byte(`{"should_show":true,"remaining_send_capacity":8,"remaining_reward_capacity":3,"grants":[{"grant_type":"rate_limit_reset_credit","amount":1,"recipient":"referrer"}],"rules":["Offer rule"]}`))
|
||||
return
|
||||
}
|
||||
posts.Add(1)
|
||||
require.Equal(t, http.MethodPost, r.Method)
|
||||
require.Equal(t, "/backend-api/referrals/invite", r.URL.Path)
|
||||
var body map[string]any
|
||||
require.NoError(t, json.NewDecoder(r.Body).Decode(&body))
|
||||
require.Equal(t, map[string]any{"program_id": "codex_referral_consumer", "entrypoint": "persistent", "emails": []any{"friend@example.com"}}, body)
|
||||
_, _ = w.Write([]byte(`{"invites":[{"referral_id":"test-invite"}]}`))
|
||||
})
|
||||
result, err := client.QueryEligibility(context.Background(), referralTestCall())
|
||||
require.NoError(t, err)
|
||||
require.True(t, result.ShouldShow)
|
||||
require.Equal(t, 8, *result.RemainingSendCapacity)
|
||||
require.Equal(t, 3, *result.RemainingRewardCapacity)
|
||||
require.Equal(t, "rate_limit_reset_credit", result.Grants[0].GrantType)
|
||||
require.Equal(t, []string{"Offer rule"}, result.Rules)
|
||||
require.NoError(t, client.SendInvite(context.Background(), referralTestCall(), "friend@example.com"))
|
||||
require.EqualValues(t, 1, gets.Load())
|
||||
require.EqualValues(t, 1, posts.Load())
|
||||
}
|
||||
|
||||
func TestOpenAIReferralClientSanitizesErrorsWithoutRetry(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
status int
|
||||
reason string
|
||||
}{
|
||||
{400, "OPENAI_REFERRAL_REJECTED"}, {422, "OPENAI_REFERRAL_REJECTED"},
|
||||
{401, "OPENAI_REFERRAL_AUTH_ERROR"}, {403, "OPENAI_REFERRAL_FORBIDDEN"},
|
||||
{409, "OPENAI_REFERRAL_ALREADY_EXISTS"}, {429, "OPENAI_REFERRAL_RATE_LIMITED"},
|
||||
{500, "OPENAI_REFERRAL_UPSTREAM_ERROR"},
|
||||
} {
|
||||
t.Run(http.StatusText(tc.status), func(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
client := referralTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
|
||||
requests.Add(1)
|
||||
w.WriteHeader(tc.status)
|
||||
_, _ = w.Write([]byte(`{"detail":"secret-token someone@example.com"}`))
|
||||
})
|
||||
_, queryErr := client.QueryEligibility(context.Background(), referralTestCall())
|
||||
sendErr := client.SendInvite(context.Background(), referralTestCall(), "friend@example.com")
|
||||
for _, err := range []error{queryErr, sendErr} {
|
||||
require.Equal(t, tc.reason, infraerrors.Reason(err))
|
||||
require.NotContains(t, err.Error(), "secret-token")
|
||||
require.NotContains(t, err.Error(), "someone@example.com")
|
||||
}
|
||||
require.EqualValues(t, 2, requests.Load(), "one query and one send, no retries")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIReferralClientInvalidResponses(t *testing.T) {
|
||||
for _, body := range []string{"null", "invalid json", `{"invites":[]}`, `{"invites":[null]}`, `{"invites":[{},{}]}`} {
|
||||
t.Run(body, func(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
client := referralTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
|
||||
requests.Add(1)
|
||||
_, _ = w.Write([]byte(body))
|
||||
})
|
||||
err := client.SendInvite(context.Background(), referralTestCall(), "friend@example.com")
|
||||
require.Equal(t, "OPENAI_REFERRAL_SEND_UNKNOWN", infraerrors.Reason(err))
|
||||
require.EqualValues(t, 1, requests.Load())
|
||||
if body == "null" || body == "invalid json" {
|
||||
_, err := client.QueryEligibility(context.Background(), referralTestCall())
|
||||
require.Equal(t, "OPENAI_REFERRAL_INVALID_RESPONSE", infraerrors.Reason(err))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIReferralClientLostResponseIsUnknownWithoutRetry(t *testing.T) {
|
||||
var requests atomic.Int32
|
||||
client := referralTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
|
||||
requests.Add(1)
|
||||
hijacker, ok := w.(http.Hijacker)
|
||||
require.True(t, ok)
|
||||
conn, _, err := hijacker.Hijack()
|
||||
require.NoError(t, err)
|
||||
_ = conn.Close() // Upstream received the request but its response was lost.
|
||||
})
|
||||
err := client.SendInvite(context.Background(), referralTestCall(), "friend@example.com")
|
||||
require.Equal(t, "OPENAI_REFERRAL_SEND_UNKNOWN", infraerrors.Reason(err))
|
||||
require.EqualValues(t, 1, requests.Load())
|
||||
}
|
||||
|
||||
func TestOpenAIReferralClientFactoryErrorIsSanitized(t *testing.T) {
|
||||
client := NewOpenAIReferralClient(func(string) (*req.Client, error) { return nil, errors.New("secret proxy credentials") })
|
||||
err := client.SendInvite(context.Background(), referralTestCall(), "friend@example.com")
|
||||
require.Equal(t, "OPENAI_REFERRAL_CLIENT_ERROR", infraerrors.Reason(err))
|
||||
require.NotContains(t, err.Error(), "secret")
|
||||
}
|
||||
@@ -162,6 +162,7 @@ var ProviderSet = wire.NewSet(
|
||||
NewClaudeOAuthClient,
|
||||
NewHTTPUpstream,
|
||||
NewOpenAIOAuthClient,
|
||||
NewOpenAIReferralClient,
|
||||
NewGrokOAuthClient,
|
||||
NewGeminiOAuthClient,
|
||||
NewGeminiCliCodeAssistClient,
|
||||
|
||||
@@ -117,7 +117,7 @@ func TestGetOpenAIUsage_SparkShadow_WritesExtraAndReturnsNonEmptyWindows(t *test
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
quotaService := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv))
|
||||
quotaService := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv), nil)
|
||||
svc := &AccountUsageService{
|
||||
accountRepo: repo,
|
||||
openAIQuotaService: quotaService,
|
||||
|
||||
@@ -44,7 +44,7 @@ func TestQueryUsageCodexCredits(t *testing.T) {
|
||||
http.Error(w, "unavailable", http.StatusServiceUnavailable)
|
||||
}))
|
||||
defer srv.Close()
|
||||
svc := NewOpenAIQuotaService(repo, nil, NewOpenAITokenProvider(repo, tokens, nil), newQuotaRedirectingFactory(srv))
|
||||
svc := NewOpenAIQuotaService(repo, nil, NewOpenAITokenProvider(repo, tokens, nil), newQuotaRedirectingFactory(srv), nil)
|
||||
usage, err := svc.QueryUsage(context.Background(), 100)
|
||||
require.NoError(t, err)
|
||||
require.Positive(t, usage.FetchedAt)
|
||||
|
||||
@@ -175,7 +175,7 @@ func TestQueryUsageResetCreditCountPrecedence(t *testing.T) {
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv))
|
||||
svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv), nil)
|
||||
usage, err := svc.QueryUsage(context.Background(), 100)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, usage)
|
||||
|
||||
@@ -136,6 +136,7 @@ type OpenAIQuotaService struct {
|
||||
proxyRepo ProxyRepository
|
||||
tokenProvider *OpenAITokenProvider
|
||||
privacyClientFactory PrivacyClientFactory
|
||||
referralClient OpenAIReferralClient
|
||||
agentIdentityTaskMu sync.Mutex
|
||||
agentIdentityWS agentIdentityWSConnectionInvalidator
|
||||
}
|
||||
@@ -148,12 +149,14 @@ func NewOpenAIQuotaService(
|
||||
proxyRepo ProxyRepository,
|
||||
tokenProvider *OpenAITokenProvider,
|
||||
privacyClientFactory PrivacyClientFactory,
|
||||
referralClient OpenAIReferralClient,
|
||||
) *OpenAIQuotaService {
|
||||
return &OpenAIQuotaService{
|
||||
accountRepo: accountRepo,
|
||||
proxyRepo: proxyRepo,
|
||||
tokenProvider: tokenProvider,
|
||||
privacyClientFactory: privacyClientFactory,
|
||||
referralClient: referralClient,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -208,7 +208,7 @@ func TestResetCreditTargetedSendsStableCreditAndRedeemIDs(t *testing.T) {
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv))
|
||||
svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv), nil)
|
||||
result, err := svc.ResetCreditTargeted(context.Background(), account.ID, "credit-123", "redeem-456")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "ok", result.Code)
|
||||
@@ -261,7 +261,7 @@ func TestResetCreditAgentIdentityUsesAssertionAndRecoversInvalidTaskOnce(t *test
|
||||
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
|
||||
|
||||
invalidator := &agentIdentityWSInvalidationRecorder{}
|
||||
svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv))
|
||||
svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv), nil)
|
||||
svc.agentIdentityWS = invalidator
|
||||
|
||||
result, err := svc.ResetCredit(context.Background(), account.ID)
|
||||
@@ -323,7 +323,7 @@ func TestResetCreditAgentIdentityReusesConcurrentlyRecoveredTask(t *testing.T) {
|
||||
openAIAgentIdentityAuthAPIBaseURL = srv.URL
|
||||
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
|
||||
|
||||
svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv))
|
||||
svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv), nil)
|
||||
result, err := svc.ResetCredit(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "ok", result.Code)
|
||||
@@ -376,7 +376,7 @@ func TestPrepareUpstreamCallShadowResolve(t *testing.T) {
|
||||
// privacyClientFactory 可以是任意合法工厂;prepareUpstreamCall 在返回前不调用它
|
||||
svc := NewOpenAIQuotaService(repo, nil, tokenProvider, func(_ string) (*req.Client, error) {
|
||||
return req.C(), nil
|
||||
})
|
||||
}, nil)
|
||||
|
||||
_, chatGPTAccountID, _, _, err := svc.prepareUpstreamCall(ctx, 200)
|
||||
require.NoError(t, err, "shadow resolve should succeed; got error: %v", err)
|
||||
@@ -414,7 +414,7 @@ func TestQueryUsageAgentIdentityUsesAssertionWithoutOAuthToken(t *testing.T) {
|
||||
_, _ = w.Write([]byte(`{"plan_type":"pro","rate_limit":{"allowed":true}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv))
|
||||
svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv), nil)
|
||||
usage, err := svc.QueryUsage(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, usage)
|
||||
@@ -468,7 +468,7 @@ func TestQueryUsageAgentIdentityRecoversInvalidTaskOnce(t *testing.T) {
|
||||
t.Cleanup(func() { openAIAgentIdentityAuthAPIBaseURL = oldBase })
|
||||
|
||||
invalidator := &agentIdentityWSInvalidationRecorder{}
|
||||
svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv))
|
||||
svc := NewOpenAIQuotaService(repo, nil, nil, newQuotaRedirectingFactory(srv), nil)
|
||||
svc.agentIdentityWS = invalidator
|
||||
usage, err := svc.QueryUsage(context.Background(), account.ID)
|
||||
require.NoError(t, err)
|
||||
@@ -564,7 +564,7 @@ func TestQueryUsageIncludesResetCreditExpirations_EndToEnd(t *testing.T) {
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv))
|
||||
svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv), nil)
|
||||
usage, err := svc.QueryUsage(ctx, 100)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, usage)
|
||||
@@ -625,7 +625,7 @@ func TestQueryUsageResetCreditDetails401NonFatal(t *testing.T) {
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv))
|
||||
svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv), nil)
|
||||
usage, err := svc.QueryUsage(ctx, 100)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, usage)
|
||||
@@ -774,7 +774,7 @@ func TestQueryUsageShadowResolve_EndToEnd(t *testing.T) {
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv))
|
||||
svc := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv), nil)
|
||||
usage, err := svc.QueryUsage(ctx, 200)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, usage)
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
package service
|
||||
|
||||
import "context"
|
||||
|
||||
// OpenAIReferralCall carries the account's refreshed authentication and proxy
|
||||
// settings to the transport adapter. Headers must never be logged or persisted.
|
||||
type OpenAIReferralCall struct {
|
||||
ProxyURL string
|
||||
Headers map[string]string
|
||||
ProgramID string
|
||||
}
|
||||
|
||||
// OpenAIReferralClient isolates invitation transport from eligibility policy.
|
||||
// QueryEligibility returns a non-nil result on success. SendInvite must not
|
||||
// retry and returns nil only when the upstream confirms one invitation.
|
||||
type OpenAIReferralClient interface {
|
||||
QueryEligibility(context.Context, OpenAIReferralCall) (*OpenAIReferralEligibility, error)
|
||||
SendInvite(context.Context, OpenAIReferralCall, string) error
|
||||
}
|
||||
@@ -2,24 +2,18 @@ package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/mail"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/imroc/req/v3"
|
||||
)
|
||||
|
||||
// Desktop 26.908.40834, app-primary-44ec287874b7.js: SIt, EIt, PIt, FIt.
|
||||
// Referrals use /backend-api/referrals, not /backend-api/wham.
|
||||
const (
|
||||
openAIReferralURL = "https://chatgpt.com/backend-api/referrals/invite"
|
||||
openAIReferralSnapshotKey = "codex_referral_snapshot"
|
||||
openAIReferralConsumer = "codex_referral_consumer"
|
||||
openAIReferralWorkspace = "codex_referral_workspace"
|
||||
openAIReferralEntrypoint = "persistent"
|
||||
)
|
||||
|
||||
type OpenAIReferralGrant struct {
|
||||
@@ -98,21 +92,19 @@ func referralProgram(a *Account) string {
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OpenAIQuotaService) referralRequest(ctx context.Context, id int64) (*req.Request, error) {
|
||||
func (s *OpenAIQuotaService) referralCall(ctx context.Context, id int64, program string) (OpenAIReferralCall, error) {
|
||||
if s.referralClient == nil {
|
||||
return OpenAIReferralCall{}, infraerrors.New(http.StatusServiceUnavailable, "OPENAI_REFERRAL_NOT_CONFIGURED", "referral service is unavailable")
|
||||
}
|
||||
token, accountID, proxy, fedRAMP, err := s.prepareUpstreamCall(ctx, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
client, err := s.privacyClientFactory(proxy)
|
||||
if err != nil {
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_CLIENT_ERROR", "failed to build referral client")
|
||||
return OpenAIReferralCall{}, err
|
||||
}
|
||||
headers, _, err := s.buildCodexQuotaHeaders(ctx, id, token, accountID, fedRAMP)
|
||||
if err != nil {
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_AUTH_ERROR", "failed to authenticate referral request")
|
||||
return OpenAIReferralCall{}, infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_AUTH_ERROR", "failed to authenticate referral request")
|
||||
}
|
||||
// Sending an invitation has no documented idempotency key. Never retry it.
|
||||
return client.R().SetContext(ctx).SetHeaders(headers).SetRetryCount(0), nil
|
||||
return OpenAIReferralCall{ProxyURL: proxy, Headers: headers, ProgramID: program}, nil
|
||||
}
|
||||
|
||||
func (s *OpenAIQuotaService) QueryReferralEligibility(ctx context.Context, id int64) (*OpenAIReferralEligibility, error) {
|
||||
@@ -122,21 +114,14 @@ func (s *OpenAIQuotaService) QueryReferralEligibility(ctx context.Context, id in
|
||||
}
|
||||
callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout)
|
||||
defer cancel()
|
||||
r, err := s.referralRequest(callCtx, id)
|
||||
program := referralProgram(account)
|
||||
call, err := s.referralCall(callCtx, id, program)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
program := referralProgram(account)
|
||||
resp, err := r.SetQueryParams(map[string]string{"program_id": program, "entrypoint": openAIReferralEntrypoint}).Get(openAIReferralURL + "/eligibility")
|
||||
result, err := s.referralClient.QueryEligibility(callCtx, call)
|
||||
if err != nil {
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_QUERY_FAILED", "failed to query invitation eligibility")
|
||||
}
|
||||
if !resp.IsSuccessState() {
|
||||
return nil, referralHTTPError(resp.StatusCode)
|
||||
}
|
||||
var result *OpenAIReferralEligibility
|
||||
if err := json.Unmarshal(resp.Bytes(), &result); err != nil || result == nil {
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_INVALID_RESPONSE", "invalid invitation eligibility response")
|
||||
return nil, err
|
||||
}
|
||||
result.ProgramID = program
|
||||
result.AvailableInvites = referralCapacity(result)
|
||||
@@ -185,43 +170,12 @@ func (s *OpenAIQuotaService) SendReferralInvite(ctx context.Context, id int64, i
|
||||
}
|
||||
callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout)
|
||||
defer cancel()
|
||||
r, err := s.referralRequest(callCtx, id)
|
||||
call, err := s.referralCall(callCtx, id, eligibility.ProgramID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
resp, err := r.SetBody(map[string]any{
|
||||
"program_id": eligibility.ProgramID, "entrypoint": openAIReferralEntrypoint, "emails": []string{email},
|
||||
}).Post(openAIReferralURL)
|
||||
if err != nil {
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_SEND_UNKNOWN", "invitation outcome is unknown; check the invitation status in Codex before sending again")
|
||||
}
|
||||
if !resp.IsSuccessState() {
|
||||
return nil, referralHTTPError(resp.StatusCode)
|
||||
}
|
||||
var payload struct {
|
||||
Invites []json.RawMessage `json:"invites"`
|
||||
}
|
||||
if err := json.Unmarshal(resp.Bytes(), &payload); err != nil || len(payload.Invites) != 1 || string(payload.Invites[0]) == "null" {
|
||||
return nil, infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_SEND_UNKNOWN", "invitation outcome is unknown; check the invitation status in Codex before sending again")
|
||||
if err := s.referralClient.SendInvite(callCtx, call, email); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &OpenAIReferralSendResult{Email: email, Sent: true}, nil
|
||||
}
|
||||
|
||||
// Avoid exposing upstream response bodies, which can contain personal data or
|
||||
// credentials. The status still distinguishes validation, duplicates and limits.
|
||||
func referralHTTPError(status int) error {
|
||||
switch status {
|
||||
case http.StatusBadRequest, http.StatusUnprocessableEntity:
|
||||
return infraerrors.New(http.StatusBadRequest, "OPENAI_REFERRAL_REJECTED", "the invitation request was rejected; check the email and eligibility")
|
||||
case http.StatusUnauthorized:
|
||||
return infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_AUTH_ERROR", "upstream authentication failed; refresh the account credentials")
|
||||
case http.StatusForbidden:
|
||||
return infraerrors.New(http.StatusForbidden, "OPENAI_REFERRAL_FORBIDDEN", "this account is not eligible for invitations")
|
||||
case http.StatusConflict:
|
||||
return infraerrors.New(http.StatusConflict, "OPENAI_REFERRAL_ALREADY_EXISTS", "an invitation already exists for this recipient")
|
||||
case http.StatusTooManyRequests:
|
||||
return infraerrors.New(http.StatusTooManyRequests, "OPENAI_REFERRAL_RATE_LIMITED", "invitation limit reached; try again later")
|
||||
default:
|
||||
return infraerrors.New(http.StatusBadGateway, "OPENAI_REFERRAL_UPSTREAM_ERROR", "invitation service is unavailable")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,23 +3,44 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/imroc/req/v3"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func referralTestService(t *testing.T, plan string, handler http.HandlerFunc) (*OpenAIQuotaService, *stubQuotaAccountRepo) {
|
||||
type referralClientStub struct {
|
||||
eligibility *OpenAIReferralEligibility
|
||||
queryErr, sendErr error
|
||||
calls []OpenAIReferralCall
|
||||
emails []string
|
||||
}
|
||||
|
||||
func (s *referralClientStub) QueryEligibility(_ context.Context, call OpenAIReferralCall) (*OpenAIReferralEligibility, error) {
|
||||
s.calls = append(s.calls, call)
|
||||
return s.eligibility, s.queryErr
|
||||
}
|
||||
func (s *referralClientStub) SendInvite(_ context.Context, call OpenAIReferralCall, email string) error {
|
||||
s.calls = append(s.calls, call)
|
||||
s.emails = append(s.emails, email)
|
||||
return s.sendErr
|
||||
}
|
||||
|
||||
func referralTestService(t *testing.T, plan string, client OpenAIReferralClient) (*OpenAIQuotaService, *stubQuotaAccountRepo) {
|
||||
t.Helper()
|
||||
a := &Account{ID: 100, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive,
|
||||
Credentials: map[string]any{"chatgpt_account_id": "workspace-test", "plan_type": plan}}
|
||||
repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{100: a}}
|
||||
tokens := &stubQuotaTokenCache{tokens: map[string]string{OpenAITokenCacheKey(a): "test-token"}}
|
||||
srv := httptest.NewServer(handler)
|
||||
t.Cleanup(srv.Close)
|
||||
return NewOpenAIQuotaService(repo, nil, NewOpenAITokenProvider(repo, tokens, nil), newQuotaRedirectingFactory(srv)), repo
|
||||
factory := func(string) (*req.Client, error) {
|
||||
t.Fatal("referral service must use the business interface, not the HTTP factory")
|
||||
return nil, nil
|
||||
}
|
||||
return NewOpenAIQuotaService(repo, nil, NewOpenAITokenProvider(repo, tokens, nil), factory, client), repo
|
||||
}
|
||||
|
||||
func TestOpenAIReferralSend(t *testing.T) {
|
||||
@@ -28,29 +49,13 @@ func TestOpenAIReferralSend(t *testing.T) {
|
||||
{"self_serve_business_usage_based", openAIReferralWorkspace},
|
||||
} {
|
||||
t.Run(tc.plan, func(t *testing.T) {
|
||||
var gets, posts int
|
||||
svc, repo := referralTestService(t, tc.plan, func(w http.ResponseWriter, r *http.Request) {
|
||||
require.Equal(t, "Bearer test-token", r.Header.Get("Authorization"))
|
||||
require.Equal(t, "workspace-test", r.Header.Get("ChatGPT-Account-ID"))
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
switch r.URL.Path {
|
||||
case "/backend-api/referrals/invite/eligibility":
|
||||
gets++
|
||||
require.Equal(t, http.MethodGet, r.Method)
|
||||
require.Equal(t, tc.program, r.URL.Query().Get("program_id"))
|
||||
require.Equal(t, "persistent", r.URL.Query().Get("entrypoint"))
|
||||
_, _ = w.Write([]byte(`{"should_show":true,"remaining_send_capacity":8,"remaining_reward_capacity":3,"grants":[{"grant_type":"rate_limit_reset_credit","amount":1,"recipient":"referrer"}],"rules":["Offer rule"]}`))
|
||||
case "/backend-api/referrals/invite":
|
||||
posts++
|
||||
require.Equal(t, http.MethodPost, r.Method)
|
||||
var body map[string]any
|
||||
require.NoError(t, json.NewDecoder(r.Body).Decode(&body))
|
||||
require.Equal(t, map[string]any{"program_id": tc.program, "entrypoint": "persistent", "emails": []any{"friend@example.com"}}, body)
|
||||
_, _ = w.Write([]byte(`{"invites":[{"referral_id":"test-invite","email":"friend@example.com"}]}`))
|
||||
default:
|
||||
t.Errorf("unexpected path: %s", r.URL.Path)
|
||||
}
|
||||
})
|
||||
send, reward := 8, 3
|
||||
client := &referralClientStub{eligibility: &OpenAIReferralEligibility{
|
||||
ShouldShow: true, RemainingSendCapacity: &send, RemainingRewardCapacity: &reward,
|
||||
Grants: []OpenAIReferralGrant{{GrantType: "rate_limit_reset_credit", Amount: 1, Recipient: "referrer"}},
|
||||
Rules: []string{"Offer rule"},
|
||||
}}
|
||||
svc, repo := referralTestService(t, tc.plan, client)
|
||||
eligibility, err := svc.QueryReferralEligibility(context.Background(), 100)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 3, *eligibility.AvailableInvites)
|
||||
@@ -64,8 +69,17 @@ func TestOpenAIReferralSend(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
require.True(t, result.Sent)
|
||||
require.Equal(t, "friend@example.com", result.Email)
|
||||
require.Equal(t, 2, gets, "send must recheck eligibility")
|
||||
require.Equal(t, 1, posts)
|
||||
require.Equal(t, []string{"friend@example.com"}, client.emails)
|
||||
require.Len(t, client.calls, 3, "send must recheck eligibility")
|
||||
for _, call := range client.calls {
|
||||
require.Equal(t, tc.program, call.ProgramID)
|
||||
headers := make(http.Header)
|
||||
for key, value := range call.Headers {
|
||||
headers.Set(key, value)
|
||||
}
|
||||
require.Equal(t, "Bearer test-token", headers.Get("Authorization"))
|
||||
require.Equal(t, "workspace-test", headers.Get("ChatGPT-Account-ID"))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -85,48 +99,39 @@ func TestOpenAIReferralSendGuards(t *testing.T) {
|
||||
{"unknown capacity", "a@example.com", `{"should_show":true}`, "OPENAI_REFERRAL_UNAVAILABLE", true, false},
|
||||
{"reward exhausted", "a@example.com", `{"should_show":true,"remaining_send_capacity":3,"offer_id":"credits_250","remaining_reward_capacity":0}`, "OPENAI_REFERRAL_UNAVAILABLE", true, false},
|
||||
{"unknown reward capacity", "a@example.com", `{"should_show":true,"remaining_send_capacity":3,"grants":[{}]}`, "OPENAI_REFERRAL_UNAVAILABLE", true, false},
|
||||
{"null response", "a@example.com", `null`, "OPENAI_REFERRAL_INVALID_RESPONSE", true, false},
|
||||
{"shadow", "a@example.com", `{}`, "OPENAI_REFERRAL_SHADOW_ACCOUNT", true, true},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
posts := 0
|
||||
svc, repo := referralTestService(t, "plus", func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == http.MethodPost {
|
||||
posts++
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(tc.body))
|
||||
})
|
||||
client := &referralClientStub{}
|
||||
require.NoError(t, json.Unmarshal([]byte(tc.body), &client.eligibility))
|
||||
svc, repo := referralTestService(t, "plus", client)
|
||||
if tc.shadow {
|
||||
parentID := int64(200)
|
||||
repo.accounts[100].ParentAccountID = &parentID
|
||||
}
|
||||
_, err := svc.SendReferralInvite(context.Background(), 100, OpenAIReferralSendRequest{Email: tc.email, ProgramID: openAIReferralConsumer, Confirmed: tc.confirmed})
|
||||
require.Equal(t, tc.reason, infraerrors.Reason(err))
|
||||
require.Zero(t, posts)
|
||||
require.Empty(t, client.emails)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIReferralSendDoesNotRetryOrExposeUpstreamErrors(t *testing.T) {
|
||||
for _, status := range []int{400, 401, 403, 409, 429, 500} {
|
||||
t.Run(http.StatusText(status), func(t *testing.T) {
|
||||
posts := 0
|
||||
svc, _ := referralTestService(t, "plus", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if r.Method == http.MethodGet {
|
||||
_, _ = w.Write([]byte(`{"should_show":true,"remaining_send_capacity":2,"requires_explicit_confirmation":false,"offer_id":"none"}`))
|
||||
return
|
||||
}
|
||||
posts++
|
||||
w.WriteHeader(status)
|
||||
_, _ = w.Write([]byte(`{"detail":"secret-token someone@example.com"}`))
|
||||
})
|
||||
_, err := svc.SendReferralInvite(context.Background(), 100, OpenAIReferralSendRequest{Email: "friend@example.com", ProgramID: openAIReferralConsumer})
|
||||
require.Error(t, err)
|
||||
require.NotContains(t, err.Error(), "secret-token")
|
||||
require.NotContains(t, err.Error(), "someone@example.com")
|
||||
require.Equal(t, 1, posts)
|
||||
})
|
||||
func TestOpenAIReferralSendStopsOnQueryError(t *testing.T) {
|
||||
client := &referralClientStub{queryErr: errors.New("eligibility unavailable")}
|
||||
svc, _ := referralTestService(t, "plus", client)
|
||||
_, err := svc.SendReferralInvite(context.Background(), 100, OpenAIReferralSendRequest{Email: "friend@example.com", ProgramID: openAIReferralConsumer, Confirmed: true})
|
||||
require.ErrorIs(t, err, client.queryErr)
|
||||
require.Empty(t, client.emails)
|
||||
}
|
||||
|
||||
func TestOpenAIReferralSendPreservesUnknownOutcomeWithoutRetry(t *testing.T) {
|
||||
count := 2
|
||||
client := &referralClientStub{
|
||||
eligibility: &OpenAIReferralEligibility{ShouldShow: true, RemainingSendCapacity: &count},
|
||||
sendErr: infraerrors.New(502, "OPENAI_REFERRAL_SEND_UNKNOWN", "outcome unknown"),
|
||||
}
|
||||
svc, _ := referralTestService(t, "plus", client)
|
||||
_, err := svc.SendReferralInvite(context.Background(), 100, OpenAIReferralSendRequest{Email: "friend@example.com", ProgramID: openAIReferralConsumer, Confirmed: true})
|
||||
require.ErrorIs(t, err, client.sendErr)
|
||||
require.Len(t, client.emails, 1)
|
||||
}
|
||||
|
||||
@@ -174,6 +174,13 @@ func ProvideOpenAITokenProvider(
|
||||
return p
|
||||
}
|
||||
|
||||
// ProvidePluginManager preserves account-directory wiring when regenerating Wire.
|
||||
func ProvidePluginManager(repo PluginRepository, encryptor SecretEncryptor, cfg *config.Config, hostInfo PluginHostInfo, kvStore PluginKVStore, gateway *OpenAIGatewayService) *PluginManager {
|
||||
manager := NewPluginManager(repo, encryptor, cfg, hostInfo, kvStore)
|
||||
manager.SetAccountDirectory(gateway)
|
||||
return manager
|
||||
}
|
||||
|
||||
// ProvideOpenAIQuotaService wires the OpenAI quota query/reset service.
|
||||
// It depends on the OpenAI token provider for refreshed access tokens and the
|
||||
// privacy client factory for the impersonated upstream HTTP client.
|
||||
@@ -182,9 +189,10 @@ func ProvideOpenAIQuotaService(
|
||||
proxyRepo ProxyRepository,
|
||||
tokenProvider *OpenAITokenProvider,
|
||||
privacyClientFactory PrivacyClientFactory,
|
||||
referralClient OpenAIReferralClient,
|
||||
openAIGatewayService *OpenAIGatewayService,
|
||||
) *OpenAIQuotaService {
|
||||
service := NewOpenAIQuotaService(accountRepo, proxyRepo, tokenProvider, privacyClientFactory)
|
||||
service := NewOpenAIQuotaService(accountRepo, proxyRepo, tokenProvider, privacyClientFactory, referralClient)
|
||||
service.agentIdentityWS = openAIGatewayService
|
||||
return service
|
||||
}
|
||||
@@ -924,7 +932,7 @@ var ProviderSet = wire.NewSet(
|
||||
NewTotpService,
|
||||
NewErrorPassthroughService,
|
||||
NewTLSFingerprintProfileService,
|
||||
NewPluginManager,
|
||||
ProvidePluginManager,
|
||||
NewDigestSessionStore,
|
||||
ProvideIdempotencyCoordinator,
|
||||
ProvideSystemOperationLockService,
|
||||
|
||||
Reference in New Issue
Block a user