Brisbanehuangandshaw
20ad5ec506
Group pricing (rate multiplier, peak windows, per-user overrides) and account cost (accounts.rate_multiplier) already live side by side, but nothing stops the scheduler from handing a request to an account whose cost multiplier exceeds what the group's pricing can profitably serve. Add an opt-in per-group profit gate that filters scheduling candidates by a margin rule, while ordering, scoring, stickiness and breakers keep working unchanged among qualified accounts. Admission rule: an account qualifies iff U <= D * (1 - min_margin - safety_buffer) within a small relative epsilon, where U is accounts.rate_multiplier (0 is legal; missing/negative/NaN/Inf are conservatively rejected as invalid) and D is the requester's effective downstream multiplier (user-group override ?? group default, times the group peak factor) frozen at the request's pricing instant. - groups gain profit_control_enabled / profit_min_margin / profit_safety_buffer (migration 191); the durable auth-cache invalidation trigger additionally watches the profit and pricing columns (migration 192) so out-of-band group edits cannot leave stale auth snapshots; GetByKeyForAuth explicitly projects the new columns and the API-key auth snapshot version is bumped to force a refresh of pre-existing snapshots - request-level pricing instant: token entry points install pricingAt into ctx; the profit threshold D and the RecordUsage peak factor read the same instant, so one request never changes price mid-flight across waits/retries/failover (media and unwired paths keep the existing record-time semantics) - the gate covers token requests on openai, anthropic, gemini, grok and antigravity groups: OpenAI-family handlers via WithOpenAIRequestPricingContext (responses incl. WS bridge, chat completions, messages, embeddings, alpha search), the shared gateway via WithGatewayTokenRequestPricing (messages, chat completions, responses, gemini model actions); composite groups cannot enable it directly; image/video/models/usage/count_tokens stay ungated and an explicit image-generation intent suppresses the gate end to end - post-slot recheck: after a slot is acquired the account is re-read via SchedulerSnapshotService.GetAccount (scheduler cache, then DB; only when both fail the check fails open with WARN + metric); a vetoed account releases its slot and joins the request's exclusion set for reselection; sticky bindings are written only after the final check passes, and an over-threshold sticky account is skipped, not deleted, so it comes back once its rate recovers - sticky-session cache contract: GatewayCache.GetSessionAccountID now returns ErrStickySessionNotFound on a miss (mapped from redis.Nil in the repository implementation, mirroring ErrRefreshTokenNotFound) so the profit sticky path can distinguish "no binding yet" from a real read failure without importing the cache driver in service code - cross-group re-entry (composite parent -> member group) resolves the gate against the member group and clears a stale parent gate instead of letting a foreign threshold veto accounts - per-platform/group activity counters (installs, threshold vetoes, invalid-rate vetoes, refresh failures) for observability - admin UI: profit-control section on the five platforms' group forms with percent input, validation and platform-switch reset; group create/update/duplicate normalize and validate the config at a single choke point - cmd/profit-preview: offline what-if tool that replays the production admission semantics over an exported config/account/override/model dump, reports per-model admitted-account counts under the default and the worst-case (lowest user override) D, and surfaces probe-sync staleness as warnings without affecting admission Tests: service unit coverage for gate resolution/veto/threshold epsilon/pricing instant/suppress marker/scheduler filtering and post-slot recheck (incl. -race on the profit surface), unit-tagged handler slot-recheck and capability-mapping regressions, sqlmock and real-PostgreSQL integration regressions for the GetByKeyForAuth projection and the migration-192 trigger watch list, API contract update, and frontend specs for the five-platform form helpers.
94 lines
3.5 KiB
Go
94 lines
3.5 KiB
Go
package service
|
||
|
||
// 投影漏列回归（service 半程）：认证快照 build → L2 JSON 序列化
|
||
// → 反序列化 → 还原 apiKey.Group → 请求 ctx → 利润门解析，全链路保真。
|
||
// repository 半程（真实 GetByKeyForAuth 投影）见
|
||
// internal/repository/api_key_repo_profit_projection_integration_test.go。
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"testing"
|
||
|
||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||
"github.com/stretchr/testify/require"
|
||
)
|
||
|
||
func profitAuthTestAPIKey() *APIKey {
|
||
groupID := int64(50)
|
||
return &APIKey{
|
||
ID: 82,
|
||
UserID: 40,
|
||
GroupID: &groupID,
|
||
Name: "profit-auth-roundtrip",
|
||
Status: StatusActive,
|
||
User: &User{
|
||
ID: 40,
|
||
Email: "profit@test.local",
|
||
Status: StatusActive,
|
||
Concurrency: 5,
|
||
},
|
||
Group: &Group{
|
||
ID: groupID,
|
||
Name: "VIP-roundtrip",
|
||
Platform: PlatformOpenAI,
|
||
Status: StatusActive,
|
||
Hydrated: true,
|
||
RateMultiplier: 0.06,
|
||
SubscriptionType: SubscriptionTypeStandard,
|
||
PeakRateEnabled: false,
|
||
ProfitControlEnabled: true,
|
||
ProfitMinMargin: 0.2,
|
||
ProfitSafetyBuffer: 0.05,
|
||
},
|
||
}
|
||
}
|
||
|
||
// 快照构建 → L2 JSON 往返 → 还原 → 装门：利润字段必须全程保真，阈值与
|
||
// 计费同源（0.06 × (1−0.25) = 0.045）。
|
||
func TestAPIKeyAuthSnapshotProfitControlRoundtrip(t *testing.T) {
|
||
svc := &APIKeyService{}
|
||
apiKey := profitAuthTestAPIKey()
|
||
|
||
snapshot := svc.snapshotFromAPIKey(context.Background(), apiKey)
|
||
require.NotNil(t, snapshot)
|
||
require.Equal(t, apiKeyAuthSnapshotVersion, snapshot.Version)
|
||
require.Equal(t, 18, snapshot.Version, "v18 起认证快照携带利润控制字段")
|
||
|
||
// 模拟 L2 缓存的完整 JSON 往返（与 apiKeyCache.SetAuthCache/GetAuthCache 同构）。
|
||
payload, err := json.Marshal(&APIKeyAuthCacheEntry{Snapshot: snapshot})
|
||
require.NoError(t, err)
|
||
var restored APIKeyAuthCacheEntry
|
||
require.NoError(t, json.Unmarshal(payload, &restored))
|
||
|
||
materialized, used, err := svc.applyAuthCacheEntry(apiKey.Key, &restored)
|
||
require.NoError(t, err)
|
||
require.True(t, used)
|
||
require.NotNil(t, materialized.Group)
|
||
require.True(t, materialized.Group.Hydrated)
|
||
require.True(t, materialized.Group.ProfitControlEnabled)
|
||
require.InDelta(t, 0.2, materialized.Group.ProfitMinMargin, 1e-12)
|
||
require.InDelta(t, 0.05, materialized.Group.ProfitSafetyBuffer, 1e-12)
|
||
require.InDelta(t, 0.06, materialized.Group.RateMultiplier, 1e-12)
|
||
|
||
// 中间件语义：materialized.Group 进请求 ctx → 门必须按快照配置装上。
|
||
ctx := context.WithValue(context.Background(), ctxkey.Group, materialized.Group)
|
||
gwSvc := &OpenAIGatewayService{}
|
||
gate := gwSvc.resolveOpenAIProfitControlGate(ctx, materialized.GroupID)
|
||
require.NotNil(t, gate, "还原后的认证分组必须能装门（投影漏列时本断言最先失败）")
|
||
require.InDelta(t, 0.06*(1-0.25), gate.threshold, 1e-12)
|
||
}
|
||
|
||
// 旧版本快照（v16 及更早，无利润字段保真保证）必须被淘汰回源，不得复用。
|
||
func TestAPIKeyAuthSnapshotOldVersionEvicted(t *testing.T) {
|
||
svc := &APIKeyService{}
|
||
snapshot := svc.snapshotFromAPIKey(context.Background(), profitAuthTestAPIKey())
|
||
require.NotNil(t, snapshot)
|
||
snapshot.Version = 16
|
||
|
||
materialized, used, err := svc.applyAuthCacheEntry("sk-old", &APIKeyAuthCacheEntry{Snapshot: snapshot})
|
||
require.NoError(t, err)
|
||
require.False(t, used, "版本不匹配的缓存条目必须淘汰并回源重建")
|
||
require.Nil(t, materialized)
|
||
}
|