mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-11-02 16:32:14 +00:00
feat(openai): show Codex credits and manage referral invitations
This commit is contained in:
1 parent
bbdcfbac0a
commit
5090ffe05e
21 files changed
+1421
-21
No files matched your search
@@ -21,12 +21,14 @@ type OpenAIOAuthHandler struct {
|
||||
openaiOAuthService *service.OpenAIOAuthService
|
||||
adminService service.AdminService
|
||||
quotaService openAIQuotaService
|
||||
referralService openAIReferralService
|
||||
rateLimitService openAIAccountStateRecoverer
|
||||
}
|
||||
|
||||
type openAIQuotaService interface {
|
||||
QueryUsage(ctx context.Context, accountID int64) (*service.OpenAIQuotaUsage, error)
|
||||
CacheResetCreditsSnapshot(ctx context.Context, accountID int64, credits *service.OpenAIRateLimitResetCredits) error
|
||||
CacheCreditsSnapshot(ctx context.Context, accountID int64, usage *service.OpenAIQuotaUsage) error
|
||||
CachePostResetSnapshot(ctx context.Context, accountID int64, usage *service.OpenAIQuotaUsage) error
|
||||
ResetCredit(ctx context.Context, accountID int64) (*service.OpenAIQuotaResetResult, error)
|
||||
}
|
||||
@@ -57,7 +59,8 @@ type openAIQuotaResetResponse struct {
|
||||
// failed display-cache write must never discard a successful upstream read.
|
||||
type openAIQuotaRefreshResponse struct {
|
||||
service.OpenAIQuotaUsage
|
||||
CachePersisted bool `json:"cache_persisted"`
|
||||
CachePersisted bool `json:"cache_persisted"`
|
||||
CreditsCachePersisted bool `json:"credits_cache_persisted"`
|
||||
}
|
||||
|
||||
// openAIQuotaResetPostProcessContext detaches the post-reset bookkeeping from the
|
||||
@@ -92,6 +95,7 @@ func NewOpenAIOAuthHandler(
|
||||
// `== nil` capability guards below and panic instead of returning 400.
|
||||
if quotaService != nil {
|
||||
h.quotaService = quotaService
|
||||
h.referralService = quotaService
|
||||
}
|
||||
if rateLimitService != nil {
|
||||
h.rateLimitService = rateLimitService
|
||||
@@ -522,6 +526,11 @@ func (h *OpenAIOAuthHandler) RefreshQuota(c *gin.Context) {
|
||||
service.NotifyOpenAIAutoResetCredit(accountID)
|
||||
|
||||
refreshResponse := openAIQuotaRefreshResponse{OpenAIQuotaUsage: *usage}
|
||||
if err := h.quotaService.CacheCreditsSnapshot(c.Request.Context(), accountID, usage); err != nil {
|
||||
slog.Warn("openai_quota_credits_cache_persist_failed", "account_id", accountID, "error", err)
|
||||
} else {
|
||||
refreshResponse.CreditsCachePersisted = true
|
||||
}
|
||||
// A failed snapshot write leaves the previous cache intact — report it as a
|
||||
// partial success instead of discarding the usage payload we just fetched,
|
||||
// which would leave the card without a credit count at all.
|
||||
|
||||
@@ -17,11 +17,14 @@ import (
|
||||
)
|
||||
|
||||
type openAIQuotaWorkflowStub struct {
|
||||
resetResult *service.OpenAIQuotaResetResult
|
||||
resetErr error
|
||||
queryResult *service.OpenAIQuotaUsage
|
||||
queryErr error
|
||||
cacheErr error
|
||||
resetResult *service.OpenAIQuotaResetResult
|
||||
resetErr error
|
||||
queryResult *service.OpenAIQuotaUsage
|
||||
queryErr error
|
||||
cacheErr error
|
||||
creditsCacheErr error
|
||||
creditsCacheCalls int
|
||||
cachedCreditsUsage *service.OpenAIQuotaUsage
|
||||
|
||||
resetCalls int
|
||||
queryCalls int
|
||||
@@ -48,6 +51,12 @@ func (s *openAIQuotaWorkflowStub) CacheResetCreditsSnapshot(ctx context.Context,
|
||||
return s.cacheErr
|
||||
}
|
||||
|
||||
func (s *openAIQuotaWorkflowStub) CacheCreditsSnapshot(_ context.Context, _ int64, usage *service.OpenAIQuotaUsage) error {
|
||||
s.creditsCacheCalls++
|
||||
s.cachedCreditsUsage = usage
|
||||
return s.creditsCacheErr
|
||||
}
|
||||
|
||||
func (s *openAIQuotaWorkflowStub) CachePostResetSnapshot(ctx context.Context, _ int64, _ *service.OpenAIQuotaUsage) error {
|
||||
s.cacheCalls++
|
||||
s.cacheCtxErr = ctx.Err()
|
||||
@@ -410,6 +419,34 @@ func TestOpenAIRefreshQuota_PersistFailureStillReturnsUsage(t *testing.T) {
|
||||
require.Equal(t, 1, quota.cacheCalls)
|
||||
}
|
||||
|
||||
func TestOpenAIRefreshQuota_CreditsPersistIndependently(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
resetErr error
|
||||
creditsErr error
|
||||
}{
|
||||
{name: "both saved"},
|
||||
{name: "reset details missing", resetErr: errors.New("missing expirations")},
|
||||
{name: "points cache failed", creditsErr: errors.New("write failed")},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
quota := successfulOpenAIQuotaWorkflowStub()
|
||||
balance := "1250.75"
|
||||
quota.queryResult.Credits = &service.OpenAICredits{HasCredits: true, Balance: &balance}
|
||||
quota.cacheErr = tc.resetErr
|
||||
quota.creditsCacheErr = tc.creditsErr
|
||||
status, envelope := performOpenAIQuotaRefreshRequest(t, &OpenAIOAuthHandler{quotaService: quota})
|
||||
require.Equal(t, http.StatusOK, status)
|
||||
require.Equal(t, tc.resetErr == nil, envelope.Data.CachePersisted)
|
||||
require.Equal(t, tc.creditsErr == nil, envelope.Data.CreditsCachePersisted)
|
||||
require.Equal(t, quota.queryResult.Credits, envelope.Data.Credits)
|
||||
require.Equal(t, quota.queryResult, quota.cachedCreditsUsage)
|
||||
require.Equal(t, 1, quota.creditsCacheCalls)
|
||||
require.Zero(t, quota.resetCalls)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// An empty-but-successful upstream read must not be dereferenced blindly.
|
||||
func TestOpenAIQuotaEmptyUsageIsHandledWithoutPanic(t *testing.T) {
|
||||
t.Run("refresh reports an internal error", func(t *testing.T) {
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type openAIReferralService interface {
|
||||
QueryReferralEligibility(context.Context, int64) (*service.OpenAIReferralEligibility, error)
|
||||
CacheReferralSnapshot(context.Context, int64, *service.OpenAIReferralEligibility) error
|
||||
SendReferralInvite(context.Context, int64, service.OpenAIReferralSendRequest) (*service.OpenAIReferralSendResult, error)
|
||||
}
|
||||
|
||||
type openAIReferralRefreshResponse struct {
|
||||
Eligibility *service.OpenAIReferralEligibility `json:"eligibility"`
|
||||
CachePersisted bool `json:"cache_persisted"`
|
||||
}
|
||||
|
||||
type openAIReferralSendResponse struct {
|
||||
service.OpenAIReferralSendResult
|
||||
openAIReferralRefreshResponse
|
||||
RefreshFailed bool `json:"refresh_failed"`
|
||||
}
|
||||
|
||||
func (h *OpenAIOAuthHandler) referralAccountID(c *gin.Context) (int64, bool) {
|
||||
id, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil || id <= 0 {
|
||||
response.BadRequest(c, "Invalid account ID")
|
||||
return 0, false
|
||||
}
|
||||
if h.referralService == nil {
|
||||
response.BadRequest(c, "OpenAI referral service is not enabled")
|
||||
return 0, false
|
||||
}
|
||||
return id, true
|
||||
}
|
||||
|
||||
// RefreshReferrals persists a display snapshot, hence POST and admin audit.
|
||||
func (h *OpenAIOAuthHandler) RefreshReferrals(c *gin.Context) {
|
||||
id, ok := h.referralAccountID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
eligibility, err := h.referralService.QueryReferralEligibility(c.Request.Context(), id)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
if eligibility == nil {
|
||||
response.Error(c, http.StatusBadGateway, "Empty invitation eligibility response")
|
||||
return
|
||||
}
|
||||
cacheErr := h.referralService.CacheReferralSnapshot(c.Request.Context(), id, eligibility)
|
||||
response.Success(c, openAIReferralRefreshResponse{Eligibility: eligibility, CachePersisted: cacheErr == nil})
|
||||
}
|
||||
|
||||
func (h *OpenAIOAuthHandler) SendReferralInvite(c *gin.Context) {
|
||||
id, ok := h.referralAccountID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var input service.OpenAIReferralSendRequest
|
||||
if err := c.ShouldBindJSON(&input); err != nil {
|
||||
response.BadRequest(c, "Invalid invitation request")
|
||||
return
|
||||
}
|
||||
result, err := h.referralService.SendReferralInvite(c.Request.Context(), id, input)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
if result == nil || !result.Sent {
|
||||
response.Error(c, http.StatusBadGateway, "Invitation outcome is unknown; check Codex before sending again")
|
||||
return
|
||||
}
|
||||
// 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)
|
||||
if refreshErr != nil {
|
||||
eligibility = nil
|
||||
}
|
||||
cacheErr := h.referralService.CacheReferralSnapshot(ctx, id, eligibility)
|
||||
response.Success(c, openAIReferralSendResponse{
|
||||
OpenAIReferralSendResult: *result,
|
||||
openAIReferralRefreshResponse: openAIReferralRefreshResponse{Eligibility: eligibility, CachePersisted: cacheErr == nil},
|
||||
RefreshFailed: eligibility == nil,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
//go:build unit
|
||||
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type referralHandlerStub struct {
|
||||
eligibility *service.OpenAIReferralEligibility
|
||||
queryErr, cacheErr error
|
||||
input service.OpenAIReferralSendRequest
|
||||
cached *service.OpenAIReferralEligibility
|
||||
sends int
|
||||
}
|
||||
|
||||
func (s *referralHandlerStub) QueryReferralEligibility(context.Context, int64) (*service.OpenAIReferralEligibility, error) {
|
||||
return s.eligibility, s.queryErr
|
||||
}
|
||||
func (s *referralHandlerStub) CacheReferralSnapshot(_ context.Context, _ int64, e *service.OpenAIReferralEligibility) error {
|
||||
s.cached = e
|
||||
return s.cacheErr
|
||||
}
|
||||
func (s *referralHandlerStub) SendReferralInvite(_ context.Context, _ int64, input service.OpenAIReferralSendRequest) (*service.OpenAIReferralSendResult, error) {
|
||||
s.sends++
|
||||
s.input = input
|
||||
return &service.OpenAIReferralSendResult{Email: input.Email, Sent: true}, nil
|
||||
}
|
||||
|
||||
func referralHandlerRequest(t *testing.T, stub *referralHandlerStub, path, body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
h := &OpenAIOAuthHandler{referralService: stub}
|
||||
router := gin.New()
|
||||
router.POST("/:id/referrals/refresh", h.RefreshReferrals)
|
||||
router.POST("/:id/referrals/invite", h.SendReferralInvite)
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, httptest.NewRequest(http.MethodPost, path, strings.NewReader(body)))
|
||||
return rec
|
||||
}
|
||||
|
||||
func TestOpenAIReferralHandlerRefresh(t *testing.T) {
|
||||
count := 2
|
||||
stub := &referralHandlerStub{eligibility: &service.OpenAIReferralEligibility{ShouldShow: true, AvailableInvites: &count}, cacheErr: errors.New("cache failed")}
|
||||
rec := referralHandlerRequest(t, stub, "/100/referrals/refresh", "")
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var body struct {
|
||||
Data openAIReferralRefreshResponse `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &body))
|
||||
require.Equal(t, 2, *body.Data.Eligibility.AvailableInvites)
|
||||
require.False(t, body.Data.CachePersisted)
|
||||
require.Zero(t, stub.sends)
|
||||
}
|
||||
|
||||
func TestOpenAIReferralHandlerSentSurvivesRefreshFailure(t *testing.T) {
|
||||
stub := &referralHandlerStub{queryErr: errors.New("upstream timed out")}
|
||||
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, body.Data.Sent)
|
||||
require.Equal(t, "friend@example.com", body.Data.Email)
|
||||
require.True(t, body.Data.RefreshFailed)
|
||||
require.True(t, body.Data.CachePersisted)
|
||||
require.Nil(t, stub.cached, "stale remaining count must be invalidated")
|
||||
require.True(t, stub.input.Confirmed)
|
||||
require.Equal(t, 1, stub.sends)
|
||||
}
|
||||
|
||||
func TestOpenAIReferralHandlerRejectsInvalidInput(t *testing.T) {
|
||||
stub := &referralHandlerStub{}
|
||||
for _, path := range []string{"/invalid/referrals/invite", "/0/referrals/invite", "/100/referrals/invite"} {
|
||||
rec := referralHandlerRequest(t, stub, path, "not-json")
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code)
|
||||
}
|
||||
require.Zero(t, stub.sends)
|
||||
}
|
||||
@@ -454,6 +454,8 @@ func registerOpenAIOAuthRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
|
||||
openai.GET("/accounts/:id/quota", h.Admin.OpenAIOAuth.QueryQuota)
|
||||
openai.POST("/accounts/:id/quota/refresh", h.Admin.OpenAIOAuth.RefreshQuota)
|
||||
openai.POST("/accounts/:id/reset-quota", h.Admin.OpenAIOAuth.ResetQuota)
|
||||
openai.POST("/accounts/:id/referrals/refresh", h.Admin.OpenAIOAuth.RefreshReferrals)
|
||||
openai.POST("/accounts/:id/referrals/invite", h.Admin.OpenAIOAuth.SendReferralInvite)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestQueryUsageCodexCredits(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
body string
|
||||
want string
|
||||
}{
|
||||
{"decimal balance", `{"credits":{"has_credits":true,"unlimited":false,"balance":"12345678901234567890.0123"}}`, `{"has_credits":true,"unlimited":false,"balance":"12345678901234567890.0123"}`},
|
||||
{"zero", `{"credits":{"has_credits":false,"unlimited":false,"balance":"0"}}`, `{"has_credits":false,"unlimited":false,"balance":"0"}`},
|
||||
{"unlimited", `{"credits":{"has_credits":false,"unlimited":true,"balance":null}}`, `{"has_credits":false,"unlimited":true,"balance":null}`},
|
||||
{"hidden balance", `{"credits":{"has_credits":true,"unlimited":false}}`, `{"has_credits":true,"unlimited":false,"balance":null}`},
|
||||
{"absent", `{}`, `null`},
|
||||
{"null", `{"credits":null}`, `null`},
|
||||
} {
|
||||
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"}}
|
||||
repo := &stubQuotaAccountRepo{accounts: map[int64]*Account{100: account}}
|
||||
tokens := &stubQuotaTokenCache{tokens: map[string]string{OpenAITokenCacheKey(account): "test-token"}}
|
||||
var paths []string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
paths = append(paths, r.URL.Path)
|
||||
require.Equal(t, http.MethodGet, r.Method)
|
||||
require.Equal(t, "Bearer test-token", r.Header.Get("Authorization"))
|
||||
require.Equal(t, "test-workspace", r.Header.Get("ChatGPT-Account-ID"))
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
if r.URL.Path == "/backend-api/wham/usage" {
|
||||
_, _ = w.Write([]byte(tc.body))
|
||||
return
|
||||
}
|
||||
// Reset-card details may be unavailable without losing points.
|
||||
http.Error(w, "unavailable", http.StatusServiceUnavailable)
|
||||
}))
|
||||
defer srv.Close()
|
||||
svc := NewOpenAIQuotaService(repo, nil, NewOpenAITokenProvider(repo, tokens, nil), newQuotaRedirectingFactory(srv))
|
||||
usage, err := svc.QueryUsage(context.Background(), 100)
|
||||
require.NoError(t, err)
|
||||
require.Positive(t, usage.FetchedAt)
|
||||
encoded, err := json.Marshal(usage.Credits)
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, tc.want, string(encoded))
|
||||
require.Equal(t, []string{"/backend-api/wham/usage", "/backend-api/wham/rate-limit-reset-credits"}, paths)
|
||||
require.Zero(t, repo.extraUpdateCalls, "read-only query must not write account state")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCacheCodexCreditsSnapshot(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
repo := &stubQuotaAccountRepo{}
|
||||
svc := &OpenAIQuotaService{accountRepo: repo}
|
||||
balance := "1200.50"
|
||||
usage := &OpenAIQuotaUsage{FetchedAt: 123, Credits: &OpenAICredits{HasCredits: true, Balance: &balance}}
|
||||
|
||||
// No reset-card details are required to persist the balance on the queried row.
|
||||
require.NoError(t, svc.CacheCreditsSnapshot(ctx, 200, usage))
|
||||
encoded, err := json.Marshal(repo.extraUpdates[200][openaiQuotaCreditsKey])
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, `{"credits":{"has_credits":true,"unlimited":false,"balance":"1200.50"},"fetched_at":123}`, string(encoded))
|
||||
require.NotContains(t, repo.extraUpdates, int64(100), "a shadow row must not overwrite its parent")
|
||||
require.NotContains(t, repo.extraUpdates[200], openaiQuotaResetCreditsKey)
|
||||
|
||||
// Missing credits from a successful read invalidate the previous balance.
|
||||
require.NoError(t, svc.CacheCreditsSnapshot(ctx, 200, &OpenAIQuotaUsage{FetchedAt: 456}))
|
||||
encoded, err = json.Marshal(repo.extraUpdates[200][openaiQuotaCreditsKey])
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, `{"credits":null,"fetched_at":456}`, string(encoded))
|
||||
|
||||
require.Error(t, svc.CacheCreditsSnapshot(ctx, 200, nil))
|
||||
require.Equal(t, 2, repo.extraUpdateCalls)
|
||||
repo.extraUpdateErr = errors.New("database unavailable")
|
||||
require.ErrorContains(t, svc.CacheCreditsSnapshot(ctx, 200, usage), "database unavailable")
|
||||
}
|
||||
@@ -35,6 +35,7 @@ const (
|
||||
openaiQuotaSecFetchMode = "no-cors"
|
||||
openaiQuotaSecFetchDest = "empty"
|
||||
openaiQuotaResetCreditsKey = "codex_reset_credit_snapshot"
|
||||
openaiQuotaCreditsKey = "codex_credits_snapshot"
|
||||
)
|
||||
|
||||
// OpenAIRateLimitWindow describes a single rate-limit window returned by
|
||||
@@ -75,6 +76,21 @@ type OpenAIRateLimitResetCredits struct {
|
||||
Credits []OpenAIRateLimitResetCreditDetail `json:"credits,omitempty"`
|
||||
}
|
||||
|
||||
// OpenAICredits is the spendable Codex credit balance from /wham/usage.
|
||||
// It is separate from reset credits. Upstream represents the balance as a
|
||||
// nullable decimal string; keep that representation to preserve precision.
|
||||
// Source: Codex 41ece455b7fa, codex-backend-openapi-models/src/models/credit_status_details.rs.
|
||||
type OpenAICredits struct {
|
||||
HasCredits bool `json:"has_credits"`
|
||||
Unlimited bool `json:"unlimited"`
|
||||
Balance *string `json:"balance"`
|
||||
}
|
||||
|
||||
type openAICreditsSnapshot struct {
|
||||
Credits *OpenAICredits `json:"credits"`
|
||||
FetchedAt int64 `json:"fetched_at"`
|
||||
}
|
||||
|
||||
// OpenAIQuotaUsage is the typed projection of /wham/usage we expose to the UI.
|
||||
// Fields not relevant to the quota card are intentionally omitted to keep the
|
||||
// surface narrow; full upstream payload preservation is unnecessary.
|
||||
@@ -86,6 +102,7 @@ type OpenAIQuotaUsage struct {
|
||||
RateLimit *OpenAIRateLimit `json:"rate_limit,omitempty"`
|
||||
AdditionalRateLimits []OpenAIAdditionalRateLimit `json:"additional_rate_limits,omitempty"`
|
||||
RateLimitResetCredits *OpenAIRateLimitResetCredits `json:"rate_limit_reset_credits,omitempty"`
|
||||
Credits *OpenAICredits `json:"credits,omitempty"`
|
||||
FetchedAt int64 `json:"fetched_at"`
|
||||
autoResetCandidates []openAIAutoResetCreditCandidate
|
||||
}
|
||||
@@ -228,16 +245,36 @@ func (s *OpenAIQuotaService) CacheResetCreditsSnapshot(ctx context.Context, acco
|
||||
return s.cacheResetCreditsSnapshot(ctx, accountID, credits, nil)
|
||||
}
|
||||
|
||||
// CacheCreditsSnapshot stores the queried row's display snapshot independently
|
||||
// of reset-credit expiration details. A successful read with absent credits
|
||||
// replaces the previous balance with unknown, never with a fabricated zero.
|
||||
func (s *OpenAIQuotaService) CacheCreditsSnapshot(ctx context.Context, accountID int64, usage *OpenAIQuotaUsage) error {
|
||||
if usage == nil {
|
||||
return infraerrors.New(http.StatusBadGateway, "OPENAI_QUOTA_EMPTY_USAGE", "openai quota query returned an empty result")
|
||||
}
|
||||
if err := s.accountRepo.UpdateExtra(ctx, accountID, map[string]any{
|
||||
openaiQuotaCreditsKey: openAICreditsSnapshot{Credits: usage.Credits, FetchedAt: usage.FetchedAt},
|
||||
}); err != nil {
|
||||
return infraerrors.New(http.StatusInternalServerError, "OPENAI_QUOTA_CACHE_WRITE_FAILED", "failed to cache Codex credits").WithCause(err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CachePostResetSnapshot persists the credits and usage windows observed after a reset.
|
||||
func (s *OpenAIQuotaService) CachePostResetSnapshot(ctx context.Context, accountID int64, usage *OpenAIQuotaUsage) error {
|
||||
if usage == nil {
|
||||
return s.cacheResetCreditsSnapshot(ctx, accountID, nil, nil)
|
||||
}
|
||||
updates := buildOpenAIAutoResetUsageUpdates(usage, time.Now())
|
||||
if updates == nil {
|
||||
updates = make(map[string]any)
|
||||
}
|
||||
updates[openaiQuotaCreditsKey] = openAICreditsSnapshot{Credits: usage.Credits, FetchedAt: usage.FetchedAt}
|
||||
return s.cacheResetCreditsSnapshot(
|
||||
ctx,
|
||||
accountID,
|
||||
usage.RateLimitResetCredits,
|
||||
buildOpenAIAutoResetUsageUpdates(usage, time.Now()),
|
||||
updates,
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -700,7 +700,10 @@ func TestCachePostResetSnapshot(t *testing.T) {
|
||||
repo := &stubQuotaAccountRepo{}
|
||||
svc := &OpenAIQuotaService{accountRepo: repo}
|
||||
credits := &OpenAIRateLimitResetCredits{AvailableCount: 0}
|
||||
balance := "1200.50"
|
||||
usage := &OpenAIQuotaUsage{
|
||||
Credits: &OpenAICredits{HasCredits: true, Balance: &balance},
|
||||
FetchedAt: 123,
|
||||
RateLimitResetCredits: credits,
|
||||
RateLimit: &OpenAIRateLimit{
|
||||
PrimaryWindow: &OpenAIRateLimitWindow{
|
||||
@@ -715,6 +718,7 @@ func TestCachePostResetSnapshot(t *testing.T) {
|
||||
require.NoError(t, svc.CachePostResetSnapshot(context.Background(), 100, usage))
|
||||
require.Equal(t, 1, repo.extraUpdateCalls)
|
||||
require.Equal(t, credits, repo.extraUpdates[100][openaiQuotaResetCreditsKey])
|
||||
require.Equal(t, openAICreditsSnapshot{Credits: usage.Credits, FetchedAt: 123}, repo.extraUpdates[100][openaiQuotaCreditsKey])
|
||||
require.Equal(t, 0.0, repo.extraUpdates[100]["codex_5h_used_percent"])
|
||||
require.Equal(t, 0.0, repo.extraUpdates[100]["codex_7d_used_percent"])
|
||||
}
|
||||
|
||||
@@ -0,0 +1,227 @@
|
||||
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 {
|
||||
GrantType string `json:"grant_type"`
|
||||
Recipient string `json:"recipient"`
|
||||
Amount float64 `json:"amount"`
|
||||
}
|
||||
|
||||
type OpenAIReferralEligibility struct {
|
||||
ShouldShow bool `json:"should_show"`
|
||||
RemainingSendCapacity *int `json:"remaining_send_capacity"`
|
||||
RemainingRewardCapacity *int `json:"remaining_reward_capacity"`
|
||||
RequiresExplicitConfirmation *bool `json:"requires_explicit_confirmation"`
|
||||
OfferID string `json:"offer_id,omitempty"`
|
||||
Title string `json:"title,omitempty"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Rules []string `json:"rules,omitempty"`
|
||||
Grants []OpenAIReferralGrant `json:"grants,omitempty"`
|
||||
ProgramID string `json:"program_id"`
|
||||
AvailableInvites *int `json:"available_invites"`
|
||||
FetchedAt int64 `json:"fetched_at"`
|
||||
}
|
||||
|
||||
// referralCapacity follows the desktop's send/reward capacity constraints.
|
||||
// The desktop additionally caps one batch at five; this API sends one email.
|
||||
func referralCapacity(e *OpenAIReferralEligibility) *int {
|
||||
if !e.ShouldShow {
|
||||
zero := 0
|
||||
return &zero
|
||||
}
|
||||
if e.RemainingSendCapacity == nil {
|
||||
return nil
|
||||
}
|
||||
count := max(0, *e.RemainingSendCapacity)
|
||||
if len(e.Grants) > 0 || (e.OfferID != "" && e.OfferID != "none") {
|
||||
if e.RemainingRewardCapacity == nil {
|
||||
return nil
|
||||
}
|
||||
count = min(count, max(0, *e.RemainingRewardCapacity))
|
||||
}
|
||||
return &count
|
||||
}
|
||||
|
||||
func (s *OpenAIQuotaService) referralAccount(ctx context.Context, id int64, sending bool) (*Account, error) {
|
||||
if s == nil || s.accountRepo == nil {
|
||||
return nil, infraerrors.New(http.StatusServiceUnavailable, "OPENAI_REFERRAL_NOT_CONFIGURED", "referral service is unavailable")
|
||||
}
|
||||
a, err := s.accountRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if a == nil {
|
||||
return nil, ErrAccountNotFound
|
||||
}
|
||||
if a.Platform != PlatformOpenAI || a.Type != AccountTypeOAuth {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "OPENAI_REFERRAL_INVALID_ACCOUNT", "referrals require an OpenAI OAuth account")
|
||||
}
|
||||
if a.IsShadow() {
|
||||
if sending {
|
||||
return nil, infraerrors.New(http.StatusConflict, "OPENAI_REFERRAL_SHADOW_ACCOUNT", "send invitations from the parent account")
|
||||
}
|
||||
return resolveCredentialAccount(ctx, s.accountRepo, a)
|
||||
}
|
||||
return a, nil
|
||||
}
|
||||
|
||||
func referralProgram(a *Account) string {
|
||||
if strings.EqualFold(a.GetCredential("account_type"), "workspace") {
|
||||
return openAIReferralWorkspace
|
||||
}
|
||||
switch strings.ToLower(strings.TrimSpace(a.GetCredential("plan_type"))) {
|
||||
case "team", "business", "self_serve_business_prolite", "self_serve_business_usage_based", "free_workspace", "enterprise", "enterprise_cbp_usage_based", "enterprise_cbp_automation", "edu", "education":
|
||||
return openAIReferralWorkspace
|
||||
default:
|
||||
return openAIReferralConsumer
|
||||
}
|
||||
}
|
||||
|
||||
func (s *OpenAIQuotaService) referralRequest(ctx context.Context, id int64) (*req.Request, error) {
|
||||
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")
|
||||
}
|
||||
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")
|
||||
}
|
||||
// Sending an invitation has no documented idempotency key. Never retry it.
|
||||
return client.R().SetContext(ctx).SetHeaders(headers).SetRetryCount(0), nil
|
||||
}
|
||||
|
||||
func (s *OpenAIQuotaService) QueryReferralEligibility(ctx context.Context, id int64) (*OpenAIReferralEligibility, error) {
|
||||
account, err := s.referralAccount(ctx, id, false)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout)
|
||||
defer cancel()
|
||||
r, err := s.referralRequest(callCtx, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
program := referralProgram(account)
|
||||
resp, err := r.SetQueryParams(map[string]string{"program_id": program, "entrypoint": openAIReferralEntrypoint}).Get(openAIReferralURL + "/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 *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")
|
||||
}
|
||||
result.ProgramID = program
|
||||
result.AvailableInvites = referralCapacity(result)
|
||||
result.FetchedAt = time.Now().Unix()
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *OpenAIQuotaService) CacheReferralSnapshot(ctx context.Context, id int64, eligibility *OpenAIReferralEligibility) error {
|
||||
return s.accountRepo.UpdateExtra(ctx, id, map[string]any{openAIReferralSnapshotKey: eligibility})
|
||||
}
|
||||
|
||||
type OpenAIReferralSendRequest struct {
|
||||
Email string `json:"email"`
|
||||
ProgramID string `json:"program_id"`
|
||||
Confirmed bool `json:"confirmed"`
|
||||
}
|
||||
|
||||
type OpenAIReferralSendResult struct {
|
||||
Email string `json:"email"`
|
||||
Sent bool `json:"sent"`
|
||||
}
|
||||
|
||||
func (s *OpenAIQuotaService) SendReferralInvite(ctx context.Context, id int64, input OpenAIReferralSendRequest) (*OpenAIReferralSendResult, error) {
|
||||
email := strings.TrimSpace(input.Email)
|
||||
address, err := mail.ParseAddress(email)
|
||||
if err != nil || address.Address != email || len(email) > 254 || strings.ContainsAny(email, "\r\n") {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "OPENAI_REFERRAL_INVALID_EMAIL", "enter one valid email address")
|
||||
}
|
||||
account, err := s.referralAccount(ctx, id, true)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if input.ProgramID != referralProgram(account) {
|
||||
return nil, infraerrors.New(http.StatusConflict, "OPENAI_REFERRAL_PROGRAM_CHANGED", "refresh invitation eligibility before sending")
|
||||
}
|
||||
// Recheck the server's current eligibility rather than trusting a UI cache.
|
||||
eligibility, err := s.QueryReferralEligibility(ctx, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if !eligibility.ShouldShow || eligibility.AvailableInvites == nil || *eligibility.AvailableInvites <= 0 {
|
||||
return nil, infraerrors.New(http.StatusConflict, "OPENAI_REFERRAL_UNAVAILABLE", "no invitations are currently available")
|
||||
}
|
||||
if (eligibility.RequiresExplicitConfirmation == nil || *eligibility.RequiresExplicitConfirmation) && !input.Confirmed {
|
||||
return nil, infraerrors.New(http.StatusBadRequest, "OPENAI_REFERRAL_CONFIRMATION_REQUIRED", "confirm the recipient's consent before sending")
|
||||
}
|
||||
callCtx, cancel := context.WithTimeout(ctx, openaiQuotaUpstreamTimeout)
|
||||
defer cancel()
|
||||
r, err := s.referralRequest(callCtx, id)
|
||||
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")
|
||||
}
|
||||
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")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
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) {
|
||||
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
|
||||
}
|
||||
|
||||
func TestOpenAIReferralSend(t *testing.T) {
|
||||
for _, tc := range []struct{ plan, program string }{
|
||||
{"plus", openAIReferralConsumer}, {"team", openAIReferralWorkspace},
|
||||
{"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)
|
||||
}
|
||||
})
|
||||
eligibility, err := svc.QueryReferralEligibility(context.Background(), 100)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 3, *eligibility.AvailableInvites)
|
||||
require.Equal(t, []string{"Offer rule"}, eligibility.Rules)
|
||||
require.Positive(t, eligibility.FetchedAt)
|
||||
require.NoError(t, svc.CacheReferralSnapshot(context.Background(), 100, eligibility))
|
||||
require.Equal(t, eligibility, repo.extraUpdates[100][openAIReferralSnapshotKey])
|
||||
result, err := svc.SendReferralInvite(context.Background(), 100, OpenAIReferralSendRequest{
|
||||
Email: " friend@example.com ", ProgramID: tc.program, Confirmed: true,
|
||||
})
|
||||
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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIReferralSendGuards(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, email, body, reason string
|
||||
confirmed, shadow bool
|
||||
}{
|
||||
{"invalid email", "bad-email", `{}`, "OPENAI_REFERRAL_INVALID_EMAIL", true, false},
|
||||
{"multiple emails", "a@example.com,b@example.com", `{}`, "OPENAI_REFERRAL_INVALID_EMAIL", true, false},
|
||||
{"display name", "User <a@example.com>", `{}`, "OPENAI_REFERRAL_INVALID_EMAIL", true, false},
|
||||
{"header injection", "a@example.com\r\nBcc: b@example.com", `{}`, "OPENAI_REFERRAL_INVALID_EMAIL", true, false},
|
||||
{"no consent", "a@example.com", `{"should_show":true,"remaining_send_capacity":2}`, "OPENAI_REFERRAL_CONFIRMATION_REQUIRED", false, false},
|
||||
{"ineligible", "a@example.com", `{"should_show":false,"remaining_send_capacity":2}`, "OPENAI_REFERRAL_UNAVAILABLE", true, false},
|
||||
{"exhausted", "a@example.com", `{"should_show":true,"remaining_send_capacity":0}`, "OPENAI_REFERRAL_UNAVAILABLE", true, false},
|
||||
{"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))
|
||||
})
|
||||
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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user