mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
fix(grok): 二轮评审残留 — Search 叠加计费与 fail-closed 安全
- SearchCost 叠加 token(openai/gateway),未定价 warn - Token URL 校验失败回落 DefaultTokenURL,禁止 Effective 旁路 - free 判定收窄 paidSignal(仅 plan/月额度),usage% 不否决 free - web_search: mandatory 计费、uuid request_id、上游 failover 重选账号 - 调度阈值:7d/30d 不跨期 until + 48h stale 可选跳过 - SanitizeStoredCredentials 接入 create/update/bulk/SSO/ApplyOAuth - ApplyOAuth 成功清除 grok_needs_reauth;SSO 允许 header_override_enabled - VideoModelPrices 视为媒体定价完整;realtime 正常关闭仍计费
This commit is contained in:
@@ -1415,6 +1415,9 @@ func (h *AccountHandler) ApplyOAuthCredentials(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// Drop SSO/password residue; re-auth must leave only OAuth tokens on disk.
|
||||
req.Credentials = service.SanitizeStoredCredentials(existing.Platform, req.Credentials)
|
||||
|
||||
updatedAccount, err := h.adminService.UpdateAccount(ctx, accountID, &service.UpdateAccountInput{
|
||||
Type: req.Type,
|
||||
Credentials: req.Credentials,
|
||||
@@ -1442,6 +1445,19 @@ func (h *AccountHandler) ApplyOAuthCredentials(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
// Successful re-auth clears the soft spending-limit reauth flag for Grok.
|
||||
if existing.Platform == service.PlatformGrok {
|
||||
if clearErr := h.adminService.UpdateAccountExtra(ctx, accountID, map[string]any{
|
||||
"grok_needs_reauth": false,
|
||||
"grok_needs_reauth_reason": "",
|
||||
}); clearErr != nil {
|
||||
slog.Warn("apply_oauth_credentials.clear_grok_reauth_failed",
|
||||
"account_id", accountID,
|
||||
"err", clearErr,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
if cleared, clearErr := h.adminService.ClearAccountError(ctx, accountID); clearErr != nil {
|
||||
slog.Warn("apply_oauth_credentials.clear_error_failed",
|
||||
"account_id", accountID,
|
||||
|
||||
@@ -462,8 +462,9 @@ func grokSSOImportCredentials(built map[string]any, reqCredentials map[string]an
|
||||
// Only merge operator config from the request — never free-form secrets
|
||||
// (password / sso_token / cookie / etc.) into stored credentials.
|
||||
allowedReqKeys := map[string]struct{}{
|
||||
"base_url": {}, "model_mapping": {}, "header_override": {},
|
||||
"header_overrides": {}, "custom_headers": {},
|
||||
"base_url": {}, "model_mapping": {},
|
||||
"header_override": {}, "header_overrides": {}, "header_override_enabled": {},
|
||||
"custom_headers": {},
|
||||
}
|
||||
ops := map[string]any{}
|
||||
for k, v := range reqCredentials {
|
||||
@@ -489,7 +490,7 @@ func grokSSOImportCredentials(built map[string]any, reqCredentials map[string]an
|
||||
if reqBaseURL, ok := reqCredentials["base_url"].(string); ok && strings.TrimSpace(reqBaseURL) != "" {
|
||||
credentials["base_url"] = strings.TrimSpace(reqBaseURL)
|
||||
}
|
||||
return credentials
|
||||
return service.SanitizeStoredCredentials(service.PlatformGrok, credentials)
|
||||
}
|
||||
|
||||
func grokSSOImportExpiry(requestExpiresAt *int64, requestAutoPause *bool, tokenInfo *service.GrokTokenInfo) (*int64, *bool) {
|
||||
|
||||
@@ -2416,6 +2416,33 @@ func (h *GatewayHandler) submitUsageRecordTask(parent context.Context, task serv
|
||||
task(ctx)
|
||||
}
|
||||
|
||||
// submitMandatoryUsageRecordTask never silently drops billing work on pool overflow.
|
||||
func (h *GatewayHandler) submitMandatoryUsageRecordTask(parent context.Context, task service.UsageRecordTask) {
|
||||
if task == nil {
|
||||
return
|
||||
}
|
||||
task = wrapUsageRecordTaskContext(parent, task)
|
||||
if h.usageRecordWorkerPool != nil {
|
||||
if mode := h.usageRecordWorkerPool.Submit(task); !mode.Dropped() {
|
||||
return
|
||||
}
|
||||
logger.L().With(
|
||||
zap.String("component", "handler.gateway.usage"),
|
||||
).Warn("gateway.usage_record_task_mandatory_sync_fallback")
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
logger.L().With(
|
||||
zap.String("component", "handler.gateway.usage"),
|
||||
zap.Any("panic", recovered),
|
||||
).Error("gateway.usage_record_task_panic_recovered")
|
||||
}
|
||||
}()
|
||||
task(ctx)
|
||||
}
|
||||
|
||||
// getUserMsgQueueMode 获取当前请求的 UMQ 模式
|
||||
// 返回 "serialize" | "throttle" | ""
|
||||
func (h *GatewayHandler) getUserMsgQueueMode(account *service.Account, parsed *service.ParsedRequest) string {
|
||||
|
||||
@@ -3,19 +3,21 @@ package handler
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/websearch"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"github.com/tidwall/gjson"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
@@ -79,7 +81,8 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
selected, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), groupID, "", xai.DefaultTextModel, nil, "", 0)
|
||||
failedAccounts := make(map[int64]struct{})
|
||||
selected, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), groupID, "", xai.DefaultTextModel, failedAccounts, "", 0)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": gin.H{
|
||||
"type": "scheduling_error",
|
||||
@@ -143,7 +146,11 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
accountReleaseFunc = release
|
||||
}
|
||||
if accountReleaseFunc != nil {
|
||||
defer accountReleaseFunc()
|
||||
defer func() {
|
||||
if accountReleaseFunc != nil {
|
||||
accountReleaseFunc()
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// Scheduling is 100% the same as other requests:
|
||||
@@ -156,11 +163,40 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
|
||||
nativeResp, providerName, err := h.doGrokNativeWebSearch(c.Request.Context(), c, account, req.Query, req.MaxResults)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": gin.H{
|
||||
"type": "web_search_error",
|
||||
"message": err.Error(),
|
||||
}})
|
||||
return
|
||||
var failoverErr *service.UpstreamFailoverError
|
||||
if errors.As(err, &failoverErr) && failoverErr.ShouldRetryNextAccount() {
|
||||
failedAccounts[account.ID] = struct{}{}
|
||||
if accountReleaseFunc != nil {
|
||||
accountReleaseFunc()
|
||||
accountReleaseFunc = nil
|
||||
}
|
||||
for attempt := 1; attempt < 4; attempt++ {
|
||||
selected, selectErr := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), groupID, "", xai.DefaultTextModel, failedAccounts, "", 0)
|
||||
if selectErr != nil || selected == nil || selected.Account == nil {
|
||||
break
|
||||
}
|
||||
account = selected.Account
|
||||
if selected.Acquired {
|
||||
accountReleaseFunc = selected.ReleaseFunc
|
||||
} else {
|
||||
break
|
||||
}
|
||||
nativeResp, providerName, err = h.doGrokNativeWebSearch(c.Request.Context(), c, account, req.Query, req.MaxResults)
|
||||
if err == nil {
|
||||
break
|
||||
}
|
||||
if !errors.As(err, &failoverErr) || !failoverErr.ShouldRetryNextAccount() {
|
||||
break
|
||||
}
|
||||
failedAccounts[account.ID] = struct{}{}
|
||||
accountReleaseFunc()
|
||||
accountReleaseFunc = nil
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadGateway, gin.H{"error": gin.H{"type": "web_search_error", "message": err.Error()}})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
userAgent := c.GetHeader("User-Agent")
|
||||
@@ -169,15 +205,16 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
|
||||
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
|
||||
requestPayloadHash := service.HashUsageRequestPayload([]byte(req.Query))
|
||||
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
|
||||
// Unique per invocation (not hash(query) alone) so identical queries still bill.
|
||||
searchRequestID := "web_search:" + service.HashUsageRequestPayload([]byte(req.Query+"|"+clientIP+"|"+userAgent))
|
||||
// Request IDs are billing idempotency keys, so they must be unique per invocation.
|
||||
// Query/IP/UA hashes would collapse repeated identical searches into one charge.
|
||||
searchRequestID := "web_search:" + uuid.NewString()
|
||||
if apiKey.Group != nil && (apiKey.Group.GetSearchPricePer1k() == nil || *apiKey.Group.GetSearchPricePer1k() <= 0) {
|
||||
logger.L().With(
|
||||
zap.String("component", "handler.gateway.web_search"),
|
||||
zap.Int64("group_id", apiKey.Group.ID),
|
||||
).Warn("gateway.web_search.search_price_per_1k_unset_free")
|
||||
}
|
||||
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
|
||||
h.submitMandatoryUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
|
||||
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
|
||||
Result: &service.ForwardResult{
|
||||
RequestID: searchRequestID,
|
||||
|
||||
@@ -92,11 +92,13 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) {
|
||||
elapsed := time.Since(started)
|
||||
if proxyErr != nil {
|
||||
reqLog.Info("grok_realtime.proxy_failed", zap.Error(proxyErr))
|
||||
_ = conn.Close(coderws.StatusInternalError, "upstream realtime websocket failed")
|
||||
// Do not bill failed realtime sessions (avoids charging idle/error wall-clock).
|
||||
return
|
||||
if !isExpectedGrokRealtimeClose(proxyErr) {
|
||||
_ = conn.Close(coderws.StatusInternalError, "upstream realtime websocket failed")
|
||||
return
|
||||
}
|
||||
}
|
||||
// Bill successful realtime sessions by elapsed minutes when group price is set.
|
||||
// A relay normally returns a close error when either side closes normally.
|
||||
// Those sessions still consumed upstream audio time and must be billed.
|
||||
if elapsed > 0 {
|
||||
result := &service.OpenAIForwardResult{
|
||||
Model: model,
|
||||
@@ -107,6 +109,19 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) {
|
||||
}
|
||||
}
|
||||
|
||||
func isExpectedGrokRealtimeClose(err error) bool {
|
||||
if err == nil {
|
||||
return true
|
||||
}
|
||||
switch coderws.CloseStatus(err) {
|
||||
case coderws.StatusNormalClosure, coderws.StatusGoingAway,
|
||||
coderws.StatusNoStatusRcvd, coderws.StatusAbnormalClosure:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// GrokVoice handles xAI Voice HTTP endpoints. endpoint is "tts", "stt", or "custom-voices".
|
||||
func (h *OpenAIGatewayHandler) GrokVoice(c *gin.Context, endpoint string) {
|
||||
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
coderws "github.com/coder/websocket"
|
||||
)
|
||||
|
||||
func TestIsExpectedGrokRealtimeClose(t *testing.T) {
|
||||
for _, status := range []coderws.StatusCode{
|
||||
coderws.StatusNormalClosure,
|
||||
coderws.StatusGoingAway,
|
||||
coderws.StatusNoStatusRcvd,
|
||||
coderws.StatusAbnormalClosure,
|
||||
} {
|
||||
if !isExpectedGrokRealtimeClose(coderws.CloseError{Code: status}) {
|
||||
t.Fatalf("status %v should be treated as an expected session close", status)
|
||||
}
|
||||
}
|
||||
if isExpectedGrokRealtimeClose(coderws.CloseError{Code: coderws.StatusPolicyViolation}) {
|
||||
t.Fatal("policy violations must not be treated as billable normal closes")
|
||||
}
|
||||
}
|
||||
@@ -35,11 +35,12 @@ const (
|
||||
)
|
||||
|
||||
func NewGrokOAuthClient() service.GrokOAuthClient {
|
||||
// Prefer validated official/allowlisted token URL so env misconfig cannot
|
||||
// exfiltrate auth codes / refresh tokens to an arbitrary host.
|
||||
// Fail closed: never fall back to an unvalidated EffectiveTokenURL (env can
|
||||
// point at an attacker host and steal code/refresh tokens).
|
||||
tokenURL, err := xai.ValidatedTokenURL()
|
||||
if err != nil || strings.TrimSpace(tokenURL) == "" {
|
||||
tokenURL = xai.EffectiveTokenURL()
|
||||
// Official allowlisted endpoint only — never EffectiveTokenURL() (raw env).
|
||||
tokenURL = xai.DefaultTokenURL
|
||||
}
|
||||
return &grokOAuthClient{tokenURL: tokenURL}
|
||||
}
|
||||
|
||||
@@ -47,6 +47,8 @@ func TestGrokOAuthClientExchangeAndRefreshUseFormFields(t *testing.T) {
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
// Tests inject a loopback token endpoint; allowlist requires unsafe override.
|
||||
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
|
||||
t.Setenv(xai.EnvTokenURL, server.URL)
|
||||
|
||||
client := NewGrokOAuthClient()
|
||||
@@ -80,6 +82,7 @@ func TestGrokOAuthClientRefreshForbiddenClassifiesOnlyExplicitEntitlement(t *tes
|
||||
_, _ = w.Write([]byte(tt.body))
|
||||
}))
|
||||
defer server.Close()
|
||||
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
|
||||
t.Setenv(xai.EnvTokenURL, server.URL)
|
||||
|
||||
client := NewGrokOAuthClient()
|
||||
@@ -96,6 +99,7 @@ func TestGrokOAuthClientStatusErrorRedactsSensitiveResponseBody(t *testing.T) {
|
||||
_, _ = w.Write([]byte(`{"error":"invalid_grant","access_token":"access-secret","refresh_token":"refresh-secret","code_verifier":"verifier-secret"}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
|
||||
t.Setenv(xai.EnvTokenURL, server.URL)
|
||||
|
||||
client := NewGrokOAuthClient()
|
||||
@@ -119,3 +123,12 @@ func TestGrokOAuthEntitlementDenialRequiresExplicitEvidence(t *testing.T) {
|
||||
require.False(t, grokOAuthHasExplicitEntitlementDenial(`{"error":"forbidden","message":"request forbidden"}`))
|
||||
require.False(t, grokOAuthHasExplicitEntitlementDenial(`<html>403 Forbidden</html>`))
|
||||
}
|
||||
|
||||
func TestNewGrokOAuthClient_UnvalidatedTokenURLFallsBackToDefault(t *testing.T) {
|
||||
// Without unsafe overrides, a random env TokenURL must not be used (fail-closed).
|
||||
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "")
|
||||
t.Setenv(xai.EnvTokenURL, "https://evil.example/oauth/token")
|
||||
|
||||
client := NewGrokOAuthClient().(*grokOAuthClient)
|
||||
require.Equal(t, xai.DefaultTokenURL, client.tokenURL)
|
||||
}
|
||||
|
||||
@@ -297,36 +297,61 @@ func grokThresholdCandidates(account *Account) []*accountSchedulingThresholdCand
|
||||
// Official billing windows (weekly credits / monthly used%) from probe snapshot.
|
||||
// Complements header-based grok_sched_* so paid SuperGrok accounts can auto-pause
|
||||
// when request/token headers are sparse but billing % is high.
|
||||
// Do not cross-fallback period ends (weekly until ≠ monthly end) to avoid multi-week parks.
|
||||
const grokBillingThresholdStaleAfter = 48 * time.Hour
|
||||
if billing, err := grokBillingSnapshotFromExtra(account.Extra); err == nil && billing != nil {
|
||||
if billing.UsagePercent != nil {
|
||||
candidates = append(candidates, &accountSchedulingThresholdCandidate{
|
||||
window: "seven_day",
|
||||
scope: "grok_billing",
|
||||
usedPercent: *billing.UsagePercent,
|
||||
until: parseSchedulingResetAt(firstNonEmpty(billing.PeriodEnd, billing.BillingPeriodEnd)),
|
||||
})
|
||||
}
|
||||
monthlyUtil := 0.0
|
||||
hasMonthly := false
|
||||
if billing.UsedPercent != nil {
|
||||
monthlyUtil = *billing.UsedPercent
|
||||
hasMonthly = true
|
||||
} else if billing.MonthlyLimitCents != nil && *billing.MonthlyLimitCents > 0 && billing.UsedCents != nil {
|
||||
monthlyUtil = (*billing.UsedCents / *billing.MonthlyLimitCents) * 100
|
||||
hasMonthly = true
|
||||
}
|
||||
if hasMonthly {
|
||||
candidates = append(candidates, &accountSchedulingThresholdCandidate{
|
||||
window: "thirty_day",
|
||||
scope: "grok_billing",
|
||||
usedPercent: monthlyUtil,
|
||||
until: parseSchedulingResetAt(firstNonEmpty(billing.BillingPeriodEnd, billing.PeriodEnd)),
|
||||
})
|
||||
fetchedAt := firstNonEmpty(billing.UpdatedAt, billing.FetchedAt, billing.WeeklyUpdatedAt, billing.MonthlyUpdatedAt)
|
||||
if !grokBillingSnapshotStale(fetchedAt, grokBillingThresholdStaleAfter) {
|
||||
if billing.UsagePercent != nil {
|
||||
if until := parseSchedulingResetAt(billing.PeriodEnd); until != nil {
|
||||
candidates = append(candidates, &accountSchedulingThresholdCandidate{
|
||||
window: "seven_day",
|
||||
scope: "grok_billing",
|
||||
usedPercent: *billing.UsagePercent,
|
||||
until: until,
|
||||
})
|
||||
}
|
||||
}
|
||||
monthlyUtil := 0.0
|
||||
hasMonthly := false
|
||||
if billing.UsedPercent != nil {
|
||||
monthlyUtil = *billing.UsedPercent
|
||||
hasMonthly = true
|
||||
} else if billing.MonthlyLimitCents != nil && *billing.MonthlyLimitCents > 0 && billing.UsedCents != nil {
|
||||
monthlyUtil = (*billing.UsedCents / *billing.MonthlyLimitCents) * 100
|
||||
hasMonthly = true
|
||||
}
|
||||
if hasMonthly {
|
||||
if until := parseSchedulingResetAt(billing.BillingPeriodEnd); until != nil {
|
||||
candidates = append(candidates, &accountSchedulingThresholdCandidate{
|
||||
window: "thirty_day",
|
||||
scope: "grok_billing",
|
||||
usedPercent: monthlyUtil,
|
||||
until: until,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return candidates
|
||||
}
|
||||
|
||||
func grokBillingSnapshotStale(fetchedAtRFC3339 string, maxAge time.Duration) bool {
|
||||
if maxAge <= 0 {
|
||||
return false
|
||||
}
|
||||
raw := strings.TrimSpace(fetchedAtRFC3339)
|
||||
if raw == "" {
|
||||
// Unknown freshness: still allow (probe may omit timestamps on older snapshots).
|
||||
return false
|
||||
}
|
||||
ts, err := time.Parse(time.RFC3339, raw)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
return time.Since(ts) > maxAge
|
||||
}
|
||||
|
||||
func pickLatestResetSchedulingCandidate(candidates []*accountSchedulingThresholdCandidate, threshold int, now time.Time) *accountSchedulingThresholdCandidate {
|
||||
var winner *accountSchedulingThresholdCandidate
|
||||
for _, candidate := range candidates {
|
||||
|
||||
@@ -232,7 +232,7 @@ func (s *AccountService) Create(ctx context.Context, req CreateAccountRequest) (
|
||||
Notes: normalizeAccountNotes(req.Notes),
|
||||
Platform: req.Platform,
|
||||
Type: req.Type,
|
||||
Credentials: req.Credentials,
|
||||
Credentials: SanitizeStoredCredentials(req.Platform, req.Credentials),
|
||||
Extra: req.Extra,
|
||||
ProxyID: req.ProxyID,
|
||||
Concurrency: req.Concurrency,
|
||||
@@ -325,7 +325,7 @@ func (s *AccountService) Update(ctx context.Context, id int64, req UpdateAccount
|
||||
}
|
||||
|
||||
if req.Credentials != nil {
|
||||
account.Credentials = *req.Credentials
|
||||
account.Credentials = SanitizeStoredCredentials(account.Platform, *req.Credentials)
|
||||
}
|
||||
|
||||
if req.Extra != nil {
|
||||
|
||||
@@ -495,6 +495,8 @@ func (s *adminServiceImpl) CreateAccount(ctx context.Context, input *CreateAccou
|
||||
if err := NormalizeHeaderOverrideCredentials(input.Credentials); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Never persist ephemeral SSO/password secrets after OAuth conversion.
|
||||
input.Credentials = SanitizeStoredCredentials(input.Platform, input.Credentials)
|
||||
|
||||
account, err := buildAccountForCreate(input, accountExtra)
|
||||
if err != nil {
|
||||
@@ -605,6 +607,8 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
|
||||
if err := NormalizeHeaderOverrideCredentials(account.Credentials); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Strip SSO/password residue that must never sit next to OAuth tokens.
|
||||
account.Credentials = SanitizeStoredCredentials(account.Platform, account.Credentials)
|
||||
}
|
||||
// Extra 使用 map:需要区分“未提供(nil)”与“显式清空({})”。
|
||||
// 关闭配额限制时前端会删除 quota_* 键并提交 extra:{},此时也必须落库。
|
||||
@@ -1015,6 +1019,11 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp
|
||||
if err := NormalizeHeaderOverrideCredentials(input.Credentials); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Bulk may mix platforms; always drop ephemeral SSO/password keys (cookie
|
||||
// only when platform is known Grok — empty platform still strips password/*).
|
||||
if input.Credentials != nil {
|
||||
input.Credentials = SanitizeStoredCredentials("", input.Credentials)
|
||||
}
|
||||
|
||||
// Prepare bulk updates for columns and JSONB fields.
|
||||
repoUpdates := AccountBulkUpdate{
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
package service
|
||||
|
||||
import "strings"
|
||||
|
||||
// SanitizeStoredCredentials strips secrets that must never be persisted on the
|
||||
// account credentials map after conversion to OAuth tokens (Grok Web SSO / password).
|
||||
// Call from admin create/update/import/apply-oauth paths.
|
||||
func SanitizeStoredCredentials(platform string, creds map[string]any) map[string]any {
|
||||
if creds == nil {
|
||||
return nil
|
||||
}
|
||||
// Always strip Grok ephemeral secrets regardless of platform label (import
|
||||
// may set platform later; better never store them).
|
||||
for _, key := range []string{
|
||||
"password", "sso_token", "sso", "sso-rw", "clearTextPassword",
|
||||
} {
|
||||
delete(creds, key)
|
||||
}
|
||||
// Platform-specific: also strip cookie jar residue on Grok OAuth stores.
|
||||
if strings.EqualFold(strings.TrimSpace(platform), PlatformGrok) {
|
||||
delete(creds, "cookie")
|
||||
}
|
||||
return creds
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSanitizeStoredCredentials_StripsEphemeralSSOSecrets(t *testing.T) {
|
||||
creds := map[string]any{
|
||||
"access_token": "at",
|
||||
"refresh_token": "rt",
|
||||
"password": "secret",
|
||||
"sso_token": "sso",
|
||||
"sso": "cookie-sso",
|
||||
"sso-rw": "rw",
|
||||
"clearTextPassword": "plain",
|
||||
"cookie": "jar",
|
||||
"base_url": "https://api.x.ai",
|
||||
}
|
||||
out := SanitizeStoredCredentials(PlatformGrok, creds)
|
||||
require.Equal(t, "at", out["access_token"])
|
||||
require.Equal(t, "rt", out["refresh_token"])
|
||||
require.Equal(t, "https://api.x.ai", out["base_url"])
|
||||
require.NotContains(t, out, "password")
|
||||
require.NotContains(t, out, "sso_token")
|
||||
require.NotContains(t, out, "sso")
|
||||
require.NotContains(t, out, "sso-rw")
|
||||
require.NotContains(t, out, "clearTextPassword")
|
||||
require.NotContains(t, out, "cookie")
|
||||
}
|
||||
|
||||
func TestSanitizeStoredCredentials_NonGrokKeepsCookieButDropsPassword(t *testing.T) {
|
||||
creds := map[string]any{
|
||||
"cookie": "session",
|
||||
"password": "x",
|
||||
"api_key": "k",
|
||||
}
|
||||
out := SanitizeStoredCredentials(PlatformOpenAI, creds)
|
||||
require.Equal(t, "session", out["cookie"])
|
||||
require.Equal(t, "k", out["api_key"])
|
||||
require.NotContains(t, out, "password")
|
||||
}
|
||||
|
||||
func TestSanitizeStoredCredentials_NilSafe(t *testing.T) {
|
||||
require.Nil(t, SanitizeStoredCredentials(PlatformGrok, nil))
|
||||
}
|
||||
@@ -1307,7 +1307,7 @@ func (s *GatewayService) DoGrokNativeResponsesJSON(ctx context.Context, account
|
||||
}
|
||||
resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("grok native search upstream: %w", err)
|
||||
return nil, &UpstreamFailoverError{StatusCode: http.StatusBadGateway, Reason: GatewayFailureReason("grok_search_transport")}
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
respBytes, readErr := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
|
||||
@@ -1319,6 +1319,9 @@ func (s *GatewayService) DoGrokNativeResponsesJSON(ctx context.Context, account
|
||||
if len(msg) > 200 {
|
||||
msg = msg[:200]
|
||||
}
|
||||
if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusPaymentRequired || resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 {
|
||||
return nil, &UpstreamFailoverError{StatusCode: resp.StatusCode, ResponseBody: respBytes}
|
||||
}
|
||||
return nil, fmt.Errorf("grok upstream %d: %s", resp.StatusCode, msg)
|
||||
}
|
||||
return respBytes, nil
|
||||
|
||||
@@ -813,19 +813,29 @@ func (s *GatewayService) calculateRecordUsageCost(
|
||||
return s.calculateImageCost(ctx, result, apiKey, billingModel, imageMultiplier)
|
||||
}
|
||||
|
||||
// Grok native web_search / tool search: group search_price_per_1k when set.
|
||||
if result.SearchCount > 0 {
|
||||
price := groupSearchPricePer1kFromAPIKey(apiKey)
|
||||
return s.billingService.CalculateSearchCost(result.SearchCount, price, multiplier)
|
||||
}
|
||||
// Voice audio (TTS / STT / realtime) when present on the forward result.
|
||||
if result.AudioUsage != nil {
|
||||
cfg := groupAudioPriceConfigFromAPIKey(apiKey)
|
||||
return s.billingService.CalculateAudioCost(result.AudioUsage.Mode, result.AudioUsage.DurationOrUnits, cfg, multiplier)
|
||||
}
|
||||
|
||||
// Token 计费
|
||||
return s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, opts)
|
||||
// Token 计费;SearchCount 为叠加 surcharge(不替代 token)。
|
||||
tokenCost := s.calculateTokenCost(ctx, result, apiKey, billingModel, multiplier, opts)
|
||||
if result.SearchCount > 0 {
|
||||
price := groupSearchPricePer1kFromAPIKey(apiKey)
|
||||
if price == nil || *price <= 0 {
|
||||
logger.LegacyPrintf("service.gateway", "[Billing] search_price_per_1k unset; search calls free group_model=%s count=%d", billingModel, result.SearchCount)
|
||||
}
|
||||
searchCost := s.billingService.CalculateSearchCost(result.SearchCount, price, multiplier)
|
||||
if searchCost != nil && (searchCost.TotalCost > 0 || searchCost.ActualCost > 0) {
|
||||
if tokenCost == nil {
|
||||
return searchCost
|
||||
}
|
||||
tokenCost.TotalCost += searchCost.TotalCost
|
||||
tokenCost.ActualCost += searchCost.ActualCost
|
||||
}
|
||||
}
|
||||
return tokenCost
|
||||
}
|
||||
|
||||
// compositeBillableModel 决定 composite 分组请求的计费模型:来源覆盖把计费模型
|
||||
|
||||
@@ -319,8 +319,9 @@ func isKnownGrokFreeAccount(account *Account) bool {
|
||||
paidSignal = true
|
||||
}
|
||||
}
|
||||
if billing.UsagePercent != nil || billing.UsedPercent != nil ||
|
||||
(billing.MonthlyLimitCents != nil && *billing.MonthlyLimitCents > 0) {
|
||||
// Paid evidence: explicit plan name or monthly dollar cap — NOT mere presence of
|
||||
// weekly/monthly usage percent (Free accounts also report creditUsagePercent).
|
||||
if billing.MonthlyLimitCents != nil && *billing.MonthlyLimitCents > 0 {
|
||||
paidSignal = true
|
||||
}
|
||||
// xAI deliberately reports an empty plan for Free accounts; only paid
|
||||
@@ -362,10 +363,23 @@ func isKnownGrokFreeAccount(account *Account) bool {
|
||||
}
|
||||
}
|
||||
}
|
||||
// Explicit free credentials (subscription_tier / plan_type free|basic) win over
|
||||
// weak paid inference so weekly credit % cannot disable soft-gate / media block
|
||||
// for real Free OAuth accounts. Strong paid evidence (named paid plan, monthly
|
||||
// limit cents) still wins over *inferred* free only.
|
||||
if freeSignal {
|
||||
// Explicit free/basic tier from credentials or free plan string.
|
||||
// Only override if we saw a clearly paid plan name (not mere usage %).
|
||||
if paidSignal {
|
||||
// paidSignal from MonthlyLimitCents or non-free Plan — treat as paid.
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
// Explicit paid evidence always wins over an inferred Free signal. This
|
||||
// protects upgraded/stale accounts whose previous quota snapshot still
|
||||
// carries the historical 2M Free token limit.
|
||||
return !paidSignal && (freeSignal || inferredFreeSignal)
|
||||
return !paidSignal && inferredFreeSignal
|
||||
}
|
||||
|
||||
func isGrokFreeSubscriptionTier(tier string) bool {
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestCalculateOpenAIRecordUsageCost_SearchIsAdditiveToTokens(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
price := 10.0 // $10 / 1k searches → 100 searches = $1.0
|
||||
svc := &OpenAIGatewayService{
|
||||
billingService: newTestBillingService(),
|
||||
}
|
||||
apiKey := &APIKey{
|
||||
Group: &Group{
|
||||
SearchPricePer1k: &price,
|
||||
},
|
||||
}
|
||||
|
||||
// claude-sonnet-4 fallback: Input $3/MTok, Output $15/MTok
|
||||
// 1000 in + 500 out → 0.003 + 0.0075 = 0.0105
|
||||
// + 100 searches → +1.0 → total 1.0105
|
||||
cost, err := svc.calculateOpenAIRecordUsageCost(
|
||||
context.Background(),
|
||||
&OpenAIForwardResult{SearchCount: 100},
|
||||
apiKey,
|
||||
[]string{"claude-sonnet-4"},
|
||||
1.0,
|
||||
1.0,
|
||||
1.0,
|
||||
1.0,
|
||||
UsageTokens{InputTokens: 1000, OutputTokens: 500},
|
||||
"",
|
||||
false,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, cost)
|
||||
require.InDelta(t, 1.0105, cost.ActualCost, 1e-9)
|
||||
require.InDelta(t, 1.0105, cost.TotalCost, 1e-9)
|
||||
}
|
||||
|
||||
func TestCalculateOpenAIRecordUsageCost_SearchOnlyWhenNoTokenPricing(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
price := 10.0
|
||||
svc := &OpenAIGatewayService{
|
||||
billingService: newTestBillingService(),
|
||||
}
|
||||
apiKey := &APIKey{
|
||||
Group: &Group{SearchPricePer1k: &price},
|
||||
}
|
||||
// Empty model list: token path fails; search-only surcharge still bills.
|
||||
cost, err := svc.calculateOpenAIRecordUsageCost(
|
||||
context.Background(),
|
||||
&OpenAIForwardResult{SearchCount: 100},
|
||||
apiKey,
|
||||
nil,
|
||||
1.0,
|
||||
1.0,
|
||||
1.0,
|
||||
1.0,
|
||||
UsageTokens{},
|
||||
"",
|
||||
false,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, cost)
|
||||
require.InDelta(t, 1.0, cost.ActualCost, 1e-9)
|
||||
}
|
||||
|
||||
func TestGroupMediaPricingLooksIncomplete_VideoModelPricesComplete(t *testing.T) {
|
||||
t.Parallel()
|
||||
require.True(t, groupMediaPricingLooksIncomplete(nil))
|
||||
require.True(t, groupMediaPricingLooksIncomplete(&Group{}))
|
||||
require.False(t, groupMediaPricingLooksIncomplete(&Group{
|
||||
VideoModelPrices: map[string]map[string]float64{
|
||||
"grok-imagine-video": {"720p": 0.1},
|
||||
},
|
||||
}))
|
||||
}
|
||||
@@ -447,11 +447,6 @@ func (s *OpenAIGatewayService) calculateOpenAIRecordUsageCost(
|
||||
return s.calculateOpenAIVideoCost(ctx, billingModel, apiKey, result, videoMultiplier), nil
|
||||
}
|
||||
}
|
||||
if result != nil && result.SearchCount > 0 {
|
||||
// Grok native web_search / tool search: group search_price_per_1k when set.
|
||||
price := groupSearchPricePer1kFromAPIKey(apiKey)
|
||||
return s.billingService.CalculateSearchCost(result.SearchCount, price, webSearchMultiplier), nil
|
||||
}
|
||||
if result != nil && result.AudioUsage != nil {
|
||||
cfg := groupAudioPriceConfigFromAPIKey(apiKey)
|
||||
return s.billingService.CalculateAudioCost(result.AudioUsage.Mode, result.AudioUsage.DurationOrUnits, cfg, webSearchMultiplier), nil
|
||||
@@ -463,33 +458,64 @@ func (s *OpenAIGatewayService) calculateOpenAIRecordUsageCost(
|
||||
return s.calculateOpenAIImageCost(ctx, billingModel, apiKey, result, imageMultiplier), nil
|
||||
}
|
||||
}
|
||||
if len(billingModels) == 0 || billingModel == "" {
|
||||
return nil, errors.New("openai usage billing model is empty")
|
||||
}
|
||||
|
||||
// Token path (optional search surcharge is additive — never replaces token cost).
|
||||
var tokenCost *CostBreakdown
|
||||
var lastErr error
|
||||
for _, candidate := range billingModels {
|
||||
candidate = strings.TrimSpace(candidate)
|
||||
if candidate == "" {
|
||||
continue
|
||||
if len(billingModels) > 0 && billingModel != "" {
|
||||
for _, candidate := range billingModels {
|
||||
candidate = strings.TrimSpace(candidate)
|
||||
if candidate == "" {
|
||||
continue
|
||||
}
|
||||
cost, err := s.calculateOpenAIRecordUsageTokenCost(
|
||||
ctx,
|
||||
apiKey,
|
||||
candidate,
|
||||
multiplier,
|
||||
tokens,
|
||||
serviceTier,
|
||||
longContextBillingEnabled,
|
||||
)
|
||||
if err == nil {
|
||||
tokenCost = cost
|
||||
break
|
||||
}
|
||||
lastErr = err
|
||||
}
|
||||
cost, err := s.calculateOpenAIRecordUsageTokenCost(
|
||||
ctx,
|
||||
apiKey,
|
||||
candidate,
|
||||
multiplier,
|
||||
tokens,
|
||||
serviceTier,
|
||||
longContextBillingEnabled,
|
||||
)
|
||||
if err == nil {
|
||||
return cost, nil
|
||||
}
|
||||
// Search-only (e.g. no model / zero tokens): still bill search when priced.
|
||||
searchCost := (*CostBreakdown)(nil)
|
||||
if result != nil && result.SearchCount > 0 {
|
||||
price := groupSearchPricePer1kFromAPIKey(apiKey)
|
||||
if price == nil || *price <= 0 {
|
||||
logger.L().Warn("openai_usage.grok_search_price_per_1k_unset_free",
|
||||
zap.Int("search_count", result.SearchCount),
|
||||
zap.String("model", billingModel),
|
||||
)
|
||||
}
|
||||
lastErr = err
|
||||
searchCost = s.billingService.CalculateSearchCost(result.SearchCount, price, webSearchMultiplier)
|
||||
}
|
||||
if lastErr == nil {
|
||||
lastErr = errors.New("no non-empty billing model candidates")
|
||||
|
||||
if tokenCost == nil && searchCost == nil {
|
||||
if lastErr == nil {
|
||||
if len(billingModels) == 0 || billingModel == "" {
|
||||
return nil, errors.New("openai usage billing model is empty")
|
||||
}
|
||||
lastErr = errors.New("no non-empty billing model candidates")
|
||||
}
|
||||
return nil, fmt.Errorf("calculate OpenAI usage cost failed for billing models %s: %w", strings.Join(billingModels, ","), lastErr)
|
||||
}
|
||||
return nil, fmt.Errorf("calculate OpenAI usage cost failed for billing models %s: %w", strings.Join(billingModels, ","), lastErr)
|
||||
if tokenCost == nil {
|
||||
return searchCost, nil
|
||||
}
|
||||
if searchCost == nil || (searchCost.TotalCost == 0 && searchCost.ActualCost == 0) {
|
||||
return tokenCost, nil
|
||||
}
|
||||
// Additive: tokens + search surcharge.
|
||||
tokenCost.TotalCost += searchCost.TotalCost
|
||||
tokenCost.ActualCost += searchCost.ActualCost
|
||||
return tokenCost, nil
|
||||
}
|
||||
|
||||
func isGrokVideoBillingModel(model string) bool {
|
||||
@@ -500,6 +526,8 @@ func isGrokVideoUsageResult(result *OpenAIForwardResult, billingModels []string)
|
||||
if result == nil || result.VideoCount <= 0 {
|
||||
return false
|
||||
}
|
||||
// VideoCount alone is authoritative for async video completion billing.
|
||||
// Prefer model-family match when present; never drop video mode on rename/mapping.
|
||||
candidates := append([]string{}, billingModels...)
|
||||
candidates = append(candidates, result.BillingModel, result.Model, result.UpstreamModel)
|
||||
for _, candidate := range candidates {
|
||||
@@ -507,7 +535,7 @@ func isGrokVideoUsageResult(result *OpenAIForwardResult, billingModels []string)
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
return true
|
||||
}
|
||||
|
||||
func isUsagePricingUnavailableError(err error) bool {
|
||||
@@ -675,6 +703,11 @@ func groupMediaPricingLooksIncomplete(group *Group) bool {
|
||||
if group.ImageRateMultiplier != 0 || group.VideoRateMultiplier != 0 {
|
||||
return false
|
||||
}
|
||||
// Per-model video prices are first-class billing config; a projection that
|
||||
// already carries them is complete enough to skip a DB refresh.
|
||||
if len(group.VideoModelPrices) > 0 {
|
||||
return false
|
||||
}
|
||||
return group.ImagePrice1K == nil && group.ImagePrice2K == nil && group.ImagePrice4K == nil &&
|
||||
group.VideoPrice480P == nil && group.VideoPrice720P == nil && group.VideoPrice1080P == nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user