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.
109 lines
3.7 KiB
Go
109 lines
3.7 KiB
Go
//go:build integration
|
|
|
|
package repository
|
|
|
|
import (
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
|
"github.com/stretchr/testify/require"
|
|
"github.com/stretchr/testify/suite"
|
|
)
|
|
|
|
type GatewayCacheSuite struct {
|
|
IntegrationRedisSuite
|
|
cache service.GatewayCache
|
|
}
|
|
|
|
func (s *GatewayCacheSuite) SetupTest() {
|
|
s.IntegrationRedisSuite.SetupTest()
|
|
s.cache = NewGatewayCache(s.rdb)
|
|
}
|
|
|
|
func (s *GatewayCacheSuite) TestGetSessionAccountID_Missing() {
|
|
_, err := s.cache.GetSessionAccountID(s.ctx, 1, "nonexistent")
|
|
require.True(s.T(), errors.Is(err, service.ErrStickySessionNotFound), "expected ErrStickySessionNotFound for missing session")
|
|
}
|
|
|
|
func (s *GatewayCacheSuite) TestSetAndGetSessionAccountID() {
|
|
sessionID := "s1"
|
|
accountID := int64(99)
|
|
groupID := int64(1)
|
|
sessionTTL := 1 * time.Minute
|
|
|
|
require.NoError(s.T(), s.cache.SetSessionAccountID(s.ctx, groupID, sessionID, accountID, sessionTTL), "SetSessionAccountID")
|
|
|
|
sid, err := s.cache.GetSessionAccountID(s.ctx, groupID, sessionID)
|
|
require.NoError(s.T(), err, "GetSessionAccountID")
|
|
require.Equal(s.T(), accountID, sid, "session id mismatch")
|
|
}
|
|
|
|
func (s *GatewayCacheSuite) TestSessionAccountID_TTL() {
|
|
sessionID := "s2"
|
|
accountID := int64(100)
|
|
groupID := int64(1)
|
|
sessionTTL := 1 * time.Minute
|
|
|
|
require.NoError(s.T(), s.cache.SetSessionAccountID(s.ctx, groupID, sessionID, accountID, sessionTTL), "SetSessionAccountID")
|
|
|
|
sessionKey := buildSessionKey(groupID, sessionID)
|
|
ttl, err := s.rdb.TTL(s.ctx, sessionKey).Result()
|
|
require.NoError(s.T(), err, "TTL sessionKey after Set")
|
|
s.AssertTTLWithin(ttl, 1*time.Second, sessionTTL)
|
|
}
|
|
|
|
func (s *GatewayCacheSuite) TestRefreshSessionTTL() {
|
|
sessionID := "s3"
|
|
accountID := int64(101)
|
|
groupID := int64(1)
|
|
initialTTL := 1 * time.Minute
|
|
refreshTTL := 3 * time.Minute
|
|
|
|
require.NoError(s.T(), s.cache.SetSessionAccountID(s.ctx, groupID, sessionID, accountID, initialTTL), "SetSessionAccountID")
|
|
|
|
require.NoError(s.T(), s.cache.RefreshSessionTTL(s.ctx, groupID, sessionID, refreshTTL), "RefreshSessionTTL")
|
|
|
|
sessionKey := buildSessionKey(groupID, sessionID)
|
|
ttl, err := s.rdb.TTL(s.ctx, sessionKey).Result()
|
|
require.NoError(s.T(), err, "TTL after Refresh")
|
|
s.AssertTTLWithin(ttl, 1*time.Second, refreshTTL)
|
|
}
|
|
|
|
func (s *GatewayCacheSuite) TestRefreshSessionTTL_MissingKey() {
|
|
// RefreshSessionTTL on a missing key should not error (no-op)
|
|
err := s.cache.RefreshSessionTTL(s.ctx, 1, "missing-session", 1*time.Minute)
|
|
require.NoError(s.T(), err, "RefreshSessionTTL on missing key should not error")
|
|
}
|
|
|
|
func (s *GatewayCacheSuite) TestDeleteSessionAccountID() {
|
|
sessionID := "openai:s4"
|
|
accountID := int64(102)
|
|
groupID := int64(1)
|
|
sessionTTL := 1 * time.Minute
|
|
|
|
require.NoError(s.T(), s.cache.SetSessionAccountID(s.ctx, groupID, sessionID, accountID, sessionTTL), "SetSessionAccountID")
|
|
require.NoError(s.T(), s.cache.DeleteSessionAccountID(s.ctx, groupID, sessionID), "DeleteSessionAccountID")
|
|
|
|
_, err := s.cache.GetSessionAccountID(s.ctx, groupID, sessionID)
|
|
require.True(s.T(), errors.Is(err, service.ErrStickySessionNotFound), "expected ErrStickySessionNotFound after delete")
|
|
}
|
|
|
|
func (s *GatewayCacheSuite) TestGetSessionAccountID_CorruptedValue() {
|
|
sessionID := "corrupted"
|
|
groupID := int64(1)
|
|
sessionKey := buildSessionKey(groupID, sessionID)
|
|
|
|
// Set a non-integer value
|
|
require.NoError(s.T(), s.rdb.Set(s.ctx, sessionKey, "not-a-number", 1*time.Minute).Err(), "Set invalid value")
|
|
|
|
_, err := s.cache.GetSessionAccountID(s.ctx, groupID, sessionID)
|
|
require.Error(s.T(), err, "expected error for corrupted value")
|
|
require.False(s.T(), errors.Is(err, service.ErrStickySessionNotFound), "expected parsing error, not a miss")
|
|
}
|
|
|
|
func TestGatewayCacheSuite(t *testing.T) {
|
|
suite.Run(t, new(GatewayCacheSuite))
|
|
}
|