Merge remote-tracking branch 'origin/main' into fix/composite-route-prefix-passthrough

This commit is contained in:
visa2 committed 2026-07-25 22:20:27 +08:00
commit 1614ae9c99
168 files changed
+7169 -412

No files matched your search

+1 -1
View File
@@ -1 +1 @@
0.1.164
0.1.165
+12 -1
View File
@@ -105,6 +105,8 @@ type Group struct {
SortOrder int `json:"sort_order,omitempty"`
// 是否允许 /v1/messages 调度到此 OpenAI 分组
AllowMessagesDispatch bool `json:"allow_messages_dispatch,omitempty"`
// 是否允许此 OpenAI 分组访问 Live 接口
AllowLive bool `json:"allow_live,omitempty"`
// 仅允许非 apikey 类型账号关联到此分组
RequireOauthOnly bool `json:"require_oauth_only,omitempty"`
// 调度时仅允许 privacy 已成功设置的账号
@@ -229,7 +231,7 @@ func (*Group) scanValues(columns []string) ([]any, error) {
switch columns[i] {
case group.FieldModelRouting, group.FieldSupportedModelScopes, group.FieldMessagesDispatchModelConfig, group.FieldModelsListConfig, group.FieldReasoningEffortMappings:
values[i] = new([]byte)
case group.FieldPeakRateEnabled, group.FieldIsExclusive, group.FieldAllowImageGeneration, group.FieldAllowBatchImageGeneration, group.FieldImageRateIndependent, group.FieldVideoRateIndependent, group.FieldClaudeCodeOnly, group.FieldModelRoutingEnabled, group.FieldMcpXMLInject, group.FieldAllowMessagesDispatch, group.FieldRequireOauthOnly, group.FieldRequirePrivacySet:
case group.FieldPeakRateEnabled, group.FieldIsExclusive, group.FieldAllowImageGeneration, group.FieldAllowBatchImageGeneration, group.FieldImageRateIndependent, group.FieldVideoRateIndependent, group.FieldClaudeCodeOnly, group.FieldModelRoutingEnabled, group.FieldMcpXMLInject, group.FieldAllowMessagesDispatch, group.FieldAllowLive, group.FieldRequireOauthOnly, group.FieldRequirePrivacySet:
values[i] = new(sql.NullBool)
case group.FieldRateMultiplier, group.FieldPeakRateMultiplier, group.FieldDailyLimitUsd, group.FieldWeeklyLimitUsd, group.FieldMonthlyLimitUsd, group.FieldImageRateMultiplier, group.FieldImagePrice1k, group.FieldImagePrice2k, group.FieldImagePrice4k, group.FieldBatchImageDiscountMultiplier, group.FieldBatchImageHoldMultiplier, group.FieldVideoRateMultiplier, group.FieldVideoPrice480p, group.FieldVideoPrice720p, group.FieldVideoPrice1080p, group.FieldWebSearchPricePerCall:
values[i] = new(sql.NullFloat64)
@@ -537,6 +539,12 @@ func (_m *Group) assignValues(columns []string, values []any) error {
} else if value.Valid {
_m.AllowMessagesDispatch = value.Bool
}
case group.FieldAllowLive:
if value, ok := values[i].(*sql.NullBool); !ok {
return fmt.Errorf("unexpected type %T for field allow_live", values[i])
} else if value.Valid {
_m.AllowLive = value.Bool
}
case group.FieldRequireOauthOnly:
if value, ok := values[i].(*sql.NullBool); !ok {
return fmt.Errorf("unexpected type %T for field require_oauth_only", values[i])
@@ -826,6 +834,9 @@ func (_m *Group) String() string {
builder.WriteString("allow_messages_dispatch=")
builder.WriteString(fmt.Sprintf("%v", _m.AllowMessagesDispatch))
builder.WriteString(", ")
builder.WriteString("allow_live=")
builder.WriteString(fmt.Sprintf("%v", _m.AllowLive))
builder.WriteString(", ")
builder.WriteString("require_oauth_only=")
builder.WriteString(fmt.Sprintf("%v", _m.RequireOauthOnly))
builder.WriteString(", ")
+10
View File
@@ -102,6 +102,8 @@ const (
FieldSortOrder = "sort_order"
// FieldAllowMessagesDispatch holds the string denoting the allow_messages_dispatch field in the database.
FieldAllowMessagesDispatch = "allow_messages_dispatch"
// FieldAllowLive holds the string denoting the allow_live field in the database.
FieldAllowLive = "allow_live"
// FieldRequireOauthOnly holds the string denoting the require_oauth_only field in the database.
FieldRequireOauthOnly = "require_oauth_only"
// FieldRequirePrivacySet holds the string denoting the require_privacy_set field in the database.
@@ -236,6 +238,7 @@ var Columns = []string{
FieldSupportedModelScopes,
FieldSortOrder,
FieldAllowMessagesDispatch,
FieldAllowLive,
FieldRequireOauthOnly,
FieldRequirePrivacySet,
FieldDefaultMappedModel,
@@ -341,6 +344,8 @@ var (
DefaultSortOrder int
// DefaultAllowMessagesDispatch holds the default value on creation for the "allow_messages_dispatch" field.
DefaultAllowMessagesDispatch bool
// DefaultAllowLive holds the default value on creation for the "allow_live" field.
DefaultAllowLive bool
// DefaultRequireOauthOnly holds the default value on creation for the "require_oauth_only" field.
DefaultRequireOauthOnly bool
// DefaultRequirePrivacySet holds the default value on creation for the "require_privacy_set" field.
@@ -576,6 +581,11 @@ func ByAllowMessagesDispatch(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldAllowMessagesDispatch, opts...).ToFunc()
}
// ByAllowLive orders the results by the allow_live field.
func ByAllowLive(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldAllowLive, opts...).ToFunc()
}
// ByRequireOauthOnly orders the results by the require_oauth_only field.
func ByRequireOauthOnly(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldRequireOauthOnly, opts...).ToFunc()
+15
View File
@@ -260,6 +260,11 @@ func AllowMessagesDispatch(v bool) predicate.Group {
return predicate.Group(sql.FieldEQ(FieldAllowMessagesDispatch, v))
}
// AllowLive applies equality check predicate on the "allow_live" field. It's identical to AllowLiveEQ.
func AllowLive(v bool) predicate.Group {
return predicate.Group(sql.FieldEQ(FieldAllowLive, v))
}
// RequireOauthOnly applies equality check predicate on the "require_oauth_only" field. It's identical to RequireOauthOnlyEQ.
func RequireOauthOnly(v bool) predicate.Group {
return predicate.Group(sql.FieldEQ(FieldRequireOauthOnly, v))
@@ -1985,6 +1990,16 @@ func AllowMessagesDispatchNEQ(v bool) predicate.Group {
return predicate.Group(sql.FieldNEQ(FieldAllowMessagesDispatch, v))
}
// AllowLiveEQ applies the EQ predicate on the "allow_live" field.
func AllowLiveEQ(v bool) predicate.Group {
return predicate.Group(sql.FieldEQ(FieldAllowLive, v))
}
// AllowLiveNEQ applies the NEQ predicate on the "allow_live" field.
func AllowLiveNEQ(v bool) predicate.Group {
return predicate.Group(sql.FieldNEQ(FieldAllowLive, v))
}
// RequireOauthOnlyEQ applies the EQ predicate on the "require_oauth_only" field.
func RequireOauthOnlyEQ(v bool) predicate.Group {
return predicate.Group(sql.FieldEQ(FieldRequireOauthOnly, v))
+65
View File
@@ -607,6 +607,20 @@ func (_c *GroupCreate) SetNillableAllowMessagesDispatch(v *bool) *GroupCreate {
return _c
}
// SetAllowLive sets the "allow_live" field.
func (_c *GroupCreate) SetAllowLive(v bool) *GroupCreate {
_c.mutation.SetAllowLive(v)
return _c
}
// SetNillableAllowLive sets the "allow_live" field if the given value is not nil.
func (_c *GroupCreate) SetNillableAllowLive(v *bool) *GroupCreate {
if v != nil {
_c.SetAllowLive(*v)
}
return _c
}
// SetRequireOauthOnly sets the "require_oauth_only" field.
func (_c *GroupCreate) SetRequireOauthOnly(v bool) *GroupCreate {
_c.mutation.SetRequireOauthOnly(v)
@@ -948,6 +962,10 @@ func (_c *GroupCreate) defaults() error {
v := group.DefaultAllowMessagesDispatch
_c.mutation.SetAllowMessagesDispatch(v)
}
if _, ok := _c.mutation.AllowLive(); !ok {
v := group.DefaultAllowLive
_c.mutation.SetAllowLive(v)
}
if _, ok := _c.mutation.RequireOauthOnly(); !ok {
v := group.DefaultRequireOauthOnly
_c.mutation.SetRequireOauthOnly(v)
@@ -1101,6 +1119,9 @@ func (_c *GroupCreate) check() error {
if _, ok := _c.mutation.AllowMessagesDispatch(); !ok {
return &ValidationError{Name: "allow_messages_dispatch", err: errors.New(`ent: missing required field "Group.allow_messages_dispatch"`)}
}
if _, ok := _c.mutation.AllowLive(); !ok {
return &ValidationError{Name: "allow_live", err: errors.New(`ent: missing required field "Group.allow_live"`)}
}
if _, ok := _c.mutation.RequireOauthOnly(); !ok {
return &ValidationError{Name: "require_oauth_only", err: errors.New(`ent: missing required field "Group.require_oauth_only"`)}
}
@@ -1334,6 +1355,10 @@ func (_c *GroupCreate) createSpec() (*Group, *sqlgraph.CreateSpec) {
_spec.SetField(group.FieldAllowMessagesDispatch, field.TypeBool, value)
_node.AllowMessagesDispatch = value
}
if value, ok := _c.mutation.AllowLive(); ok {
_spec.SetField(group.FieldAllowLive, field.TypeBool, value)
_node.AllowLive = value
}
if value, ok := _c.mutation.RequireOauthOnly(); ok {
_spec.SetField(group.FieldRequireOauthOnly, field.TypeBool, value)
_node.RequireOauthOnly = value
@@ -2224,6 +2249,18 @@ func (u *GroupUpsert) UpdateAllowMessagesDispatch() *GroupUpsert {
return u
}
// SetAllowLive sets the "allow_live" field.
func (u *GroupUpsert) SetAllowLive(v bool) *GroupUpsert {
u.Set(group.FieldAllowLive, v)
return u
}
// UpdateAllowLive sets the "allow_live" field to the value that was provided on create.
func (u *GroupUpsert) UpdateAllowLive() *GroupUpsert {
u.SetExcluded(group.FieldAllowLive)
return u
}
// SetRequireOauthOnly sets the "require_oauth_only" field.
func (u *GroupUpsert) SetRequireOauthOnly(v bool) *GroupUpsert {
u.Set(group.FieldRequireOauthOnly, v)
@@ -3193,6 +3230,20 @@ func (u *GroupUpsertOne) UpdateAllowMessagesDispatch() *GroupUpsertOne {
})
}
// SetAllowLive sets the "allow_live" field.
func (u *GroupUpsertOne) SetAllowLive(v bool) *GroupUpsertOne {
return u.Update(func(s *GroupUpsert) {
s.SetAllowLive(v)
})
}
// UpdateAllowLive sets the "allow_live" field to the value that was provided on create.
func (u *GroupUpsertOne) UpdateAllowLive() *GroupUpsertOne {
return u.Update(func(s *GroupUpsert) {
s.UpdateAllowLive()
})
}
// SetRequireOauthOnly sets the "require_oauth_only" field.
func (u *GroupUpsertOne) SetRequireOauthOnly(v bool) *GroupUpsertOne {
return u.Update(func(s *GroupUpsert) {
@@ -4345,6 +4396,20 @@ func (u *GroupUpsertBulk) UpdateAllowMessagesDispatch() *GroupUpsertBulk {
})
}
// SetAllowLive sets the "allow_live" field.
func (u *GroupUpsertBulk) SetAllowLive(v bool) *GroupUpsertBulk {
return u.Update(func(s *GroupUpsert) {
s.SetAllowLive(v)
})
}
// UpdateAllowLive sets the "allow_live" field to the value that was provided on create.
func (u *GroupUpsertBulk) UpdateAllowLive() *GroupUpsertBulk {
return u.Update(func(s *GroupUpsert) {
s.UpdateAllowLive()
})
}
// SetRequireOauthOnly sets the "require_oauth_only" field.
func (u *GroupUpsertBulk) SetRequireOauthOnly(v bool) *GroupUpsertBulk {
return u.Update(func(s *GroupUpsert) {
+34
View File
@@ -822,6 +822,20 @@ func (_u *GroupUpdate) SetNillableAllowMessagesDispatch(v *bool) *GroupUpdate {
return _u
}
// SetAllowLive sets the "allow_live" field.
func (_u *GroupUpdate) SetAllowLive(v bool) *GroupUpdate {
_u.mutation.SetAllowLive(v)
return _u
}
// SetNillableAllowLive sets the "allow_live" field if the given value is not nil.
func (_u *GroupUpdate) SetNillableAllowLive(v *bool) *GroupUpdate {
if v != nil {
_u.SetAllowLive(*v)
}
return _u
}
// SetRequireOauthOnly sets the "require_oauth_only" field.
func (_u *GroupUpdate) SetRequireOauthOnly(v bool) *GroupUpdate {
_u.mutation.SetRequireOauthOnly(v)
@@ -1495,6 +1509,9 @@ func (_u *GroupUpdate) sqlSave(ctx context.Context) (_node int, err error) {
if value, ok := _u.mutation.AllowMessagesDispatch(); ok {
_spec.SetField(group.FieldAllowMessagesDispatch, field.TypeBool, value)
}
if value, ok := _u.mutation.AllowLive(); ok {
_spec.SetField(group.FieldAllowLive, field.TypeBool, value)
}
if value, ok := _u.mutation.RequireOauthOnly(); ok {
_spec.SetField(group.FieldRequireOauthOnly, field.TypeBool, value)
}
@@ -2627,6 +2644,20 @@ func (_u *GroupUpdateOne) SetNillableAllowMessagesDispatch(v *bool) *GroupUpdate
return _u
}
// SetAllowLive sets the "allow_live" field.
func (_u *GroupUpdateOne) SetAllowLive(v bool) *GroupUpdateOne {
_u.mutation.SetAllowLive(v)
return _u
}
// SetNillableAllowLive sets the "allow_live" field if the given value is not nil.
func (_u *GroupUpdateOne) SetNillableAllowLive(v *bool) *GroupUpdateOne {
if v != nil {
_u.SetAllowLive(*v)
}
return _u
}
// SetRequireOauthOnly sets the "require_oauth_only" field.
func (_u *GroupUpdateOne) SetRequireOauthOnly(v bool) *GroupUpdateOne {
_u.mutation.SetRequireOauthOnly(v)
@@ -3330,6 +3361,9 @@ func (_u *GroupUpdateOne) sqlSave(ctx context.Context) (_node *Group, err error)
if value, ok := _u.mutation.AllowMessagesDispatch(); ok {
_spec.SetField(group.FieldAllowMessagesDispatch, field.TypeBool, value)
}
if value, ok := _u.mutation.AllowLive(); ok {
_spec.SetField(group.FieldAllowLive, field.TypeBool, value)
}
if value, ok := _u.mutation.RequireOauthOnly(); ok {
_spec.SetField(group.FieldRequireOauthOnly, field.TypeBool, value)
}
+1
View File
@@ -938,6 +938,7 @@ var (
{Name: "supported_model_scopes", Type: field.TypeJSON, SchemaType: map[string]string{"postgres": "jsonb"}},
{Name: "sort_order", Type: field.TypeInt, Default: 0},
{Name: "allow_messages_dispatch", Type: field.TypeBool, Default: false},
{Name: "allow_live", Type: field.TypeBool, Default: false},
{Name: "require_oauth_only", Type: field.TypeBool, Default: false},
{Name: "require_privacy_set", Type: field.TypeBool, Default: false},
{Name: "default_mapped_model", Type: field.TypeString, Size: 100, Default: ""},
+55 -1
View File
@@ -21911,6 +21911,7 @@ type GroupMutation struct {
sort_order *int
addsort_order *int
allow_messages_dispatch *bool
allow_live *bool
require_oauth_only *bool
require_privacy_set *bool
default_mapped_model *string
@@ -24226,6 +24227,42 @@ func (m *GroupMutation) ResetAllowMessagesDispatch() {
m.allow_messages_dispatch = nil
}
// SetAllowLive sets the "allow_live" field.
func (m *GroupMutation) SetAllowLive(b bool) {
m.allow_live = &b
}
// AllowLive returns the value of the "allow_live" field in the mutation.
func (m *GroupMutation) AllowLive() (r bool, exists bool) {
v := m.allow_live
if v == nil {
return
}
return *v, true
}
// OldAllowLive returns the old "allow_live" field's value of the Group entity.
// If the Group object wasn't provided to the builder, the object is fetched from the database.
// An error is returned if the mutation operation is not UpdateOne, or the database query fails.
func (m *GroupMutation) OldAllowLive(ctx context.Context) (v bool, err error) {
if !m.op.Is(OpUpdateOne) {
return v, errors.New("OldAllowLive is only allowed on UpdateOne operations")
}
if m.id == nil || m.oldValue == nil {
return v, errors.New("OldAllowLive requires an ID field in the mutation")
}
oldValue, err := m.oldValue(ctx)
if err != nil {
return v, fmt.Errorf("querying old value for OldAllowLive: %w", err)
}
return oldValue.AllowLive, nil
}
// ResetAllowLive resets all changes to the "allow_live" field.
func (m *GroupMutation) ResetAllowLive() {
m.allow_live = nil
}
// SetRequireOauthOnly sets the "require_oauth_only" field.
func (m *GroupMutation) SetRequireOauthOnly(b bool) {
m.require_oauth_only = &b
@@ -24907,7 +24944,7 @@ func (m *GroupMutation) Type() string {
// order to get all numeric fields that were incremented/decremented, call
// AddedFields().
func (m *GroupMutation) Fields() []string {
fields := make([]string, 0, 51)
fields := make([]string, 0, 52)
if m.created_at != nil {
fields = append(fields, group.FieldCreatedAt)
}
@@ -25037,6 +25074,9 @@ func (m *GroupMutation) Fields() []string {
if m.allow_messages_dispatch != nil {
fields = append(fields, group.FieldAllowMessagesDispatch)
}
if m.allow_live != nil {
fields = append(fields, group.FieldAllowLive)
}
if m.require_oauth_only != nil {
fields = append(fields, group.FieldRequireOauthOnly)
}
@@ -25155,6 +25195,8 @@ func (m *GroupMutation) Field(name string) (ent.Value, bool) {
return m.SortOrder()
case group.FieldAllowMessagesDispatch:
return m.AllowMessagesDispatch()
case group.FieldAllowLive:
return m.AllowLive()
case group.FieldRequireOauthOnly:
return m.RequireOauthOnly()
case group.FieldRequirePrivacySet:
@@ -25266,6 +25308,8 @@ func (m *GroupMutation) OldField(ctx context.Context, name string) (ent.Value, e
return m.OldSortOrder(ctx)
case group.FieldAllowMessagesDispatch:
return m.OldAllowMessagesDispatch(ctx)
case group.FieldAllowLive:
return m.OldAllowLive(ctx)
case group.FieldRequireOauthOnly:
return m.OldRequireOauthOnly(ctx)
case group.FieldRequirePrivacySet:
@@ -25592,6 +25636,13 @@ func (m *GroupMutation) SetField(name string, value ent.Value) error {
}
m.SetAllowMessagesDispatch(v)
return nil
case group.FieldAllowLive:
v, ok := value.(bool)
if !ok {
return fmt.Errorf("unexpected type %T for field %s", value, name)
}
m.SetAllowLive(v)
return nil
case group.FieldRequireOauthOnly:
v, ok := value.(bool)
if !ok {
@@ -26180,6 +26231,9 @@ func (m *GroupMutation) ResetField(name string) error {
case group.FieldAllowMessagesDispatch:
m.ResetAllowMessagesDispatch()
return nil
case group.FieldAllowLive:
m.ResetAllowLive()
return nil
case group.FieldRequireOauthOnly:
m.ResetRequireOauthOnly()
return nil
+12 -8
View File
@@ -1141,40 +1141,44 @@ func init() {
groupDescAllowMessagesDispatch := groupFields[39].Descriptor()
// group.DefaultAllowMessagesDispatch holds the default value on creation for the allow_messages_dispatch field.
group.DefaultAllowMessagesDispatch = groupDescAllowMessagesDispatch.Default.(bool)
// groupDescAllowLive is the schema descriptor for allow_live field.
groupDescAllowLive := groupFields[40].Descriptor()
// group.DefaultAllowLive holds the default value on creation for the allow_live field.
group.DefaultAllowLive = groupDescAllowLive.Default.(bool)
// groupDescRequireOauthOnly is the schema descriptor for require_oauth_only field.
groupDescRequireOauthOnly := groupFields[40].Descriptor()
groupDescRequireOauthOnly := groupFields[41].Descriptor()
// group.DefaultRequireOauthOnly holds the default value on creation for the require_oauth_only field.
group.DefaultRequireOauthOnly = groupDescRequireOauthOnly.Default.(bool)
// groupDescRequirePrivacySet is the schema descriptor for require_privacy_set field.
groupDescRequirePrivacySet := groupFields[41].Descriptor()
groupDescRequirePrivacySet := groupFields[42].Descriptor()
// group.DefaultRequirePrivacySet holds the default value on creation for the require_privacy_set field.
group.DefaultRequirePrivacySet = groupDescRequirePrivacySet.Default.(bool)
// groupDescDefaultMappedModel is the schema descriptor for default_mapped_model field.
groupDescDefaultMappedModel := groupFields[42].Descriptor()
groupDescDefaultMappedModel := groupFields[43].Descriptor()
// group.DefaultDefaultMappedModel holds the default value on creation for the default_mapped_model field.
group.DefaultDefaultMappedModel = groupDescDefaultMappedModel.Default.(string)
// group.DefaultMappedModelValidator is a validator for the "default_mapped_model" field. It is called by the builders before save.
group.DefaultMappedModelValidator = groupDescDefaultMappedModel.Validators[0].(func(string) error)
// groupDescMessagesDispatchModelConfig is the schema descriptor for messages_dispatch_model_config field.
groupDescMessagesDispatchModelConfig := groupFields[43].Descriptor()
groupDescMessagesDispatchModelConfig := groupFields[44].Descriptor()
// group.DefaultMessagesDispatchModelConfig holds the default value on creation for the messages_dispatch_model_config field.
group.DefaultMessagesDispatchModelConfig = groupDescMessagesDispatchModelConfig.Default.(domain.OpenAIMessagesDispatchModelConfig)
// groupDescModelsListConfig is the schema descriptor for models_list_config field.
groupDescModelsListConfig := groupFields[44].Descriptor()
groupDescModelsListConfig := groupFields[45].Descriptor()
// group.DefaultModelsListConfig holds the default value on creation for the models_list_config field.
group.DefaultModelsListConfig = groupDescModelsListConfig.Default.(domain.GroupModelsListConfig)
// groupDescRpmLimit is the schema descriptor for rpm_limit field.
groupDescRpmLimit := groupFields[45].Descriptor()
groupDescRpmLimit := groupFields[46].Descriptor()
// group.DefaultRpmLimit holds the default value on creation for the rpm_limit field.
group.DefaultRpmLimit = groupDescRpmLimit.Default.(int)
// groupDescMaxReasoningEffort is the schema descriptor for max_reasoning_effort field.
groupDescMaxReasoningEffort := groupFields[46].Descriptor()
groupDescMaxReasoningEffort := groupFields[47].Descriptor()
// group.DefaultMaxReasoningEffort holds the default value on creation for the max_reasoning_effort field.
group.DefaultMaxReasoningEffort = groupDescMaxReasoningEffort.Default.(string)
// group.MaxReasoningEffortValidator is a validator for the "max_reasoning_effort" field. It is called by the builders before save.
group.MaxReasoningEffortValidator = groupDescMaxReasoningEffort.Validators[0].(func(string) error)
// groupDescReasoningEffortMappings is the schema descriptor for reasoning_effort_mappings field.
groupDescReasoningEffortMappings := groupFields[47].Descriptor()
groupDescReasoningEffortMappings := groupFields[48].Descriptor()
// group.DefaultReasoningEffortMappings holds the default value on creation for the reasoning_effort_mappings field.
group.DefaultReasoningEffortMappings = groupDescReasoningEffortMappings.Default.([]domain.ReasoningEffortMapping)
idempotencyrecordMixin := schema.IdempotencyRecord{}.Mixin()
+3
View File
@@ -198,6 +198,9 @@ func (Group) Fields() []ent.Field {
field.Bool("allow_messages_dispatch").
Default(false).
Comment("是否允许 /v1/messages 调度到此 OpenAI 分组"),
field.Bool("allow_live").
Default(false).
Comment("是否允许此 OpenAI 分组访问 Live 接口"),
field.Bool("require_oauth_only").
Default(false).
Comment("仅允许非 apikey 类型账号关联到此分组"),
+11
View File
@@ -913,6 +913,8 @@ type GatewayConfig struct {
OpenAICompactModel string `mapstructure:"openai_compact_model"`
// OpenAIWS: OpenAI Responses WebSocket 配置(默认开启,可按需回滚到 HTTP)
OpenAIWS GatewayOpenAIWSConfig `mapstructure:"openai_ws"`
// Live: ChatGPT Frameless Live 会话配置。
Live GatewayLiveConfig `mapstructure:"live"`
// OpenAIScheduler: OpenAI 高级调度器粘性逃逸配置
OpenAIScheduler GatewayOpenAISchedulerConfig `mapstructure:"openai_scheduler"`
// OpenAIHTTP2: OpenAI HTTP 上游协议策略(默认启用 HTTP/2,可按代理能力回退 HTTP/1.1)
@@ -999,6 +1001,11 @@ type GatewayConfig struct {
UserMessageQueue UserMessageQueueConfig `mapstructure:"user_message_queue"`
}
type GatewayLiveConfig struct {
// MaxSessionDurationSeconds 是 Live 会话的硬上限。
MaxSessionDurationSeconds int `mapstructure:"max_session_duration_seconds"`
}
// GatewayOpenAIHTTP2Config OpenAI HTTP 上游协议配置。
// 默认启用 HTTP/2;在部分代理不兼容时按策略回退 HTTP/1.1。
type GatewayOpenAIHTTP2Config struct {
@@ -2188,6 +2195,7 @@ func setDefaults() {
viper.SetDefault("gateway.codex_image_generation_bridge_enabled", false)
viper.SetDefault("gateway.openai_passthrough_allow_timeout_headers", false)
viper.SetDefault("gateway.openai_compact_model", "gpt-5.4")
viper.SetDefault("gateway.live.max_session_duration_seconds", 3600)
// OpenAI Responses WebSocket(默认开启;可通过 force_http 紧急回滚)
viper.SetDefault("gateway.openai_ws.enabled", true)
viper.SetDefault("gateway.openai_ws.mode_router_v2_enabled", false)
@@ -3039,6 +3047,9 @@ func (c *Config) Validate() error {
(c.Gateway.OpenAIHighEffortFirstOutputTimeoutSeconds > 0 && c.Gateway.OpenAIHighEffortFirstOutputTimeoutSeconds < 30) {
return fmt.Errorf("gateway.openai_high_effort_first_output_timeout_seconds must be 0 or between 30-1800 seconds")
}
if c.Gateway.Live.MaxSessionDurationSeconds <= 0 {
c.Gateway.Live.MaxSessionDurationSeconds = 3600
}
if strings.TrimSpace(c.Gateway.ConnectionPoolIsolation) != "" {
switch c.Gateway.ConnectionPoolIsolation {
case ConnectionPoolIsolationProxy, ConnectionPoolIsolationAccount, ConnectionPoolIsolationAccountProxy:
+1
View File
@@ -133,6 +133,7 @@ var DefaultBedrockModelMapping = map[string]string{
// Claude Fable
"claude-fable-5": "anthropic.claude-fable-5",
// Claude Opus
"claude-opus-5": "us.anthropic.claude-opus-5-v1",
"claude-opus-4-8": "us.anthropic.claude-opus-4-8-v1",
"claude-opus-4-7": "us.anthropic.claude-opus-4-7-v1",
"claude-opus-4-6-thinking": "us.anthropic.claude-opus-4-6-v1",
@@ -61,7 +61,7 @@ func (r *ollamaCloudUsageHandlerTestRepo) UpdateOllamaCloudUsageSnapshot(context
func (r *ollamaCloudUsageHandlerTestRepo) DisableOllamaCloudUsageAutoRefresh(context.Context, *service.Account) error {
return nil
}
func (r *ollamaCloudUsageHandlerTestRepo) ListDueOllamaCloudUsageAccounts(context.Context, time.Time, int) ([]service.Account, error) {
func (r *ollamaCloudUsageHandlerTestRepo) ListDueOllamaCloudUsageAccounts(context.Context, time.Time, time.Duration, time.Duration, int) ([]service.Account, error) {
return nil, nil
}
@@ -273,4 +273,5 @@ func TestGetOllamaCloudUsageSettingsHandlerSuccess(t *testing.T) {
require.Equal(t, http.StatusOK, recorder.Code)
require.Contains(t, recorder.Body.String(), `"enabled":false`)
require.Contains(t, recorder.Body.String(), `"interval_minutes":60`)
require.Contains(t, recorder.Body.String(), `"debounce_minutes":1`)
}
@@ -13,6 +13,7 @@ import (
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
"github.com/Wei-Shaw/sub2api/internal/platform/liveattestation"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
@@ -25,6 +26,16 @@ type GroupHandler struct {
groupCapacityService *service.GroupCapacityService
}
// GetLiveCapability 返回当前服务端是否具备生成 Live attestation 的运行环境。
func (h *GroupHandler) GetLiveCapability(c *gin.Context) {
err := liveattestation.NewProvider().Check(c.Request.Context())
result := gin.H{"supported": err == nil}
if err != nil {
result["reason"] = err.Error()
}
response.Success(c, result)
}
type optionalLimitField struct {
set bool
value *float64
@@ -125,6 +136,7 @@ type CreateGroupRequest struct {
SupportedModelScopes []string `json:"supported_model_scopes"`
// OpenAI Messages 调度配置(仅 openai 平台使用)
AllowMessagesDispatch bool `json:"allow_messages_dispatch"`
AllowLive bool `json:"allow_live"`
RequireOAuthOnly bool `json:"require_oauth_only"`
RequirePrivacySet bool `json:"require_privacy_set"`
DefaultMappedModel string `json:"default_mapped_model"`
@@ -183,6 +195,7 @@ type UpdateGroupRequest struct {
SupportedModelScopes *[]string `json:"supported_model_scopes"`
// OpenAI Messages 调度配置(仅 openai 平台使用)
AllowMessagesDispatch *bool `json:"allow_messages_dispatch"`
AllowLive *bool `json:"allow_live"`
RequireOAuthOnly *bool `json:"require_oauth_only"`
RequirePrivacySet *bool `json:"require_privacy_set"`
DefaultMappedModel *string `json:"default_mapped_model"`
@@ -499,6 +512,7 @@ func (h *GroupHandler) Create(c *gin.Context) {
MCPXMLInject: req.MCPXMLInject,
SupportedModelScopes: req.SupportedModelScopes,
AllowMessagesDispatch: req.AllowMessagesDispatch,
AllowLive: req.AllowLive,
RequireOAuthOnly: req.RequireOAuthOnly,
RequirePrivacySet: req.RequirePrivacySet,
DefaultMappedModel: req.DefaultMappedModel,
@@ -617,6 +631,7 @@ func (h *GroupHandler) Update(c *gin.Context) {
MCPXMLInject: req.MCPXMLInject,
SupportedModelScopes: req.SupportedModelScopes,
AllowMessagesDispatch: req.AllowMessagesDispatch,
AllowLive: req.AllowLive,
RequireOAuthOnly: req.RequireOAuthOnly,
RequirePrivacySet: req.RequirePrivacySet,
DefaultMappedModel: req.DefaultMappedModel,
@@ -2978,6 +2978,17 @@ func (r *oauthPendingFlowUserRepo) Create(ctx context.Context, user *service.Use
return nil
}
func (r *oauthPendingFlowUserRepo) CreateWithEmailAliasGuard(ctx context.Context, user *service.User) error {
aliasExists, err := r.ExistsByEmailAlias(ctx, user.Email)
if err != nil {
return err
}
if aliasExists {
return service.ErrEmailExists
}
return r.Create(ctx, user)
}
func (r *oauthPendingFlowUserRepo) GetByID(ctx context.Context, id int64) (*service.User, error) {
entity, err := r.client.User.Get(ctx, id)
if err != nil {
@@ -3185,6 +3196,20 @@ func (r *oauthPendingFlowUserRepo) ExistsByEmail(ctx context.Context, email stri
return count > 0, err
}
func (r *oauthPendingFlowUserRepo) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) {
identity := service.NormalizeEmailForAliasDedup(email)
emails, err := r.client.User.Query().Select(dbuser.FieldEmail).Strings(ctx)
if err != nil {
return false, err
}
for _, stored := range emails {
if service.NormalizeEmailForAliasDedup(stored) == identity {
return true, nil
}
}
return false, nil
}
func (r *oauthPendingFlowUserRepo) RemoveGroupFromAllowedGroups(context.Context, int64) (int64, error) {
panic("unexpected RemoveGroupFromAllowedGroups call")
}
@@ -42,6 +42,9 @@ func (h *BatchImageHandler) Submit(c *gin.Context) {
if !h.checkSecurityAuditBeforeSubmit(c, &req) {
return
}
if sessionID := service.ExtractClientSessionID(c); sessionID != "" {
req.SessionID = &sessionID
}
got, err := h.service.Submit(c.Request.Context(), owner, req, c.GetHeader("Idempotency-Key"))
if err != nil {
batchImageError(c, err)
+2
View File
@@ -204,6 +204,7 @@ func groupFromServiceBase(g *service.Group) Group {
FallbackGroupID: g.FallbackGroupID,
FallbackGroupIDOnInvalidRequest: g.FallbackGroupIDOnInvalidRequest,
AllowMessagesDispatch: g.AllowMessagesDispatch,
AllowLive: g.AllowLive,
RequireOAuthOnly: g.RequireOAuthOnly,
RequirePrivacySet: g.RequirePrivacySet,
RPMLimit: g.RPMLimit,
@@ -669,6 +670,7 @@ func usageLogFromServiceUser(l *service.UsageLog) UsageLog {
MediaType: l.MediaType,
UserAgent: l.UserAgent,
IPAddress: l.IPAddress,
SessionID: l.SessionID,
CacheTTLOverridden: l.CacheTTLOverridden,
BillingMode: l.BillingMode,
CreatedAt: l.CreatedAt,
+5
View File
@@ -132,6 +132,8 @@ type Group struct {
// OpenAI Messages 调度开关(用户侧需要此字段判断是否展示 Claude Code 教程)
AllowMessagesDispatch bool `json:"allow_messages_dispatch"`
// OpenAI Live 接口开关
AllowLive bool `json:"allow_live"`
// 账号过滤控制(仅 OpenAI/Antigravity 平台有效)
RequireOAuthOnly bool `json:"require_oauth_only"`
@@ -521,6 +523,9 @@ type UsageLog struct {
UserAgent *string `json:"user_agent"`
// IPAddress is visible to the owner of the usage record.
IPAddress *string `json:"ip_address,omitempty"`
// SessionID is the explicit client-provided request correlation identifier
// (e.g. the session_id / X-Session-Id headers). Omitted when absent.
SessionID *string `json:"session_id,omitempty"`
// Cache TTL Override 标记
CacheTTLOverridden bool `json:"cache_ttl_overridden"`
@@ -535,6 +535,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
// ForceCacheBilling 提前拍成标量,避免 worker 闭包保活 failover 状态里的响应体。
forceCacheBilling := fs.ForceCacheBilling
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
Result: result,
@@ -547,6 +548,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
UpstreamEndpoint: upstreamEndpoint,
UserAgent: userAgent,
IPAddress: clientIP,
SessionID: sessionID,
RequestPayloadHash: requestPayloadHash,
ForceCacheBilling: forceCacheBilling,
APIKeyService: h.apiKeyService,
@@ -970,6 +972,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
// ForceCacheBilling 提前拍成标量,避免 worker 闭包保活 failover 状态里的响应体。
forceCacheBilling := fs.ForceCacheBilling
quotaPlatform := service.QuotaPlatform(c.Request.Context(), currentAPIKey)
sessionID := service.ExtractClientSessionID(c)
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
Result: result,
@@ -982,6 +985,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
UpstreamEndpoint: upstreamEndpoint,
UserAgent: userAgent,
IPAddress: clientIP,
SessionID: sessionID,
RequestPayloadHash: requestPayloadHash,
ForceCacheBilling: forceCacheBilling,
APIKeyService: h.apiKeyService,
@@ -298,6 +298,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
Result: result,
@@ -312,6 +313,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
IPAddress: clientIP,
RequestPayloadHash: requestPayloadHash,
APIKeyService: h.apiKeyService,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
}); err != nil {
reqLog.Error("gateway.cc.record_usage_failed",
@@ -276,6 +276,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
Result: result,
@@ -290,6 +291,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
IPAddress: clientIP,
RequestPayloadHash: requestPayloadHash,
APIKeyService: h.apiKeyService,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
}); err != nil {
reqLog.Error("gateway.responses.record_usage_failed",
@@ -532,6 +532,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
// ForceCacheBilling 提前拍成标量,避免 worker 闭包保活 failover 状态里的响应体。
forceCacheBilling := fs.ForceCacheBilling
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
h.submitUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
if err := h.gatewayService.RecordUsageWithLongContext(ctx, &service.RecordUsageLongContextInput{
Result: result,
@@ -549,6 +550,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
LongContextMultiplier: 2.0, // 超出部分双倍计费
ForceCacheBilling: forceCacheBilling,
APIKeyService: h.apiKeyService,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
}); err != nil {
logger.L().With(
+2
View File
@@ -465,6 +465,7 @@ func recordGrokMediaUsage(
) {
userAgent := c.GetHeader("User-Agent")
clientIP := ip.GetClientIP(c)
sessionID := service.ExtractClientSessionID(c)
payloadForHash := body
if len(payloadForHash) == 0 && strings.TrimSpace(requestID) != "" {
payloadForHash = []byte(requestID)
@@ -493,6 +494,7 @@ func recordGrokMediaUsage(
RequestPayloadHash: service.HashUsageRequestPayload(payloadForHash),
APIKeyService: h.apiKeyService,
QuotaPlatform: quotaPlatform,
SessionID: sessionID,
ChannelUsageFields: channelUsageFields,
}); err != nil {
logger.L().With(
@@ -233,6 +233,7 @@ func (h *OpenAIGatewayHandler) recordAlphaSearchUsage(
) {
userAgent := c.GetHeader("User-Agent")
clientIP := ip.GetClientIP(c)
sessionID := service.ExtractClientSessionID(c)
requestPayloadHash := service.HashUsageRequestPayload(body)
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
@@ -252,6 +253,7 @@ func (h *OpenAIGatewayHandler) recordAlphaSearchUsage(
RequestPayloadHash: requestPayloadHash,
APIKeyService: h.apiKeyService,
QuotaPlatform: quotaPlatform,
SessionID: sessionID,
ChannelUsageFields: channelMapping.ToUsageFields(requestedModel, result.UpstreamModel),
}); err != nil {
logger.L().With(
@@ -340,6 +340,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) {
@@ -355,6 +356,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
IPAddress: clientIP,
APIKeyService: h.apiKeyService,
QuotaPlatform: quotaPlatform,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
CyberBlocked: cyberBlocked,
}); err != nil {
@@ -243,6 +243,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
@@ -257,6 +258,7 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
IPAddress: clientIP,
APIKeyService: h.apiKeyService,
QuotaPlatform: quotaPlatform,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
}); err != nil {
logger.L().With(
@@ -627,6 +627,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
// 使用量记录通过有界 worker 池提交,避免请求热路径创建无界 goroutine。
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
@@ -644,6 +645,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
RequestPayloadHash: requestPayloadHash,
APIKeyService: h.apiKeyService,
QuotaPlatform: quotaPlatform,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
CyberBlocked: cyberBlocked,
}); err != nil {
@@ -1138,6 +1140,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) {
@@ -1154,6 +1157,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
RequestPayloadHash: requestPayloadHash,
APIKeyService: h.apiKeyService,
QuotaPlatform: quotaPlatform,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMappingMsg, reqModel, result.UpstreamModel),
CyberBlocked: cyberBlocked,
}); err != nil {
@@ -1862,6 +1866,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
h.submitOpenAIUsageRecordTask(ctx, result, func(taskCtx context.Context) {
if err := h.gatewayService.RecordUsage(taskCtx, &service.OpenAIRecordUsageInput{
@@ -1877,6 +1882,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
RequestPayloadHash: requestPayloadHash,
APIKeyService: h.apiKeyService,
QuotaPlatform: quotaPlatform,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMappingWS, reqModel, result.UpstreamModel),
CyberBlocked: cyberBlocked,
}); err != nil {
@@ -2884,6 +2890,8 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey
userAgent = c.GetHeader("User-Agent")
clientIPStr = strings.TrimSpace(ip.GetClientIP(c))
}
// 提前拍成标量,避免在下方 goroutine 内访问 gin.Context。
sessionID := service.ExtractClientSessionID(c)
apiKeyPrefix := ""
if apiKey != nil {
apiKeyPrefix = keyPrefix(apiKey.Key, 8)
@@ -2940,6 +2948,7 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey
UpstreamEndpoint: upstreamEndpoint,
UserAgent: userAgent,
IPAddress: clientIPStr,
SessionID: sessionID,
RequestPayloadHash: requestPayloadHash,
APIKeyService: apiKeySvc,
ChannelUsageFields: channelFields,
@@ -91,6 +91,8 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
zap.Bool("stream", parsed.Stream),
zap.Bool("multipart", parsed.Multipart),
zap.String("capability", string(parsed.RequiredCapability)),
zap.String("img_quality", parsed.Quality),
zap.String("img_size", parsed.Size),
)
if !service.GroupAllowsImageGeneration(apiKey.Group) {
@@ -374,6 +376,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
if result != nil {
upstreamModel = result.UpstreamModel
}
sessionID := service.ExtractClientSessionID(c)
h.submitMandatoryUsageRecordTask(c.Request.Context(), func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
Result: result,
@@ -388,6 +391,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
RequestPayloadHash: requestPayloadHash,
APIKeyService: h.apiKeyService,
QuotaPlatform: quotaPlatform,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, requestModel, upstreamModel),
}); err != nil {
logger.L().With(
@@ -12,11 +12,14 @@ import (
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
"go.uber.org/zap"
"go.uber.org/zap/zaptest/observer"
)
type openAIImagesFailoverAccountRepo struct {
@@ -154,8 +157,10 @@ func TestOpenAIGatewayHandlerImages_ServerErrorFailsOverAndReturnsClearErrorWhen
)
handler.maxAccountSwitches = 10
body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat"}`)
req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
body := []byte(`{"model":"gpt-image-2","prompt":"draw a cat","quality":"high","size":"1536x1024"}`)
core, observedLogs := observer.New(zap.DebugLevel)
requestCtx := logger.IntoContext(context.Background(), zap.New(core))
req := httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body)).WithContext(requestCtx)
req.Header.Set("Content-Type", "application/json")
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
@@ -173,6 +178,16 @@ func TestOpenAIGatewayHandlerImages_ServerErrorFailsOverAndReturnsClearErrorWhen
handler.Images(c)
accountSelectingLogs := observedLogs.FilterMessage("openai.images.account_selecting").All()
require.NotEmpty(t, accountSelectingLogs)
loggedFields := make(map[string]string)
for _, field := range accountSelectingLogs[0].Context {
loggedFields[field.Key] = field.String
}
require.Equal(t, "high", loggedFields["img_quality"])
require.Equal(t, "1536x1024", loggedFields["img_size"])
require.NotContains(t, loggedFields, "prompt")
require.Equal(t, []int64{1, 2}, upstream.calls())
require.Equal(t, http.StatusBadGateway, rec.Code)
require.Equal(t, "upstream_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String())
+236
View File
@@ -0,0 +1,236 @@
package handler
import (
"encoding/json"
"errors"
"io"
"net/http"
"net/url"
"strconv"
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
coderws "github.com/coder/websocket"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
"go.uber.org/zap"
)
func (h *OpenAIGatewayHandler) Live(c *gin.Context) {
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
if !ok {
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
return
}
subject, ok := middleware2.GetAuthSubjectFromContext(c)
if !ok {
h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
return
}
if apiKey.Group == nil || apiKey.Group.Platform != service.PlatformOpenAI {
h.errorResponse(c, http.StatusNotFound, "not_found_error", "Live is not supported for this platform")
return
}
if !liveEnabledForAPIKey(apiKey) {
h.errorResponse(c, http.StatusForbidden, "permission_error", "Live is not enabled for this group")
return
}
request, err := parseLiveCallRequest(c)
if err != nil {
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", err.Error())
return
}
model := strings.TrimSpace(gjson.GetBytes(request.Session, "model").String())
reqLog := requestLogger(
c,
"handler.openai_gateway.live",
zap.Int64("user_id", subject.UserID),
zap.Int64("api_key_id", apiKey.ID),
zap.Any("group_id", apiKey.GroupID),
)
if decision := h.checkSecurityAudit(
c,
reqLog,
apiKey,
subject,
service.ContentModerationProtocolOpenAIResponses,
model,
request.Session,
); decision != nil && !decision.AllowNextStage {
h.openAISecurityAuditError(c, decision)
return
}
subscription, _ := middleware2.GetSubscriptionFromContext(c)
if h.billingCacheService == nil {
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", "Billing service unavailable")
return
}
if err := h.billingCacheService.CheckBillingEligibility(
c.Request.Context(),
apiKey.User,
apiKey,
apiKey.Group,
subscription,
service.QuotaPlatform(c.Request.Context(), apiKey),
); err != nil {
status, code, message, retryAfter := billingErrorDetails(err)
if retryAfter > 0 {
c.Header("Retry-After", strconv.Itoa(retryAfter))
}
h.errorResponse(c, status, code, message)
return
}
userRelease, acquired, err := h.concurrencyHelper.TryAcquireUserSlot(
c.Request.Context(),
subject.UserID,
subject.Concurrency,
)
if err != nil {
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", "Live concurrency unavailable")
return
}
if !acquired {
h.errorResponse(c, http.StatusTooManyRequests, "rate_limit_error", "Live concurrency limit reached")
return
}
defer userRelease()
identity := liveCallIdentity(c, apiKey, subject.UserID, subscription)
created, err := h.gatewayService.CreateLiveCall(c.Request.Context(), request, identity, subject.Concurrency)
if err != nil {
h.writeLiveCreateError(c, err)
return
}
c.Header("Location", liveSidebandLocation(c.FullPath(), created.CallID))
c.Data(http.StatusOK, "application/sdp", created.SDP)
}
func parseLiveCallRequest(c *gin.Context) (*service.LiveCallRequest, error) {
contentType := strings.ToLower(c.GetHeader("Content-Type"))
if strings.HasPrefix(contentType, "multipart/form-data") {
sdp := c.PostForm("sdp")
session := json.RawMessage(c.PostForm("session"))
request := &service.LiveCallRequest{SDP: sdp, Session: session}
if err := service.ValidateLiveCallRequest(request); err != nil {
return nil, err
}
return request, nil
}
var request service.LiveCallRequest
decoder := json.NewDecoder(c.Request.Body)
if err := decoder.Decode(&request); err != nil {
return nil, errors.New("request body must be valid JSON")
}
if err := decoder.Decode(&struct{}{}); !errors.Is(err, io.EOF) {
return nil, errors.New("request body must contain one JSON object")
}
if err := service.ValidateLiveCallRequest(&request); err != nil {
return nil, err
}
return &request, nil
}
func liveSidebandLocation(fullPath, callID string) string {
prefix := "/v1/live/"
if strings.HasPrefix(fullPath, "/backend-api/codex/") {
prefix = "/backend-api/codex/"
}
return prefix + url.PathEscape(callID)
}
func liveCallIdentity(
c *gin.Context,
apiKey *service.APIKey,
userID int64,
subscription *service.UserSubscription,
) service.LiveCallIdentity {
var subscriptionID *int64
if subscription != nil {
value := subscription.ID
subscriptionID = &value
}
return service.LiveCallIdentity{
APIKeyID: apiKey.ID,
UserID: userID,
GroupID: apiKey.GroupID,
SubscriptionID: subscriptionID,
UserAgent: c.GetHeader("User-Agent"),
IPAddress: ip.GetClientIP(c),
InboundEndpoint: GetInboundEndpoint(c),
}
}
func (h *OpenAIGatewayHandler) writeLiveCreateError(c *gin.Context, err error) {
switch {
case errors.Is(err, service.ErrLiveConcurrencyFull):
h.errorResponse(c, http.StatusTooManyRequests, "rate_limit_error", "Live concurrency limit reached")
case errors.Is(err, service.ErrLiveUnavailable):
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", "Live is unavailable")
default:
var attestationErr *service.LiveAttestationUnavailableError
if errors.As(err, &attestationErr) {
h.errorResponse(c, http.StatusServiceUnavailable, "api_error", attestationErr.Error())
return
}
var upstreamErr *service.UpstreamFailoverError
if errors.As(err, &upstreamErr) && upstreamErr.StatusCode >= 400 && upstreamErr.StatusCode < 500 {
h.errorResponse(c, upstreamErr.StatusCode, "invalid_request_error", "Live upstream rejected the request")
return
}
h.errorResponse(c, http.StatusBadGateway, "api_error", "Live upstream request failed")
}
}
func (h *OpenAIGatewayHandler) LiveSideband(c *gin.Context) {
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
if !ok {
h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
return
}
subject, ok := middleware2.GetAuthSubjectFromContext(c)
if !ok {
h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
return
}
if !liveEnabledForAPIKey(apiKey) {
h.errorResponse(c, http.StatusForbidden, "permission_error", "Live is not enabled for this group")
return
}
identity := service.LiveCallIdentity{
APIKeyID: apiKey.ID,
UserID: subject.UserID,
GroupID: apiKey.GroupID,
}
record, err := h.gatewayService.GetLiveCallForIdentity(c.Request.Context(), c.Param("call_id"), identity)
if err != nil {
if errors.Is(err, service.ErrLiveIdentityMismatch) {
h.errorResponse(c, http.StatusForbidden, "permission_error", "Live call belongs to another identity")
return
}
h.errorResponse(c, http.StatusNotFound, "not_found_error", "Live call not found")
return
}
downstream, err := coderws.Accept(c.Writer, c.Request, &coderws.AcceptOptions{
InsecureSkipVerify: true,
})
if err != nil {
return
}
defer func() { _ = downstream.CloseNow() }()
if err := h.gatewayService.ProxyLiveSideband(c.Request.Context(), record, downstream); err != nil {
_ = downstream.Close(coderws.StatusInternalError, "live sideband closed")
return
}
_ = downstream.Close(coderws.StatusNormalClosure, "")
}
func liveEnabledForAPIKey(apiKey *service.APIKey) bool {
return apiKey != nil &&
apiKey.Group != nil &&
apiKey.Group.Platform == service.PlatformOpenAI &&
apiKey.Group.AllowLive
}
@@ -0,0 +1,118 @@
package handler
import (
"bytes"
"encoding/json"
"mime/multipart"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestParseLiveCallRequestMultipartPreservesSession(t *testing.T) {
gin.SetMode(gin.TestMode)
session := `{"model":"gpt-live-test","delegation":{"type":"client"},"instructions":"你好"}`
var body bytes.Buffer
writer := multipart.NewWriter(&body)
require.NoError(t, writer.WriteField("sdp", "v=0\r\n"))
require.NoError(t, writer.WriteField("session", session))
require.NoError(t, writer.Close())
request := httptest.NewRequest("POST", "/v1/live", &body)
request.Header.Set("Content-Type", writer.FormDataContentType())
context, _ := gin.CreateTestContext(httptest.NewRecorder())
context.Request = request
parsed, err := parseLiveCallRequest(context)
require.NoError(t, err)
require.Equal(t, "v=0\r\n", parsed.SDP)
require.JSONEq(t, session, string(parsed.Session))
require.Equal(t, "client", jsonPathString(t, parsed.Session, "delegation", "type"))
}
func TestParseLiveCallRequestJSONPreservesSessionWithoutDelegation(t *testing.T) {
gin.SetMode(gin.TestMode)
body := `{"sdp":"v=0\\r\\n","session":{"model":"gpt-live-test","instructions":"standalone"}}`
request := httptest.NewRequest("POST", "/backend-api/codex/realtime/calls", bytes.NewBufferString(body))
request.Header.Set("Content-Type", "application/json")
context, _ := gin.CreateTestContext(httptest.NewRecorder())
context.Request = request
parsed, err := parseLiveCallRequest(context)
require.NoError(t, err)
require.NotContains(t, string(parsed.Session), "delegation")
require.Equal(t, "standalone", jsonPathString(t, parsed.Session, "instructions"))
}
func TestParseLiveCallRequestRejectsInvalidJSONShape(t *testing.T) {
gin.SetMode(gin.TestMode)
testCases := []string{
`{"session":{"type":"quicksilver"}}`,
`{"sdp":"v=0\\r\\n","session":[]}`,
`{"sdp":"v=0\\r\\n","session":null}`,
`{"sdp":"v=0\\r\\n","session":{"type":"quicksilver"}} {}`,
}
for _, body := range testCases {
request := httptest.NewRequest("POST", "/backend-api/codex/realtime/calls", bytes.NewBufferString(body))
request.Header.Set("Content-Type", "application/json")
context, _ := gin.CreateTestContext(httptest.NewRecorder())
context.Request = request
_, err := parseLiveCallRequest(context)
require.Error(t, err)
}
}
func TestLiveSidebandLocationMatchesCreateRoute(t *testing.T) {
require.Equal(t, "/v1/live/call_123", liveSidebandLocation("/v1/live", "call_123"))
require.Equal(
t,
"/backend-api/codex/call_123",
liveSidebandLocation("/backend-api/codex/realtime/calls", "call_123"),
)
}
func TestLiveEnabledForAPIKey(t *testing.T) {
require.False(t, liveEnabledForAPIKey(nil))
require.False(t, liveEnabledForAPIKey(&service.APIKey{}))
require.False(t, liveEnabledForAPIKey(&service.APIKey{
Group: &service.Group{Platform: service.PlatformOpenAI},
}))
require.False(t, liveEnabledForAPIKey(&service.APIKey{
Group: &service.Group{Platform: service.PlatformAnthropic, AllowLive: true},
}))
require.True(t, liveEnabledForAPIKey(&service.APIKey{
Group: &service.Group{Platform: service.PlatformOpenAI, AllowLive: true},
}))
}
func TestLiveAttestationErrorIsExplicit(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
context, _ := gin.CreateTestContext(recorder)
(&OpenAIGatewayHandler{}).writeLiveCreateError(context, &service.LiveAttestationUnavailableError{
Reason: "Live attestation is only supported when Sub2API runs on macOS",
})
require.Equal(t, http.StatusServiceUnavailable, recorder.Code)
require.Contains(t, recorder.Body.String(), "Sub2API runs on macOS")
}
func jsonPathString(t *testing.T, raw json.RawMessage, keys ...string) string {
t.Helper()
var value any
require.NoError(t, json.Unmarshal(raw, &value))
current := value
for _, key := range keys {
object, ok := current.(map[string]any)
require.True(t, ok)
current = object[key]
}
result, ok := current.(string)
require.True(t, ok)
return result
}
@@ -26,6 +26,9 @@ type userHandlerRepoStub struct {
}
func (s *userHandlerRepoStub) Create(context.Context, *service.User) error { return nil }
func (s *userHandlerRepoStub) CreateWithEmailAliasGuard(context.Context, *service.User) error {
return nil
}
func (s *userHandlerRepoStub) GetByID(context.Context, int64) (*service.User, error) {
cloned := *s.user
return &cloned, nil
@@ -97,6 +100,9 @@ func (s *userHandlerRepoStub) BatchUpdateLimits(context.Context, []int64, *int,
return 0, nil
}
func (s *userHandlerRepoStub) ExistsByEmail(context.Context, string) (bool, error) { return false, nil }
func (s *userHandlerRepoStub) ExistsByEmailAlias(context.Context, string) (bool, error) {
return false, nil
}
func (s *userHandlerRepoStub) RemoveGroupFromAllowedGroups(context.Context, int64) (int64, error) {
return 0, nil
}
+6
View File
@@ -147,6 +147,12 @@ var DefaultModels = []Model{
DisplayName: "Claude Opus 4.8",
CreatedAt: "2026-05-29T00:00:00Z",
},
{
ID: "claude-opus-5",
Type: "model",
DisplayName: "Claude Opus 5",
CreatedAt: "2026-07-25T00:00:00Z",
},
{
ID: "claude-sonnet-5",
Type: "model",
@@ -0,0 +1,17 @@
package liveattestation
import (
"context"
"errors"
)
var (
ErrUnsupportedPlatform = errors.New("live attestation is only supported when Sub2API runs on macOS; Windows support is not implemented yet")
ErrChatGPTAppMissing = errors.New("live attestation requires the official ChatGPT app on the Sub2API server")
)
// Provider 在发起 Live 请求前生成 ChatGPT DeviceCheck attestation。
type Provider interface {
Check(ctx context.Context) error
Generate(ctx context.Context) (string, error)
}
@@ -0,0 +1,294 @@
//go:build darwin
package liveattestation
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"os"
"os/exec"
"path/filepath"
"runtime"
"strings"
"time"
"github.com/google/uuid"
)
const (
chatGPTApplicationPath = "/Applications/ChatGPT.app"
attestationTimeout = 5 * time.Second
)
type darwinProvider struct {
appSessionID string
appPaths []string
}
type deviceSignals struct {
SchemaVersion int `json:"schemaVersion"`
PreferredLanguages []string `json:"preferredLanguages"`
Locale string `json:"locale"`
Timezone string `json:"timezone"`
ScreenSizeSum int `json:"screenSizeSum"`
ScreenScale float64 `json:"screenScale"`
AppSessionID string `json:"appSessionId"`
}
type macOSSignals struct {
Locale string `json:"locale"`
Languages []string `json:"languages"`
Timezone string `json:"timezone"`
Width float64 `json:"width"`
Height float64 `json:"height"`
Scale float64 `json:"scale"`
}
func NewProvider() Provider {
paths := []string{chatGPTApplicationPath}
if home, err := os.UserHomeDir(); err == nil && strings.TrimSpace(home) != "" {
paths = append(paths, filepath.Join(home, "Applications", "ChatGPT.app"))
}
return &darwinProvider{
appSessionID: uuid.NewString(),
appPaths: paths,
}
}
func (p *darwinProvider) Check(ctx context.Context) error {
checkCtx, cancel := context.WithTimeout(ctx, attestationTimeout)
defer cancel()
_, _, _, err := p.resolveRuntime(checkCtx)
return err
}
func (p *darwinProvider) resolveRuntime(ctx context.Context) (string, string, string, error) {
if runtime.GOARCH != "arm64" {
return "", "", "", errors.New("live attestation currently requires Apple Silicon; Intel macOS is not supported")
}
appPath, err := p.findApplication()
if err != nil {
return "", "", "", err
}
resourcesPath := filepath.Join(appPath, "Contents", "Resources")
nodePath := filepath.Join(resourcesPath, "cua_node", "bin", "node")
modulePath := filepath.Join(resourcesPath, "native", "devicecheck.node")
for filePath, label := range map[string]string{
nodePath: "bundled Node.js runtime",
modulePath: "DeviceCheck native module",
} {
if info, statErr := os.Stat(filePath); statErr != nil || info.IsDir() {
return "", "", "", fmt.Errorf("%w: ChatGPT app is missing its %s", ErrChatGPTAppMissing, label)
}
}
bundleID, err := readBundleIdentifier(ctx, appPath)
if err != nil {
return "", "", "", err
}
return nodePath, modulePath, bundleID, nil
}
func (p *darwinProvider) Generate(ctx context.Context) (string, error) {
runCtx, cancel := context.WithTimeout(ctx, attestationTimeout)
defer cancel()
nodePath, modulePath, bundleID, err := p.resolveRuntime(runCtx)
if err != nil {
return "", err
}
signals, err := p.readSignals(runCtx)
if err != nil {
return "", err
}
signalsJSON, err := json.Marshal(signals)
if err != nil {
return "", fmt.Errorf("encode Live attestation signals: %w", err)
}
command := exec.CommandContext(runCtx, nodePath, "-e", deviceCheckScript)
command.Env = []string{
"PATH=/usr/bin:/bin",
"SUB2API_DEVICECHECK_MODULE=" + modulePath,
"SUB2API_ATTESTATION_BUNDLE_ID=" + bundleID,
"SUB2API_ATTESTATION_SIGNALS=" + string(signalsJSON),
}
var stdout bytes.Buffer
var stderr bytes.Buffer
command.Stdout = &stdout
command.Stderr = &stderr
if err := command.Run(); err != nil {
if errors.Is(runCtx.Err(), context.DeadlineExceeded) {
return "", errors.New("ChatGPT DeviceCheck token generation timed out")
}
reason := strings.TrimSpace(stderr.String())
if len(reason) > 240 {
reason = reason[:240]
}
if reason == "" {
reason = err.Error()
}
return "", fmt.Errorf("ChatGPT DeviceCheck token generation failed: %s", reason)
}
header := strings.TrimSpace(stdout.String())
if len(header) < 20 || len(header) > 16*1024 || !json.Valid([]byte(header)) {
return "", errors.New("ChatGPT DeviceCheck returned a malformed attestation")
}
return header, nil
}
func (p *darwinProvider) findApplication() (string, error) {
for _, appPath := range p.appPaths {
info, err := os.Stat(appPath)
if err == nil && info.IsDir() {
return appPath, nil
}
}
return "", ErrChatGPTAppMissing
}
func readBundleIdentifier(ctx context.Context, appPath string) (string, error) {
infoPlist := filepath.Join(appPath, "Contents", "Info.plist")
output, err := exec.CommandContext(
ctx,
"/usr/bin/plutil",
"-extract",
"CFBundleIdentifier",
"raw",
infoPlist,
).Output()
if err != nil {
return "", fmt.Errorf("%w: cannot read its bundle identifier", ErrChatGPTAppMissing)
}
bundleID := strings.TrimSpace(string(output))
if !strings.HasPrefix(bundleID, "com.openai.") {
return "", errors.New("the installed ChatGPT app has an unexpected bundle identifier")
}
return bundleID, nil
}
func (p *darwinProvider) readSignals(ctx context.Context) (deviceSignals, error) {
const script = `ObjC.import("Foundation"); ObjC.import("AppKit");
const screen = $.NSScreen.mainScreen;
const frame = screen.frame;
JSON.stringify({
locale: ObjC.unwrap($.NSLocale.currentLocale.localeIdentifier),
languages: ObjC.deepUnwrap($.NSLocale.preferredLanguages),
timezone: ObjC.unwrap($.NSTimeZone.localTimeZone.name),
width: Number(frame.size.width),
height: Number(frame.size.height),
scale: Number(screen.backingScaleFactor)
})`
output, err := exec.CommandContext(ctx, "/usr/bin/osascript", "-l", "JavaScript", "-e", script).Output()
if err != nil {
return deviceSignals{}, fmt.Errorf("read macOS signals for Live attestation: %w", err)
}
var values macOSSignals
if err := json.Unmarshal(output, &values); err != nil {
return deviceSignals{}, fmt.Errorf("decode macOS signals for Live attestation: %w", err)
}
locale := truncateSignal(values.Locale, 64, "unknown")
languages := values.Languages
if len(languages) == 0 {
languages = []string{locale}
}
if len(languages) > 16 {
languages = languages[:16]
}
for index := range languages {
languages[index] = truncateSignal(languages[index], 64, locale)
}
scale := values.Scale
if scale <= 0 {
scale = 1
}
return deviceSignals{
SchemaVersion: 1,
PreferredLanguages: languages,
Locale: locale,
Timezone: truncateSignal(values.Timezone, 64, "unknown"),
ScreenSizeSum: max(0, int(values.Width+values.Height+0.5)),
ScreenScale: scale,
AppSessionID: truncateSignal(p.appSessionID, 128, uuid.NewString()),
}, nil
}
func truncateSignal(value string, limit int, fallback string) string {
value = strings.TrimSpace(value)
if value == "" {
value = fallback
}
if len(value) > limit {
return value[:limit]
}
return value
}
const deviceCheckScript = `
const addon = require(process.env.SUB2API_DEVICECHECK_MODULE);
const signals = JSON.parse(process.env.SUB2API_ATTESTATION_SIGNALS);
const bundleID = process.env.SUB2API_ATTESTATION_BUNDLE_ID;
function head(major, value) {
if (value < 24) return Buffer.from([major + value]);
if (value <= 255) return Buffer.from([major + 24, value]);
if (value <= 65535) {
const out = Buffer.allocUnsafe(3);
out[0] = major + 25;
out.writeUInt16BE(value, 1);
return out;
}
const out = Buffer.allocUnsafe(5);
out[0] = major + 26;
out.writeUInt32BE(value, 1);
return out;
}
function uint(value) { return head(0, value); }
function text(value) {
const body = Buffer.from(value, "utf8");
return Buffer.concat([head(96, body.length), body]);
}
function float(value) {
if (Number.isSafeInteger(value) && value >= 0) return uint(value);
const out = Buffer.allocUnsafe(9);
out[0] = 251;
out.writeDoubleBE(value, 1);
return out;
}
function array(values) { return Buffer.concat([head(128, values.length), ...values]); }
function map(entries) {
return Buffer.concat([head(160, entries.length), ...entries.flatMap(([key, value]) => [uint(key), value])]);
}
function field(key, value) { return Buffer.concat([text(key), text(value)]); }
function base64url(value) {
return value.toString("base64").replaceAll("+", "-").replaceAll("/", "_").replace(/=+$/u, "");
}
(async () => {
const result = await addon.generateToken();
if (!result || !result.supported) throw new Error("DeviceCheck is not supported on this Mac");
if (!result.tokenBase64) throw new Error("DeviceCheck returned no token");
const fingerprint = map([
[0, uint(signals.schemaVersion)],
[1, array(signals.preferredLanguages.map(text))],
[2, text(signals.locale)],
[3, text(signals.timezone)],
[4, uint(signals.screenSizeSum)],
[5, float(signals.screenScale)],
[6, text(signals.appSessionId)]
]);
const fields = [
field("token", result.tokenBase64),
field("bundle_id", bundleID),
Buffer.concat([text("f"), head(64, fingerprint.length), fingerprint])
];
if (result.latencyMs != null) {
fields.push(Buffer.concat([text("t"), float(result.latencyMs)]));
}
const token = "v1." + base64url(Buffer.concat([Buffer.from([160 + fields.length]), ...fields]));
process.stdout.write(JSON.stringify({v: 1, s: 0, t: token}));
})().catch((error) => {
process.stderr.write(error instanceof Error ? error.message : String(error));
process.exitCode = 1;
});`
@@ -0,0 +1,19 @@
//go:build !darwin
package liveattestation
import "context"
type unsupportedProvider struct{}
func NewProvider() Provider {
return unsupportedProvider{}
}
func (unsupportedProvider) Check(context.Context) error {
return ErrUnsupportedPlatform
}
func (unsupportedProvider) Generate(context.Context) (string, error) {
return "", ErrUnsupportedPlatform
}
@@ -0,0 +1,20 @@
//go:build !darwin
package liveattestation
import (
"context"
"errors"
"testing"
)
func TestUnsupportedProviderReturnsExplicitPlatformError(t *testing.T) {
provider := NewProvider()
if err := provider.Check(context.Background()); !errors.Is(err, ErrUnsupportedPlatform) {
t.Fatalf("Check() error = %v, want ErrUnsupportedPlatform", err)
}
_, err := provider.Generate(context.Background())
if !errors.Is(err, ErrUnsupportedPlatform) {
t.Fatalf("Generate() error = %v, want ErrUnsupportedPlatform", err)
}
}
@@ -354,74 +354,182 @@ func canonicalJSON(raw string) string {
return string(encoded)
}
// ListDueOllamaCloudUsageAccounts returns at most one due representative per
// exact API key before hydration, preventing one shared group from consuming a
// whole runner cycle.
func (r *accountRepository) ListDueOllamaCloudUsageAccounts(ctx context.Context, now time.Time, limit int) ([]service.Account, error) {
// ollamaCloudUsageParseRFC3339SQL reuses the verified RFC3339(/Nano) parse path
// for a snapshot timestamp expression. Invalid or missing values fail open to NULL.
//
// The value is rewritten twice before it reaches jsonpath:
// 1. Sub-second precision beyond 6 digits is truncated, because .datetime()
// rejects more than microsecond resolution while Go emits 9 digits.
// 2. A trailing "Z" is rewritten to "+00:00". jsonpath .datetime() only learned
// to accept the ISO-8601 "Z" designator in PostgreSQL 17, and every timestamp
// this service writes is UTC (hence "Z"). Without this rewrite the parse
// silently yields NULL on PostgreSQL <= 16, which makes every due column NULL
// and collapses ListDueOllamaCloudUsageAccounts into its fail-open branch.
//
// jsonpath (rather than a direct ::timestamptz cast) is required so that values
// passing the shape regex but naming an impossible date (e.g. 2026-02-30) fail
// open to NULL instead of aborting the whole query.
func ollamaCloudUsageParseRFC3339SQL(expression string) string {
return `CASE
WHEN ` + expression + ` IS NULL THEN NULL
WHEN ` + expression + ` ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}(\.[0-9]+)?(Z|[+-][0-9]{2}:[0-9]{2})$'
THEN jsonb_path_query_first_tz(
to_jsonb(regexp_replace(
regexp_replace(
` + expression + `,
'(\.[0-9]{6})[0-9]+(Z|[+-][0-9]{2}:[0-9]{2})$',
'\1\2'
),
'Z$',
'+00:00'
)),
'$.datetime()', '{}'::jsonb, true
) #>> '{}'
ELSE NULL
END`
}
// ListDueOllamaCloudUsageAccounts returns at most one truly-due activity-driven
// candidate per exact API key. Due timing (debounce, max-wait, failure backoff)
// is evaluated in SQL before LIMIT so non-due active groups cannot starve due ones.
// Account.LastUsedAt is stamped with the group MAX(last_used_at) for a service
// pure-function recheck against races between list and refresh.
//
// Rules mirror service.ollamaCloudUsageAutoRefreshDueAt (keep both in sync):
// - missing/invalid snapshot or times → fail-open first due
// - success: activity after fetched_at;
// due_at = GREATEST(LEAST(last_used+debounce, fetched+maxWait), fetched+minFetchInterval)
// - failed/unauthorized: activity after last_attempt; activity_due = LEAST(...);
// final due_at is not earlier than a valid next_refresh_at (invalid/missing fail-open)
func (r *accountRepository) ListDueOllamaCloudUsageAccounts(
ctx context.Context,
now time.Time,
debounce, maxWait time.Duration,
limit int,
) ([]service.Account, error) {
if limit <= 0 {
return []service.Account{}, nil
}
if r == nil || r.sql == nil {
return nil, errors.New("account repository SQL executor not configured")
}
if debounce <= 0 {
debounce = time.Minute
}
if maxWait <= 0 {
maxWait = time.Hour
}
debounceSeconds := debounce.Seconds()
maxWaitSeconds := maxWait.Seconds()
minFetchIntervalSeconds := service.OllamaCloudUsageMinFetchInterval.Seconds()
rows, err := r.sql.QueryContext(ctx, `
WITH candidates AS (
SELECT id, credentials ->> 'api_key' AS api_key,
extra #>> '{ollama_cloud_usage_snapshot,next_refresh_at}' AS next_refresh_at
WITH eligible AS (
SELECT id,
credentials ->> 'api_key' AS api_key,
last_used_at,
extra -> 'ollama_cloud_usage_snapshot' AS snapshot
FROM accounts
WHERE deleted_at IS NULL
AND status = 'active'
AND `+ollamaCloudUsageEligibleSQL+`
AND jsonb_typeof(extra -> 'ollama_cloud_usage_session') = 'string'
AND extra @> '{"ollama_cloud_usage_auto_refresh": true}'::jsonb
), group_activity AS (
SELECT credentials ->> 'api_key' AS api_key,
MAX(last_used_at) AS group_last_used_at
FROM accounts
WHERE deleted_at IS NULL
AND `+ollamaCloudUsageEligibleSQL+`
AND jsonb_typeof(credentials -> 'api_key') = 'string'
GROUP BY credentials ->> 'api_key'
), joined AS (
SELECT e.id, e.api_key, e.snapshot, g.group_last_used_at,
e.snapshot #>> '{status}' AS status,
e.snapshot #>> '{fetched_at}' AS fetched_at,
e.snapshot #>> '{last_attempt_at}' AS last_attempt_at,
e.snapshot #>> '{next_refresh_at}' AS next_refresh_at
FROM eligible e
JOIN group_activity g ON g.api_key = e.api_key
), parsed AS MATERIALIZED (
SELECT id, api_key, next_refresh_at,
next_refresh_at ~ '^[0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}:[0-9]{2}:[0-9]{2}(\.[0-9]+)?(Z|[+-][0-9]{2}:[0-9]{2})$' AS rfc3339_shape,
jsonb_path_query_first_tz(
to_jsonb(regexp_replace(
next_refresh_at,
'(\.[0-9]{6})[0-9]+(Z|[+-][0-9]{2}:[0-9]{2})$',
'\1\2'
)),
'$.datetime()', '{}'::jsonb, true
) #>> '{}' AS parsed_next_refresh_at
FROM candidates
), due AS (
SELECT id, api_key, snapshot, group_last_used_at, status,
`+ollamaCloudUsageParseRFC3339SQL("fetched_at")+` AS parsed_fetched_at,
`+ollamaCloudUsageParseRFC3339SQL("last_attempt_at")+` AS parsed_last_attempt_at,
`+ollamaCloudUsageParseRFC3339SQL("next_refresh_at")+` AS parsed_next_refresh_at
FROM joined
), timed AS (
SELECT *,
CASE WHEN next_refresh_at IS NULL OR NOT rfc3339_shape OR parsed_next_refresh_at IS NULL THEN 0 ELSE 1 END AS due_class
CASE
WHEN status = 'ok'
AND parsed_fetched_at IS NOT NULL
AND group_last_used_at IS NOT NULL
AND group_last_used_at > parsed_fetched_at::timestamptz
THEN GREATEST(
LEAST(
group_last_used_at + make_interval(secs => $2::double precision),
parsed_fetched_at::timestamptz + make_interval(secs => $3::double precision)
),
parsed_fetched_at::timestamptz + make_interval(secs => $5::double precision)
)
WHEN status IN ('failed', 'unauthorized')
AND parsed_last_attempt_at IS NOT NULL
AND group_last_used_at IS NOT NULL
AND group_last_used_at > parsed_last_attempt_at::timestamptz
THEN GREATEST(
LEAST(
group_last_used_at + make_interval(secs => $2::double precision),
parsed_last_attempt_at::timestamptz + make_interval(secs => $3::double precision)
),
COALESCE(parsed_next_refresh_at::timestamptz, '-infinity'::timestamptz)
)
ELSE NULL
END AS activity_due_at
FROM parsed
WHERE next_refresh_at IS NULL
OR NOT rfc3339_shape
OR parsed_next_refresh_at IS NULL
OR parsed_next_refresh_at::timestamptz <= $1
), candidates AS (
SELECT *,
CASE
WHEN snapshot IS NULL OR snapshot = 'null'::jsonb OR status IS NULL
OR status NOT IN ('ok', 'failed', 'unauthorized') THEN 0
WHEN status = 'ok' AND parsed_fetched_at IS NULL THEN 0
WHEN status IN ('failed', 'unauthorized') AND parsed_last_attempt_at IS NULL THEN 0
WHEN activity_due_at IS NOT NULL AND $1 >= activity_due_at THEN 1
ELSE NULL
END AS due_class,
activity_due_at AS due_at
FROM timed
), ranked AS (
SELECT *, row_number() OVER (
PARTITION BY api_key
ORDER BY due_class,
CASE WHEN rfc3339_shape AND parsed_next_refresh_at IS NOT NULL THEN parsed_next_refresh_at::timestamptz END NULLS FIRST,
id
) AS group_rank
FROM due
SELECT id, api_key, group_last_used_at, due_class, due_at,
row_number() OVER (
PARTITION BY api_key
ORDER BY due_class,
due_at NULLS FIRST,
id
) AS group_rank
FROM candidates
WHERE due_class IS NOT NULL
)
SELECT id
SELECT id, group_last_used_at
FROM ranked
WHERE group_rank = 1
ORDER BY due_class,
CASE WHEN rfc3339_shape AND parsed_next_refresh_at IS NOT NULL THEN parsed_next_refresh_at::timestamptz END NULLS FIRST,
id
LIMIT $2
`, now.UTC(), limit)
ORDER BY due_class, due_at NULLS FIRST, id
LIMIT $4
`, now.UTC(), debounceSeconds, maxWaitSeconds, limit, minFetchIntervalSeconds)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
type dueRow struct {
id int64
groupLastUsed *time.Time
}
rowsOut := make([]dueRow, 0, limit)
ids := make([]int64, 0, limit)
for rows.Next() {
var id int64
if err := rows.Scan(&id); err != nil {
var row dueRow
if err := rows.Scan(&row.id, &row.groupLastUsed); err != nil {
return nil, err
}
ids = append(ids, id)
rowsOut = append(rowsOut, row)
ids = append(ids, row.id)
}
if err := rows.Err(); err != nil {
return nil, err
@@ -430,11 +538,26 @@ func (r *accountRepository) ListDueOllamaCloudUsageAccounts(ctx context.Context,
if err != nil {
return nil, err
}
result := make([]service.Account, 0, len(accounts))
byID := make(map[int64]*service.Account, len(accounts))
for _, account := range accounts {
if account != nil {
result = append(result, *account)
byID[account.ID] = account
}
}
result := make([]service.Account, 0, len(rowsOut))
for _, row := range rowsOut {
account := byID[row.id]
if account == nil {
continue
}
// Stamp group MAX(last_used_at) for service due evaluation.
if row.groupLastUsed != nil {
ts := row.groupLastUsed.UTC()
account.LastUsedAt = &ts
} else {
account.LastUsedAt = nil
}
result = append(result, *account)
}
return result, nil
}
@@ -23,81 +23,199 @@ func TestListDueOllamaCloudUsageAccountsOrderingLimitAndProxyHydration(t *testin
Username: "user", Password: "pass", Status: service.StatusActive,
})
createAccount := func(name, baseURL string, proxyID *int64, nextRefreshAt *time.Time) *service.Account {
createAccount := func(name, baseURL string, proxyID *int64, snapshot map[string]any, lastUsed *time.Time) *service.Account {
t.Helper()
extra := map[string]any{
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
service.OllamaCloudUsageAutoRefreshExtraKey: true,
}
if nextRefreshAt != nil {
extra[service.OllamaCloudUsageSnapshotExtraKey] = map[string]any{
"status": service.OllamaCloudUsageStatusOK, "next_refresh_at": nextRefreshAt.UTC().Format(time.RFC3339Nano),
}
if snapshot != nil {
extra[service.OllamaCloudUsageSnapshotExtraKey] = snapshot
}
return mustCreateAccount(t, tx.Client(), &service.Account{
Name: name, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
Credentials: map[string]any{"api_key": name, "base_url": baseURL},
Extra: extra, ProxyID: proxyID,
Extra: extra, ProxyID: proxyID, LastUsedAt: lastUsed,
})
}
uppercasePath := createAccount("ollama-uppercase-path", "https://ollama.com/V1", nil, nil)
missingSnapshot := createAccount("ollama-due-missing", "HTTPS://WWW.OLLAMA.COM:443/v1", &proxy.ID, nil)
oldest := now.Add(-2 * time.Hour)
due := createAccount("ollama-due-oldest", "https://ollama.com", nil, &oldest)
future := now.Add(time.Minute)
_ = createAccount("ollama-not-due", "https://ollama.com", nil, &future)
_ = createAccount("ollama-ineligible", "https://ollama.com.evil.test", nil, nil)
uppercasePath := createAccount("ollama-uppercase-path", "https://ollama.com/V1", nil, nil, nil)
missingSnapshot := createAccount("ollama-due-missing", "HTTPS://WWW.OLLAMA.COM:443/v1", &proxy.ID, nil, nil)
fetched := now.Add(-2 * time.Hour)
activity := now.Add(-5 * time.Minute)
due := createAccount("ollama-due-activity", "https://ollama.com", nil, map[string]any{
"status": service.OllamaCloudUsageStatusOK,
"fetched_at": fetched.UTC().Format(time.RFC3339Nano),
"last_attempt_at": fetched.UTC().Format(time.RFC3339Nano),
"next_refresh_at": fetched.Add(time.Hour).UTC().Format(time.RFC3339Nano),
}, &activity)
// Success snapshot without newer activity must not be listed.
_ = createAccount("ollama-not-due-idle", "https://ollama.com", nil, map[string]any{
"status": service.OllamaCloudUsageStatusOK,
"fetched_at": now.Add(-time.Hour).UTC().Format(time.RFC3339Nano),
"last_attempt_at": now.Add(-time.Hour).UTC().Format(time.RFC3339Nano),
"next_refresh_at": now.Add(-time.Minute).UTC().Format(time.RFC3339Nano),
}, nil)
_ = createAccount("ollama-ineligible", "https://ollama.com.evil.test", nil, nil, nil)
accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, 2)
accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, time.Minute, time.Hour, 2)
require.NoError(t, err)
require.Len(t, accounts, 2)
require.Equal(t, missingSnapshot.ID, accounts[0].ID)
require.Equal(t, due.ID, accounts[1].ID)
require.NotNil(t, accounts[1].LastUsedAt)
require.WithinDuration(t, activity.UTC(), accounts[1].LastUsedAt.UTC(), time.Second)
require.NotContains(t, accountIDs(accounts), uppercasePath.ID)
require.NotNil(t, accounts[0].Proxy)
require.Equal(t, proxy.ID, accounts[0].Proxy.ID)
require.Equal(t, proxy.URL(), accounts[0].Proxy.URL())
}
func TestListDueOllamaCloudUsageAccountsParsesRFC3339NanoAndFailsOpen(t *testing.T) {
// TestListDueOllamaCloudUsageAccountsParsesAllRFC3339Precisions pins the SQL
// timestamp parse path across the sub-second precisions and zone spellings that
// actually reach the database.
//
// Each fixture stores a fetched_at only two minutes old with activity 30s later,
// so a correctly parsed row is NOT due (debounce and the min fetch interval both
// place it in the future). A row whose timestamp fails to parse becomes NULL and
// falls into the fail-open branch, which makes it due. Asserting on absence is
// therefore what makes this test able to fail:
//
// - Go writes UTC times, i.e. the "Z" designator. jsonpath .datetime() only
// accepts "Z" from PostgreSQL 17 on, so without the Z -> +00:00 rewrite in
// ollamaCloudUsageParseRFC3339SQL every fixture here goes due on 14-16.
// - 7/8/9 sub-second digits exceed the microsecond resolution .datetime()
// allows and must be truncated first.
//
// Run against the oldest supported server to exercise the version-sensitive path:
//
// SUB2API_TEST_POSTGRES_IMAGE=postgres:15-alpine go test -tags integration ./internal/repository/
func TestListDueOllamaCloudUsageAccountsParsesAllRFC3339Precisions(t *testing.T) {
ctx := context.Background()
tx := testEntTx(t)
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
now := time.Date(2026, time.July, 22, 14, 0, 0, 0, time.UTC)
activity := now.Add(-30 * time.Second)
create := func(name, nextRefreshAt string) *service.Account {
t.Helper()
return mustCreateAccount(t, tx.Client(), &service.Account{
Name: name, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
Credentials: map[string]any{"api_key": name, "base_url": "https://ollama.com"},
// All three spell the same instant, now-2m, with different precision/zone.
notDue := map[string]string{
"nano-z": "2026-07-22T13:58:00.123456789Z",
"eight-positive": "2026-07-22T14:58:00.12345678+01:00",
"seven-negative": "2026-07-22T11:58:00.1234567-02:00",
}
for name, fetchedAt := range notDue {
_ = mustCreateAccount(t, tx.Client(), &service.Account{
Name: "ollama-precision-" + name, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "precision-" + name, "base_url": "https://ollama.com"},
Extra: map[string]any{
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
service.OllamaCloudUsageAutoRefreshExtraKey: true,
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{
"status": service.OllamaCloudUsageStatusOK, "next_refresh_at": nextRefreshAt,
"status": service.OllamaCloudUsageStatusOK,
"fetched_at": fetchedAt,
"last_attempt_at": fetchedAt,
},
},
LastUsedAt: &activity,
})
}
sevenDigitsOffset := create("ollama-nano-seven", "2026-07-22T11:00:00.1234567-02:00")
eightDigitsOffset := create("ollama-nano-eight", "2026-07-22T11:00:00.12345678+01:00")
nineDigitsZ := create("ollama-nano-nine", "2026-07-22T09:00:00.123456789Z")
invalidCalendar := create("ollama-nano-invalid", "2026-02-30T09:00:00.123456789Z")
future := create("ollama-nano-future", "2026-07-22T15:00:00.123456789Z")
// Guards against a vacuous pass: an genuinely due row must still come back.
staleFetched := now.Add(-2 * time.Hour)
due := mustCreateAccount(t, tx.Client(), &service.Account{
Name: "ollama-precision-due", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "precision-due", "base_url": "https://ollama.com"},
Extra: map[string]any{
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
service.OllamaCloudUsageAutoRefreshExtraKey: true,
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{
"status": service.OllamaCloudUsageStatusOK,
"fetched_at": staleFetched.UTC().Format(time.RFC3339Nano),
"last_attempt_at": staleFetched.UTC().Format(time.RFC3339Nano),
},
},
LastUsedAt: &activity,
})
accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, 10)
accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, time.Minute, time.Hour, 10)
require.NoError(t, err)
ids := accountIDs(accounts)
require.Contains(t, ids, due.ID, "a stale snapshot with fresh activity must be due")
require.Len(t, ids, 1,
"only the stale group may be due; extra rows mean a timestamp failed to parse and fell into the fail-open branch")
}
func TestListDueOllamaCloudUsageAccountsUsesGroupMaxLastUsedAndFailsOpen(t *testing.T) {
ctx := context.Background()
tx := testEntTx(t)
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
now := time.Date(2026, time.July, 22, 14, 0, 0, 0, time.UTC)
fetched := now.Add(-30 * time.Minute)
older := now.Add(-10 * time.Minute)
newer := now.Add(-2 * time.Minute)
leader := mustCreateAccount(t, tx.Client(), &service.Account{
Name: "ollama-group-leader", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "shared-key", "base_url": "https://ollama.com"},
Extra: map[string]any{
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
service.OllamaCloudUsageAutoRefreshExtraKey: true,
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{
"status": service.OllamaCloudUsageStatusOK,
"fetched_at": fetched.UTC().Format(time.RFC3339Nano),
"last_attempt_at": fetched.UTC().Format(time.RFC3339Nano),
"next_refresh_at": fetched.Add(time.Hour).UTC().Format(time.RFC3339Nano),
},
},
LastUsedAt: &older,
})
_ = mustCreateAccount(t, tx.Client(), &service.Account{
Name: "ollama-group-sibling", Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "shared-key", "base_url": "https://www.ollama.com/v1"},
LastUsedAt: &newer,
})
invalid := mustCreateAccount(t, tx.Client(), &service.Account{
Name: "ollama-invalid-snapshot", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "invalid-key", "base_url": "https://ollama.com"},
Extra: map[string]any{
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
service.OllamaCloudUsageAutoRefreshExtraKey: true,
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{
"status": service.OllamaCloudUsageStatusOK, "fetched_at": "2026-02-30T09:00:00.123456789Z",
},
},
})
idle := mustCreateAccount(t, tx.Client(), &service.Account{
Name: "ollama-idle-ok", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
Credentials: map[string]any{"api_key": "idle-key", "base_url": "https://ollama.com"},
Extra: map[string]any{
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
service.OllamaCloudUsageAutoRefreshExtraKey: true,
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{
"status": service.OllamaCloudUsageStatusOK,
"fetched_at": fetched.UTC().Format(time.RFC3339Nano),
"last_attempt_at": fetched.UTC().Format(time.RFC3339Nano),
"next_refresh_at": fetched.Add(time.Hour).UTC().Format(time.RFC3339Nano),
},
},
})
accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, time.Minute, time.Hour, 10)
require.NoError(t, err, "invalid stored values must not abort the query")
require.Equal(t, []int64{
invalidCalendar.ID,
nineDigitsZ.ID,
eightDigitsOffset.ID,
sevenDigitsOffset.ID,
}, accountIDs(accounts))
require.NotContains(t, accountIDs(accounts), future.ID)
ids := accountIDs(accounts)
require.Contains(t, ids, invalid.ID)
require.Contains(t, ids, leader.ID)
require.NotContains(t, ids, idle.ID)
for _, account := range accounts {
if account.ID == leader.ID {
require.NotNil(t, account.LastUsedAt)
require.WithinDuration(t, newer.UTC(), account.LastUsedAt.UTC(), time.Second,
"group MAX(last_used_at) must come from the sibling")
}
}
}
func TestLockAndMergeAccountProbeExtraCoalescesNullableOllamaGroupIdentity(t *testing.T) {
@@ -413,3 +531,111 @@ func TestUpdateCredentialsUnchangedCredentialsPreserveManagedExtra(t *testing.T)
require.NotContains(t, probeLoaded.Extra, service.UpstreamBillingProbeExtraKey,
"changed credentials must keep clearing the probe snapshot")
}
// TestListDueOllamaCloudUsageAccountsSQLDueRulesMatchService proves the SQL
// candidate layer applies debounce / max-wait / failure-backoff before LIMIT,
// matching service.ollamaCloudUsageIsAutoRefreshDue, and that >20 active-but-
// not-yet-due groups cannot starve a truly due max-wait group.
func TestListDueOllamaCloudUsageAccountsSQLDueRulesMatchService(t *testing.T) {
ctx := context.Background()
tx := testEntTx(t)
repo := newAccountRepositoryWithSQL(tx.Client(), tx, nil)
now := time.Date(2026, time.July, 25, 12, 0, 0, 0, time.UTC)
debounce := time.Minute
maxWait := time.Hour
createOK := func(name string, fetched, lastUsed time.Time) *service.Account {
t.Helper()
return mustCreateAccount(t, tx.Client(), &service.Account{
Name: name, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
Credentials: map[string]any{"api_key": name, "base_url": "https://ollama.com"},
Extra: map[string]any{
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
service.OllamaCloudUsageAutoRefreshExtraKey: true,
service.OllamaCloudUsageSnapshotExtraKey: map[string]any{
"status": service.OllamaCloudUsageStatusOK,
"fetched_at": fetched.UTC().Format(time.RFC3339Nano),
"last_attempt_at": fetched.UTC().Format(time.RFC3339Nano),
"next_refresh_at": fetched.Add(maxWait).UTC().Format(time.RFC3339Nano),
},
},
LastUsedAt: &lastUsed,
})
}
createFailed := func(name string, lastAttempt, lastUsed, nextRefresh time.Time, nextRefreshRaw string) *service.Account {
t.Helper()
snapshot := map[string]any{
"status": service.OllamaCloudUsageStatusFailed,
"last_attempt_at": lastAttempt.UTC().Format(time.RFC3339Nano),
"failure_count": 1,
}
if nextRefreshRaw != "" {
snapshot["next_refresh_at"] = nextRefreshRaw
} else {
snapshot["next_refresh_at"] = nextRefresh.UTC().Format(time.RFC3339Nano)
}
return mustCreateAccount(t, tx.Client(), &service.Account{
Name: name, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
Credentials: map[string]any{"api_key": name, "base_url": "https://ollama.com"},
Extra: map[string]any{
service.OllamaCloudUsageSessionExtraKey: "cipher:wos-session=fixture",
service.OllamaCloudUsageAutoRefreshExtraKey: true,
service.OllamaCloudUsageSnapshotExtraKey: snapshot,
},
LastUsedAt: &lastUsed,
})
}
// 21 groups with activity after fetch but debounce not elapsed — previously
// these alone could fill LIMIT 20 every minute and starve true due groups.
notDueIDs := make(map[int64]struct{}, 21)
for i := 0; i < 21; i++ {
// fetched 10m ago, last used 10s ago → due_at = lastUsed+debounce = now+50s (not due)
acc := createOK(fmt.Sprintf("ollama-not-due-debounce-%02d", i), now.Add(-10*time.Minute), now.Add(-10*time.Second))
notDueIDs[acc.ID] = struct{}{}
}
// Truly due via max-wait: fetched 2h ago, continuous activity 10s ago.
// due_at = min(now-10s+1m, now-2h+1h) = now-1h → due.
maxWaitDue := createOK("ollama-due-maxwait", now.Add(-2*time.Hour), now.Add(-10*time.Second))
// Success debounce elapsed: last used 2m ago with debounce 1m → due.
debounceDue := createOK("ollama-due-debounce", now.Add(-30*time.Minute), now.Add(-2*time.Minute))
// Success still within debounce → not due.
_ = createOK("ollama-not-due-fresh", now.Add(-30*time.Minute), now.Add(-20*time.Second))
// Failure blocked by next_refresh_at backoff even with new activity.
_ = createFailed("ollama-fail-backoff", now.Add(-30*time.Minute), now.Add(-2*time.Minute), now.Add(10*time.Minute), "")
// Failure after backoff with new request → due.
failDue := createFailed("ollama-fail-due", now.Add(-30*time.Minute), now.Add(-2*time.Minute), now.Add(-time.Minute), "")
// Invalid next_refresh_at must fail open (not abort query / not block activity due).
failInvalidNext := createFailed("ollama-fail-invalid-next", now.Add(-30*time.Minute), now.Add(-2*time.Minute), time.Time{}, "not-a-timestamp")
accounts, err := repo.ListDueOllamaCloudUsageAccounts(ctx, now, debounce, maxWait, 20)
require.NoError(t, err)
ids := accountIDs(accounts)
require.Contains(t, ids, maxWaitDue.ID, "max-wait due group must not be starved by not-yet-due activity groups")
require.Contains(t, ids, debounceDue.ID, "success debounce elapsed must be due in SQL")
require.Contains(t, ids, failDue.ID, "failure after backoff with new activity must be due in SQL")
require.Contains(t, ids, failInvalidNext.ID, "invalid next_refresh_at must fail open to activity due")
require.LessOrEqual(t, len(accounts), 20)
// Fixtures below match service.ollamaCloudUsageIsAutoRefreshDue semantics;
// none of the not-yet-due groups may appear even when they outnumber the limit.
for _, id := range ids {
_, isNotDue := notDueIDs[id]
require.False(t, isNotDue, "not-yet-due debounce group %d must not be returned by SQL LIMIT layer", id)
}
require.NotContains(t, ids, int64(0))
// Explicit not-due names must stay out: fresh success and failure still in backoff.
for _, account := range accounts {
require.NotContains(t, account.Name, "not-due")
require.NotEqual(t, "ollama-fail-backoff", account.Name)
require.NotEqual(t, "ollama-not-due-fresh", account.Name)
}
}
@@ -190,13 +190,15 @@ func TestListDueOllamaCloudUsageAccountsFiltersOrdersAndLimits(t *testing.T) {
require.NoError(t, err)
t.Cleanup(func() { _ = db.Close() })
now := time.Date(2026, time.July, 22, 12, 0, 0, 0, time.UTC)
debounce := time.Minute
maxWait := time.Hour
var capturedSQL string
mock.ExpectQuery("WITH candidates AS").
WithArgs(now, 20).
WillReturnRows(sqlmock.NewRows([]string{"id"}))
mock.ExpectQuery("WITH eligible AS").
WithArgs(now.UTC(), debounce.Seconds(), maxWait.Seconds(), 20, service.OllamaCloudUsageMinFetchInterval.Seconds()).
WillReturnRows(sqlmock.NewRows([]string{"id", "group_last_used_at"}))
repo := newAccountRepositoryWithSQL(nil, captureQuerySQL{db: db, captured: &capturedSQL}, nil)
accounts, err := repo.ListDueOllamaCloudUsageAccounts(context.Background(), now, 20)
accounts, err := repo.ListDueOllamaCloudUsageAccounts(context.Background(), now, debounce, maxWait, 20)
require.NoError(t, err)
require.Empty(t, accounts)
@@ -209,10 +211,24 @@ func TestListDueOllamaCloudUsageAccountsFiltersOrdersAndLimits(t *testing.T) {
ollamaCloudBaseURLMatchesSQL("credentials ->> 'base_url'"),
"jsonb_typeof(extra -> 'ollama_cloud_usage_session') = 'string'",
`extra @> '{"ollama_cloud_usage_auto_refresh": true}'::jsonb`,
"parsed_next_refresh_at::timestamptz <= $1",
"MAX(last_used_at) AS group_last_used_at",
"PARTITION BY api_key",
"WHERE group_rank = 1",
"LIMIT $2",
"LIMIT $4",
"make_interval(secs => $2::double precision)",
"make_interval(secs => $3::double precision)",
// Minimum interval floor between successful fetches.
"make_interval(secs => $5::double precision)",
// jsonpath .datetime() only accepts the ISO-8601 "Z" designator from
// PostgreSQL 17 on, and this service writes UTC timestamps. Without this
// rewrite every parsed_* column is NULL on 14-16 and the due filter
// collapses into its fail-open branch.
`regexp_replace( regexp_replace( fetched_at, '(\.[0-9]{6})[0-9]+(Z|[+-][0-9]{2}:[0-9]{2})$', '\1\2' ), 'Z$', '+00:00' )`,
"group_last_used_at > parsed_fetched_at::timestamptz",
"group_last_used_at > parsed_last_attempt_at::timestamptz",
"$1 >= activity_due_at",
"COALESCE(parsed_next_refresh_at::timestamptz, '-infinity'::timestamptz)",
"ORDER BY due_class, due_at NULLS FIRST, id",
} {
require.Contains(t, normalized, clause)
}
@@ -199,6 +199,7 @@ func (r *apiKeyRepository) GetByKeyForAuth(ctx context.Context, key string) (*se
group.FieldMcpXMLInject,
group.FieldSupportedModelScopes,
group.FieldAllowMessagesDispatch,
group.FieldAllowLive,
group.FieldDefaultMappedModel,
group.FieldMessagesDispatchModelConfig,
group.FieldModelsListConfig,
@@ -948,6 +949,7 @@ func groupEntityToService(g *dbent.Group) *service.Group {
SupportedModelScopes: g.SupportedModelScopes,
SortOrder: g.SortOrder,
AllowMessagesDispatch: g.AllowMessagesDispatch,
AllowLive: g.AllowLive,
RequireOAuthOnly: g.RequireOauthOnly,
RequirePrivacySet: g.RequirePrivacySet,
DefaultMappedModel: g.DefaultMappedModel,
@@ -750,7 +750,7 @@ INSERT INTO batch_image_jobs (
batch_discount_multiplier, hold_multiplier, billable_unit_price, hold_unit_price,
pricing_snapshot_version,
currency, hold_id,
idempotency_key, request_hash, manifest_hash, retry_count, output_expires_at
idempotency_key, request_hash, manifest_hash, retry_count, session_id, output_expires_at
) VALUES (
$1, $2, $3, $4, $5, $6, $7, $8, $9,
$10, $11, $12, $13, $14,
@@ -760,7 +760,7 @@ INSERT INTO batch_image_jobs (
$25, $26, $27, $28,
$29,
$30, $31,
$32, $33, $34, $35, $36
$32, $33, $34, $35, $36, $37
)
RETURNING `+batchImageJobColumns,
params.BatchID, params.UserID, params.APIKeyID, params.AccountID, params.Provider, params.Model, params.TaskName, params.ParentBatchID, params.Status,
@@ -771,7 +771,7 @@ RETURNING `+batchImageJobColumns,
params.BatchDiscountMultiplier, params.HoldMultiplier, params.BillableUnitPrice, params.HoldUnitPrice,
params.PricingSnapshotVersion,
params.Currency, params.HoldID,
params.IdempotencyKey, params.RequestHash, params.ManifestHash, params.RetryCount, params.OutputExpiresAt,
params.IdempotencyKey, params.RequestHash, params.ManifestHash, params.RetryCount, params.SessionID, params.OutputExpiresAt,
))
}
@@ -825,7 +825,7 @@ batch_discount_multiplier, hold_multiplier, billable_unit_price, hold_unit_price
pricing_snapshot_version,
currency, hold_id,
idempotency_key, request_hash, manifest_hash,
retry_count, version, output_expires_at, input_deleted_at, output_deleted_at, downloaded_at, user_deleted_at,
retry_count, version, session_id, output_expires_at, input_deleted_at, output_deleted_at, downloaded_at, user_deleted_at,
last_error_code, last_error_message,
created_at, updated_at, submitted_at, started_at, finished_at, settled_at`
@@ -838,6 +838,7 @@ func scanBatchImageJob(row rowScanner) (*service.BatchImageJob, error) {
var parentBatchID sql.NullString
var holdAmount, actualCost sql.NullFloat64
var holdID, idempotencyKey, requestHash, manifestHash sql.NullString
var sessionID sql.NullString
var outputExpiresAt, inputDeletedAt, outputDeletedAt, downloadedAt, userDeletedAt sql.NullTime
var lastErrorCode, lastErrorMessage sql.NullString
var submittedAt, startedAt, finishedAt, settledAt sql.NullTime
@@ -852,7 +853,7 @@ func scanBatchImageJob(row rowScanner) (*service.BatchImageJob, error) {
&job.PricingSnapshotVersion,
&job.Currency, &holdID,
&idempotencyKey, &requestHash, &manifestHash,
&job.RetryCount, &job.Version, &outputExpiresAt, &inputDeletedAt, &outputDeletedAt, &downloadedAt, &userDeletedAt,
&job.RetryCount, &job.Version, &sessionID, &outputExpiresAt, &inputDeletedAt, &outputDeletedAt, &downloadedAt, &userDeletedAt,
&lastErrorCode, &lastErrorMessage,
&job.CreatedAt, &job.UpdatedAt, &submittedAt, &startedAt, &finishedAt, &settledAt,
)
@@ -874,6 +875,7 @@ func scanBatchImageJob(row rowScanner) (*service.BatchImageJob, error) {
job.IdempotencyKey = batchImageNullStringPtr(idempotencyKey)
job.RequestHash = batchImageNullStringPtr(requestHash)
job.ManifestHash = batchImageNullStringPtr(manifestHash)
job.SessionID = batchImageNullStringPtr(sessionID)
job.OutputExpiresAt = batchImageNullTimePtr(outputExpiresAt)
job.InputDeletedAt = batchImageNullTimePtr(inputDeletedAt)
job.OutputDeletedAt = batchImageNullTimePtr(outputDeletedAt)
+152 -13
View File
@@ -29,11 +29,15 @@ const (
// 格式: concurrency:user:{userID}
userSlotKeyPrefix = "concurrency:user:"
// 格式: concurrency:api_key:{apiKeyID}
apiKeySlotKeyPrefix = "concurrency:api_key:"
apiKeySlotKeyPrefix = "concurrency:api_key:"
liveAccountSlotKeyPrefix = "concurrency:live:account:"
liveUserSlotKeyPrefix = "concurrency:live:user:"
liveAPIKeySlotKeyPrefix = "concurrency:live:api_key:"
// API-key-scoped client WebSocket ingress leases use a shorter TTL than
// ordinary request slots, because idle ingress sessions do not hold a turn slot.
openAIWSIngressLeaseKeyPrefix = "concurrency:openai_ws_ingress:api_key:"
openAIWSIngressLeaseTTLSeconds = 60
liveLeaseTTLSeconds = 60
// 等待队列计数器格式: concurrency:wait:{userID}
waitQueueKeyPrefix = "concurrency:wait:"
// 账号级等待队列计数器格式: wait:account:{accountID}
@@ -59,7 +63,7 @@ const (
var (
// acquireScript 使用有序集合计数并在未达上限时添加槽位
// 使用 Redis TIME 命令获取服务器时间,避免多实例时钟不同步问题
// KEYS[1] = 有序集合键 (concurrency:account:{id} / concurrency:user:{id})
// KEYS[1] = 普通槽位键,KEYS[2] = 对应 Live 槽位键
// ARGV[1] = maxConcurrency
// ARGV[2] = TTL(秒)
// ARGV[3] = requestID
@@ -69,6 +73,7 @@ var (
-- replicates correctly. No-op on Redis 5.0+ (effects replication is default).
redis.replicate_commands()
local key = KEYS[1]
local liveKey = KEYS[2]
local maxConcurrency = tonumber(ARGV[1])
local ttl = tonumber(ARGV[2])
local requestID = ARGV[3]
@@ -80,6 +85,7 @@ var (
-- 清理过期槽位
redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
redis.call('ZREMRANGEBYSCORE', liveKey, '-inf', now - 60)
-- 检查是否已存在(支持重试场景刷新时间戳)
local exists = redis.call('ZSCORE', key, requestID)
@@ -90,7 +96,7 @@ var (
end
-- 检查是否达到并发上限
local count = redis.call('ZCARD', key)
local count = redis.call('ZCARD', key) + redis.call('ZCARD', liveKey)
if count < maxConcurrency then
redis.call('ZADD', key, now, requestID)
redis.call('EXPIRE', key, ttl)
@@ -102,13 +108,14 @@ var (
// getCountScript 统计有序集合中的槽位数量并清理过期条目
// 使用 Redis TIME 命令获取服务器时间
// KEYS[1] = 有序集合键
// KEYS[1] = 普通槽位键,KEYS[2] = 对应 Live 槽位键
// ARGV[1] = TTL(秒)
getCountScript = redis.NewScript(`
-- Redis 3.2-4.x compat: opt into effects replication so redis.call('TIME')
-- replicates correctly. No-op on Redis 5.0+ (effects replication is default).
redis.replicate_commands()
local key = KEYS[1]
local liveKey = KEYS[2]
local ttl = tonumber(ARGV[1])
-- 使用 Redis 服务器时间
@@ -117,7 +124,60 @@ var (
local expireBefore = now - ttl
redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
return redis.call('ZCARD', key)
redis.call('ZREMRANGEBYSCORE', liveKey, '-inf', now - 60)
return redis.call('ZCARD', key) + redis.call('ZCARD', liveKey)
`)
acquireLiveLeaseScript = redis.NewScript(`
redis.replicate_commands()
local accountRegular = KEYS[1]
local accountLive = KEYS[2]
local userRegular = KEYS[3]
local userLive = KEYS[4]
local apiLive = KEYS[5]
local accountMax = tonumber(ARGV[1])
local userMax = tonumber(ARGV[2])
local ttl = tonumber(ARGV[3])
local leaseID = ARGV[4]
local replacing = tonumber(ARGV[5])
local now = tonumber(redis.call('TIME')[1])
local liveExpireBefore = now - ttl
redis.call('ZREMRANGEBYSCORE', accountLive, '-inf', liveExpireBefore)
redis.call('ZREMRANGEBYSCORE', userLive, '-inf', liveExpireBefore)
redis.call('ZREMRANGEBYSCORE', apiLive, '-inf', liveExpireBefore)
if redis.call('ZSCORE', accountLive, leaseID) ~= false then
return 1
end
local accountCount = redis.call('ZCARD', accountRegular) + redis.call('ZCARD', accountLive)
local userCount = redis.call('ZCARD', userRegular) + redis.call('ZCARD', userLive)
local allowance = 0
if replacing == 1 then allowance = 1 end
if accountMax > 0 and accountCount >= accountMax + allowance then return 0 end
if userMax > 0 and userCount >= userMax + allowance then return 0 end
redis.call('ZADD', accountLive, now, leaseID)
redis.call('ZADD', userLive, now, leaseID)
redis.call('ZADD', apiLive, now, leaseID)
redis.call('EXPIRE', accountLive, ttl)
redis.call('EXPIRE', userLive, ttl)
redis.call('EXPIRE', apiLive, ttl)
return 1
`)
refreshLiveLeaseScript = redis.NewScript(`
redis.replicate_commands()
local ttl = tonumber(ARGV[1])
local leaseID = ARGV[2]
local now = tonumber(redis.call('TIME')[1])
local expireBefore = now - ttl
for _, key in ipairs(KEYS) do
redis.call('ZREMRANGEBYSCORE', key, '-inf', expireBefore)
if redis.call('ZSCORE', key, leaseID) == false then return 0 end
end
for _, key in ipairs(KEYS) do
redis.call('ZADD', key, now, leaseID)
redis.call('EXPIRE', key, ttl)
end
return 1
`)
// trackSlotScript 记录 stats-only 槽位,不做并发上限判断。
@@ -330,6 +390,18 @@ func apiKeySlotKey(apiKeyID int64) string {
return fmt.Sprintf("%s%d", apiKeySlotKeyPrefix, apiKeyID)
}
func liveAccountSlotKey(accountID int64) string {
return fmt.Sprintf("%s%d", liveAccountSlotKeyPrefix, accountID)
}
func liveUserSlotKey(userID int64) string {
return fmt.Sprintf("%s%d", liveUserSlotKeyPrefix, userID)
}
func liveAPIKeySlotKey(apiKeyID int64) string {
return fmt.Sprintf("%s%d", liveAPIKeySlotKeyPrefix, apiKeyID)
}
func openAIWSIngressLeaseKey(apiKeyID int64) string {
return fmt.Sprintf("%s%d", openAIWSIngressLeaseKeyPrefix, apiKeyID)
}
@@ -559,7 +631,7 @@ func runScriptInt64Pair(ctx context.Context, rdb *redis.Client, script *redis.Sc
func (c *concurrencyCache) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
key := accountSlotKey(accountID)
// 时间戳在 Lua 脚本内使用 Redis TIME 命令获取,确保多实例时钟一致
result, now, err := runScriptInt64Pair(ctx, c.rdb, acquireScript, []string{key}, maxConcurrency, c.slotTTLSeconds, requestID)
result, now, err := runScriptInt64Pair(ctx, c.rdb, acquireScript, []string{key, liveAccountSlotKey(accountID)}, maxConcurrency, c.slotTTLSeconds, requestID)
if err != nil {
return false, err
}
@@ -583,7 +655,7 @@ func (c *concurrencyCache) ReleaseAccountSlot(ctx context.Context, accountID int
func (c *concurrencyCache) GetAccountConcurrency(ctx context.Context, accountID int64) (int, error) {
key := accountSlotKey(accountID)
// 时间戳在 Lua 脚本内使用 Redis TIME 命令获取
result, err := getCountScript.Run(ctx, c.rdb, []string{key}, c.slotTTLSeconds).Int()
result, err := getCountScript.Run(ctx, c.rdb, []string{key, liveAccountSlotKey(accountID)}, c.slotTTLSeconds).Int()
if err != nil {
return 0, err
}
@@ -605,14 +677,18 @@ func (c *concurrencyCache) GetAccountConcurrencyBatch(ctx context.Context, accou
type accountCmd struct {
accountID int64
zcardCmd *redis.IntCmd
liveCmd *redis.IntCmd
}
cmds := make([]accountCmd, 0, len(accountIDs))
for _, accountID := range accountIDs {
slotKey := accountSlotKeyPrefix + strconv.FormatInt(accountID, 10)
liveKey := liveAccountSlotKeyPrefix + strconv.FormatInt(accountID, 10)
pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
pipe.ZRemRangeByScore(ctx, liveKey, "-inf", strconv.FormatInt(now.Unix()-liveLeaseTTLSeconds, 10))
cmds = append(cmds, accountCmd{
accountID: accountID,
zcardCmd: pipe.ZCard(ctx, slotKey),
liveCmd: pipe.ZCard(ctx, liveKey),
})
}
@@ -622,7 +698,7 @@ func (c *concurrencyCache) GetAccountConcurrencyBatch(ctx context.Context, accou
result := make(map[int64]int, len(accountIDs))
for _, cmd := range cmds {
result[cmd.accountID] = int(cmd.zcardCmd.Val())
result[cmd.accountID] = int(cmd.zcardCmd.Val() + cmd.liveCmd.Val())
}
return result, nil
}
@@ -632,7 +708,7 @@ func (c *concurrencyCache) GetAccountConcurrencyBatch(ctx context.Context, accou
func (c *concurrencyCache) AcquireUserSlot(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error) {
key := userSlotKey(userID)
// 时间戳在 Lua 脚本内使用 Redis TIME 命令获取,确保多实例时钟一致
result, now, err := runScriptInt64Pair(ctx, c.rdb, acquireScript, []string{key}, maxConcurrency, c.slotTTLSeconds, requestID)
result, now, err := runScriptInt64Pair(ctx, c.rdb, acquireScript, []string{key, liveUserSlotKey(userID)}, maxConcurrency, c.slotTTLSeconds, requestID)
if err != nil {
return false, err
}
@@ -656,7 +732,7 @@ func (c *concurrencyCache) ReleaseUserSlot(ctx context.Context, userID int64, re
func (c *concurrencyCache) GetUserConcurrency(ctx context.Context, userID int64) (int, error) {
key := userSlotKey(userID)
// 时间戳在 Lua 脚本内使用 Redis TIME 命令获取
result, err := getCountScript.Run(ctx, c.rdb, []string{key}, c.slotTTLSeconds).Int()
result, err := getCountScript.Run(ctx, c.rdb, []string{key, liveUserSlotKey(userID)}, c.slotTTLSeconds).Int()
if err != nil {
return 0, err
}
@@ -716,6 +792,57 @@ func (c *concurrencyCache) ReleaseOpenAIWSIngressLease(ctx context.Context, apiK
return c.rdb.ZRem(ctx, openAIWSIngressLeaseKey(apiKeyID), leaseID).Err()
}
func (c *concurrencyCache) AcquireLiveLease(
ctx context.Context,
accountID int64,
accountMax int,
userID int64,
userMax int,
apiKeyID int64,
leaseID string,
replacingRegularSlots bool,
) (bool, error) {
if c == nil || c.rdb == nil || accountID <= 0 || userID <= 0 || apiKeyID <= 0 || leaseID == "" {
return false, nil
}
replacing := 0
if replacingRegularSlots {
replacing = 1
}
result, err := acquireLiveLeaseScript.Run(ctx, c.rdb, []string{
accountSlotKey(accountID),
liveAccountSlotKey(accountID),
userSlotKey(userID),
liveUserSlotKey(userID),
liveAPIKeySlotKey(apiKeyID),
}, accountMax, userMax, liveLeaseTTLSeconds, leaseID, replacing).Int()
return result == 1, err
}
func (c *concurrencyCache) RefreshLiveLease(ctx context.Context, accountID, userID, apiKeyID int64, leaseID string) (bool, error) {
if c == nil || c.rdb == nil || leaseID == "" {
return false, nil
}
result, err := refreshLiveLeaseScript.Run(ctx, c.rdb, []string{
liveAccountSlotKey(accountID),
liveUserSlotKey(userID),
liveAPIKeySlotKey(apiKeyID),
}, liveLeaseTTLSeconds, leaseID).Int()
return result == 1, err
}
func (c *concurrencyCache) ReleaseLiveLease(ctx context.Context, accountID, userID, apiKeyID int64, leaseID string) error {
if c == nil || c.rdb == nil || leaseID == "" {
return nil
}
pipe := c.rdb.TxPipeline()
pipe.ZRem(ctx, liveAccountSlotKey(accountID), leaseID)
pipe.ZRem(ctx, liveUserSlotKey(userID), leaseID)
pipe.ZRem(ctx, liveAPIKeySlotKey(apiKeyID), leaseID)
_, err := pipe.Exec(ctx)
return err
}
func (c *concurrencyCache) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKeyIDs []int64) (map[int64]int, error) {
if len(apiKeyIDs) == 0 {
return map[int64]int{}, nil
@@ -731,14 +858,18 @@ func (c *concurrencyCache) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKey
type apiKeyCmd struct {
apiKeyID int64
zcardCmd *redis.IntCmd
liveCmd *redis.IntCmd
}
cmds := make([]apiKeyCmd, 0, len(apiKeyIDs))
for _, apiKeyID := range apiKeyIDs {
slotKey := apiKeySlotKeyPrefix + strconv.FormatInt(apiKeyID, 10)
liveKey := liveAPIKeySlotKeyPrefix + strconv.FormatInt(apiKeyID, 10)
pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
pipe.ZRemRangeByScore(ctx, liveKey, "-inf", strconv.FormatInt(now.Unix()-liveLeaseTTLSeconds, 10))
cmds = append(cmds, apiKeyCmd{
apiKeyID: apiKeyID,
zcardCmd: pipe.ZCard(ctx, slotKey),
liveCmd: pipe.ZCard(ctx, liveKey),
})
}
@@ -748,7 +879,7 @@ func (c *concurrencyCache) GetAPIKeyConcurrencyBatch(ctx context.Context, apiKey
result := make(map[int64]int, len(apiKeyIDs))
for _, cmd := range cmds {
result[cmd.apiKeyID] = int(cmd.zcardCmd.Val())
result[cmd.apiKeyID] = int(cmd.zcardCmd.Val() + cmd.liveCmd.Val())
}
return result, nil
}
@@ -834,17 +965,21 @@ func (c *concurrencyCache) GetAccountsLoadBatch(ctx context.Context, accounts []
id int64
maxConcurrency int
zcardCmd *redis.IntCmd
liveCmd *redis.IntCmd
getCmd *redis.StringCmd
}
cmds := make([]accountCmds, 0, len(accounts))
for _, acc := range accounts {
slotKey := accountSlotKeyPrefix + strconv.FormatInt(acc.ID, 10)
liveKey := liveAccountSlotKeyPrefix + strconv.FormatInt(acc.ID, 10)
waitKey := accountWaitKeyPrefix + strconv.FormatInt(acc.ID, 10)
pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
pipe.ZRemRangeByScore(ctx, liveKey, "-inf", strconv.FormatInt(now.Unix()-liveLeaseTTLSeconds, 10))
ac := accountCmds{
id: acc.ID,
maxConcurrency: acc.MaxConcurrency,
zcardCmd: pipe.ZCard(ctx, slotKey),
liveCmd: pipe.ZCard(ctx, liveKey),
getCmd: pipe.Get(ctx, waitKey),
}
cmds = append(cmds, ac)
@@ -856,7 +991,7 @@ func (c *concurrencyCache) GetAccountsLoadBatch(ctx context.Context, accounts []
loadMap := make(map[int64]*service.AccountLoadInfo, len(accounts))
for _, ac := range cmds {
currentConcurrency := int(ac.zcardCmd.Val())
currentConcurrency := int(ac.zcardCmd.Val() + ac.liveCmd.Val())
waitingCount := 0
if v, err := ac.getCmd.Int(); err == nil {
waitingCount = v
@@ -894,17 +1029,21 @@ func (c *concurrencyCache) GetUsersLoadBatch(ctx context.Context, users []servic
id int64
maxConcurrency int
zcardCmd *redis.IntCmd
liveCmd *redis.IntCmd
getCmd *redis.StringCmd
}
cmds := make([]userCmds, 0, len(users))
for _, u := range users {
slotKey := userSlotKeyPrefix + strconv.FormatInt(u.ID, 10)
liveKey := liveUserSlotKeyPrefix + strconv.FormatInt(u.ID, 10)
waitKey := waitQueueKeyPrefix + strconv.FormatInt(u.ID, 10)
pipe.ZRemRangeByScore(ctx, slotKey, "-inf", strconv.FormatInt(cutoffTime, 10))
pipe.ZRemRangeByScore(ctx, liveKey, "-inf", strconv.FormatInt(now.Unix()-liveLeaseTTLSeconds, 10))
uc := userCmds{
id: u.ID,
maxConcurrency: u.MaxConcurrency,
zcardCmd: pipe.ZCard(ctx, slotKey),
liveCmd: pipe.ZCard(ctx, liveKey),
getCmd: pipe.Get(ctx, waitKey),
}
cmds = append(cmds, uc)
@@ -916,7 +1055,7 @@ func (c *concurrencyCache) GetUsersLoadBatch(ctx context.Context, users []servic
loadMap := make(map[int64]*service.UserLoadInfo, len(users))
for _, uc := range cmds {
currentConcurrency := int(uc.zcardCmd.Val())
currentConcurrency := int(uc.zcardCmd.Val() + uc.liveCmd.Val())
waitingCount := 0
if v, err := uc.getCmd.Int(); err == nil {
waitingCount = v
@@ -97,6 +97,42 @@ func (s *ConcurrencyCacheSuite) TestOpenAIWSIngressAPIKeySlot_ReapsCrashedLeaseW
require.Equal(s.T(), int64(2), count)
}
func (s *ConcurrencyCacheSuite) TestLiveLease_CountsTowardRegularAccountAndUserLimits() {
liveCache, ok := s.cache.(service.LiveConcurrencyCache)
require.True(s.T(), ok)
accountID := int64(9101)
userID := int64(9102)
apiKeyID := int64(9103)
acquired, err := liveCache.AcquireLiveLease(
s.ctx,
accountID,
1,
userID,
1,
apiKeyID,
"live-integration",
false,
)
require.NoError(s.T(), err)
require.True(s.T(), acquired)
regularAccount, err := s.cache.AcquireAccountSlot(s.ctx, accountID, 1, "regular-account")
require.NoError(s.T(), err)
require.False(s.T(), regularAccount)
regularUser, err := s.cache.AcquireUserSlot(s.ctx, userID, 1, "regular-user")
require.NoError(s.T(), err)
require.False(s.T(), regularUser)
refreshed, err := liveCache.RefreshLiveLease(s.ctx, accountID, userID, apiKeyID, "live-integration")
require.NoError(s.T(), err)
require.True(s.T(), refreshed)
require.NoError(s.T(), liveCache.ReleaseLiveLease(s.ctx, accountID, userID, apiKeyID, "live-integration"))
regularAccount, err = s.cache.AcquireAccountSlot(s.ctx, accountID, 1, "regular-account")
require.NoError(s.T(), err)
require.True(s.T(), regularAccount)
}
func (s *ConcurrencyCacheSuite) TestAccountSlot_AcquireAndRelease() {
accountID := int64(10)
reqID1, reqID2, reqID3 := "req1", "req2", "req3"
@@ -0,0 +1,73 @@
package repository
import (
"context"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/alicebob/miniredis/v2"
"github.com/redis/go-redis/v9"
"github.com/stretchr/testify/require"
)
func TestLiveLeaseReplacesRegularSlotsAndCountsTowardLimits(t *testing.T) {
redisServer := miniredis.RunT(t)
client := redis.NewClient(&redis.Options{Addr: redisServer.Addr()})
regular := NewConcurrencyCache(client, 15, 900)
live, ok := regular.(service.LiveConcurrencyCache)
require.True(t, ok)
ctx := context.Background()
accountAcquired, err := regular.AcquireAccountSlot(ctx, 10, 1, "regular-account")
require.NoError(t, err)
require.True(t, accountAcquired)
userAcquired, err := regular.AcquireUserSlot(ctx, 20, 1, "regular-user")
require.NoError(t, err)
require.True(t, userAcquired)
acquired, err := live.AcquireLiveLease(ctx, 10, 1, 20, 1, 30, "live-lease", true)
require.NoError(t, err)
require.True(t, acquired)
require.NoError(t, regular.ReleaseAccountSlot(ctx, 10, "regular-account"))
require.NoError(t, regular.ReleaseUserSlot(ctx, 20, "regular-user"))
accountCount, err := regular.GetAccountConcurrency(ctx, 10)
require.NoError(t, err)
require.Equal(t, 1, accountCount)
userCount, err := regular.GetUserConcurrency(ctx, 20)
require.NoError(t, err)
require.Equal(t, 1, userCount)
accountAcquired, err = regular.AcquireAccountSlot(ctx, 10, 1, "ordinary-blocked")
require.NoError(t, err)
require.False(t, accountAcquired)
refreshed, err := live.RefreshLiveLease(ctx, 10, 20, 30, "live-lease")
require.NoError(t, err)
require.True(t, refreshed)
require.NoError(t, live.ReleaseLiveLease(ctx, 10, 20, 30, "live-lease"))
accountAcquired, err = regular.AcquireAccountSlot(ctx, 10, 1, "ordinary-allowed")
require.NoError(t, err)
require.True(t, accountAcquired)
}
func TestLiveLeaseExpiresWithoutRefresh(t *testing.T) {
redisServer := miniredis.RunT(t)
client := redis.NewClient(&redis.Options{Addr: redisServer.Addr()})
regular := NewConcurrencyCache(client, 15, 900)
live, ok := regular.(service.LiveConcurrencyCache)
require.True(t, ok)
ctx := context.Background()
acquired, err := live.AcquireLiveLease(ctx, 10, 1, 20, 1, 30, "expired-live", false)
require.NoError(t, err)
require.True(t, acquired)
redisServer.FastForward(61 * time.Second)
acquired, err = regular.AcquireAccountSlot(ctx, 10, 1, "ordinary-after-expiry")
require.NoError(t, err)
require.True(t, acquired)
refreshed, err := live.RefreshLiveLease(ctx, 10, 20, 30, "expired-live")
require.NoError(t, err)
require.False(t, refreshed)
}
@@ -2,7 +2,10 @@ package repository
import (
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"strconv"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
@@ -10,6 +13,7 @@ import (
)
const stickySessionPrefix = "sticky_session:"
const liveCallPrefix = "live:call:"
type gatewayCache struct {
rdb *redis.Client
@@ -54,6 +58,7 @@ func (c *gatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64
// Compile-time assertion: gatewayCache must implement CyberSessionBlockStore.
var _ service.CyberSessionBlockStore = (*gatewayCache)(nil)
var _ service.LiveCallStore = (*gatewayCache)(nil)
const cyberSessionBlockPrefix = "cyber_session_block:"
@@ -71,3 +76,142 @@ func (c *gatewayCache) IsCyberSessionBlocked(ctx context.Context, key string) (b
}
return n > 0, nil
}
var claimLiveControllerScript = redis.NewScript(`
local key = KEYS[1]
local target = ARGV[1]
local owner = ARGV[2]
local current = redis.call('HGET', key, 'controller')
if current == false or current == 'closed' then
return 0
end
if target == 'observer' and current ~= 'pending' then
return 0
end
if target == 'proxy' and current ~= 'pending' and current ~= 'observer' and
(current ~= 'proxy' or redis.call('HGET', key, 'controller_owner') ~= owner) then
return 0
end
redis.call('HSET', key, 'controller', target, 'controller_owner', owner)
return 1
`)
var markLiveCallClosedScript = redis.NewScript(`
local key = KEYS[1]
if redis.call('EXISTS', key) == 0 then
return 0
end
if redis.call('HGET', key, 'controller') == 'closed' then
return 0
end
redis.call('HSET', key, 'controller', 'closed', 'controller_owner', '')
redis.call('EXPIRE', key, ARGV[1])
return 1
`)
var releaseLiveControllerScript = redis.NewScript(`
local key = KEYS[1]
if redis.call('HGET', key, 'controller') ~= 'proxy' or
redis.call('HGET', key, 'controller_owner') ~= ARGV[1] then
return 0
end
redis.call('HSET', key, 'controller', 'pending', 'controller_owner', '')
return 1
`)
func liveCallKey(callHash string) string {
return liveCallPrefix + callHash
}
func HashLiveCallID(callID string) string {
sum := sha256.Sum256([]byte(callID))
return hex.EncodeToString(sum[:])
}
func (c *gatewayCache) SaveLiveCall(ctx context.Context, record *service.LiveCallRecord, ttl time.Duration) error {
if record == nil || record.CallHash == "" || record.CallID == "" {
return fmt.Errorf("invalid live call record")
}
values := map[string]any{
"call_id": record.CallID,
"account_id": record.AccountID,
"api_key_id": record.APIKeyID,
"user_id": record.UserID,
"group_id": record.GroupID,
"subscription_id": record.SubscriptionID,
"lease_id": record.LeaseID,
"model": record.Model,
"created_at": record.CreatedAt.UnixMilli(),
"expires_at": record.ExpiresAt.UnixMilli(),
"controller": record.Controller,
"controller_owner": record.ControllerOwner,
"user_agent": record.UserAgent,
"ip_address": record.IPAddress,
"inbound_endpoint": record.InboundEndpoint,
"attestation": record.AttestationCiphertext,
}
key := liveCallKey(record.CallHash)
pipe := c.rdb.TxPipeline()
pipe.HSet(ctx, key, values)
pipe.Expire(ctx, key, ttl)
_, err := pipe.Exec(ctx)
return err
}
func (c *gatewayCache) GetLiveCall(ctx context.Context, callHash string) (*service.LiveCallRecord, error) {
values, err := c.rdb.HGetAll(ctx, liveCallKey(callHash)).Result()
if err != nil {
return nil, err
}
if len(values) == 0 {
return nil, service.ErrLiveCallNotFound
}
parseInt := func(field string) int64 {
value, _ := strconv.ParseInt(values[field], 10, 64)
return value
}
createdAt := time.UnixMilli(parseInt("created_at"))
expiresAt := time.UnixMilli(parseInt("expires_at"))
return &service.LiveCallRecord{
CallID: values["call_id"],
CallHash: callHash,
AccountID: parseInt("account_id"),
APIKeyID: parseInt("api_key_id"),
UserID: parseInt("user_id"),
GroupID: parseInt("group_id"),
SubscriptionID: parseInt("subscription_id"),
LeaseID: values["lease_id"],
Model: values["model"],
CreatedAt: createdAt,
ExpiresAt: expiresAt,
Controller: values["controller"],
ControllerOwner: values["controller_owner"],
UserAgent: values["user_agent"],
IPAddress: values["ip_address"],
InboundEndpoint: values["inbound_endpoint"],
AttestationCiphertext: values["attestation"],
}, nil
}
func (c *gatewayCache) ClaimLiveController(ctx context.Context, callHash, controller, owner string) (bool, error) {
result, err := claimLiveControllerScript.Run(ctx, c.rdb, []string{liveCallKey(callHash)}, controller, owner).Int()
return result == 1, err
}
func (c *gatewayCache) GetLiveController(ctx context.Context, callHash string) (string, error) {
value, err := c.rdb.HGet(ctx, liveCallKey(callHash), "controller").Result()
if err == redis.Nil {
return "", service.ErrLiveCallNotFound
}
return value, err
}
func (c *gatewayCache) ReleaseLiveController(ctx context.Context, callHash, owner string) (bool, error) {
result, err := releaseLiveControllerScript.Run(ctx, c.rdb, []string{liveCallKey(callHash)}, owner).Int()
return result == 1, err
}
func (c *gatewayCache) MarkLiveCallClosed(ctx context.Context, callHash string, ttl time.Duration) (bool, error) {
result, err := markLiveCallClosedScript.Run(ctx, c.rdb, []string{liveCallKey(callHash)}, int64(ttl.Seconds())).Int()
return result == 1, err
}
@@ -0,0 +1,62 @@
package repository
import (
"context"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/alicebob/miniredis/v2"
"github.com/redis/go-redis/v9"
"github.com/stretchr/testify/require"
)
func TestGatewayCacheLiveCallIdentityAndController(t *testing.T) {
redisServer := miniredis.RunT(t)
client := redis.NewClient(&redis.Options{Addr: redisServer.Addr()})
cache, ok := NewGatewayCache(client).(service.LiveCallStore)
require.True(t, ok)
otherInstance, ok := NewGatewayCache(client).(service.LiveCallStore)
require.True(t, ok)
record := &service.LiveCallRecord{
CallID: "call_secret",
CallHash: HashLiveCallID("call_secret"),
AccountID: 11,
APIKeyID: 22,
UserID: 33,
GroupID: 44,
LeaseID: "lease",
Model: "gpt-live-test",
AttestationCiphertext: "encrypted-attestation",
CreatedAt: time.Now(),
ExpiresAt: time.Now().Add(time.Hour),
Controller: service.LiveControllerPending,
}
require.NoError(t, cache.SaveLiveCall(context.Background(), record, time.Hour))
loaded, err := otherInstance.GetLiveCall(context.Background(), record.CallHash)
require.NoError(t, err)
require.Equal(t, record.CallID, loaded.CallID)
require.Equal(t, record.AccountID, loaded.AccountID)
require.Equal(t, record.AttestationCiphertext, loaded.AttestationCiphertext)
claimed, err := cache.ClaimLiveController(context.Background(), record.CallHash, service.LiveControllerObserver, "observer-1")
require.NoError(t, err)
require.True(t, claimed)
claimed, err = cache.ClaimLiveController(context.Background(), record.CallHash, service.LiveControllerProxy, "proxy-1")
require.NoError(t, err)
require.True(t, claimed)
controller, err := cache.GetLiveController(context.Background(), record.CallHash)
require.NoError(t, err)
require.Equal(t, service.LiveControllerProxy, controller)
released, err := cache.ReleaseLiveController(context.Background(), record.CallHash, "proxy-1")
require.NoError(t, err)
require.True(t, released)
closed, err := cache.MarkLiveCallClosed(context.Background(), record.CallHash, time.Hour)
require.NoError(t, err)
require.True(t, closed)
closed, err = cache.MarkLiveCallClosed(context.Background(), record.CallHash, time.Hour)
require.NoError(t, err)
require.False(t, closed)
}
@@ -90,6 +90,7 @@ func createGroupRecord(ctx context.Context, client *dbent.Client, groupIn *servi
SetModelRoutingEnabled(groupIn.ModelRoutingEnabled).
SetMcpXMLInject(groupIn.MCPXMLInject).
SetAllowMessagesDispatch(groupIn.AllowMessagesDispatch).
SetAllowLive(groupIn.AllowLive).
SetRequireOauthOnly(groupIn.RequireOAuthOnly).
SetRequirePrivacySet(groupIn.RequirePrivacySet).
SetDefaultMappedModel(groupIn.DefaultMappedModel).
@@ -255,6 +256,7 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er
SetModelRoutingEnabled(groupIn.ModelRoutingEnabled).
SetMcpXMLInject(groupIn.MCPXMLInject).
SetAllowMessagesDispatch(groupIn.AllowMessagesDispatch).
SetAllowLive(groupIn.AllowLive).
SetRequireOauthOnly(groupIn.RequireOAuthOnly).
SetRequirePrivacySet(groupIn.RequirePrivacySet).
SetDefaultMappedModel(groupIn.DefaultMappedModel).
@@ -140,7 +140,20 @@ func dockerIsAvailable(ctx context.Context) bool {
return cmd.Run() == nil
}
// selectDockerImage resolves the container image for the harness.
//
// SUB2API_TEST_POSTGRES_IMAGE overrides the PostgreSQL image so the suite can be
// run against the oldest documented-supported server, not just the newest. That
// matters for SQL that behaves differently across major versions: jsonpath
// .datetime() only accepts the ISO-8601 "Z" designator from PostgreSQL 17 on, so
// a suite pinned to 18 cannot observe breakage on 14-16.
//
// SUB2API_TEST_POSTGRES_IMAGE=postgres:15-alpine go test -tags integration ./internal/repository/
func selectDockerImage(ctx context.Context, preferred string) string {
if override := strings.TrimSpace(os.Getenv("SUB2API_TEST_POSTGRES_IMAGE")); override != "" &&
strings.HasPrefix(preferred, "postgres:") {
return override
}
if dockerImageExists(ctx, preferred) {
return preferred
}
@@ -55,6 +55,9 @@ func TestMigrationsRunner_IsIdempotent_AndSchemaIsUpToDate(t *testing.T) {
requireColumn(t, tx, "accounts", "session_window_status", "character varying", 20, true)
requireIndex(t, tx, "accounts", "idx_accounts_autopause_expiry_due")
// groups: OpenAI Live 默认关闭,管理员显式开启后才可访问。
requireColumn(t, tx, "groups", "allow_live", "boolean", 0, false)
// api_keys: key length should be 128
requireColumn(t, tx, "api_keys", "key", "character varying", 128, false)
@@ -79,6 +79,7 @@ var usageLogInsertArgTypes = [...]string{
"text", // billing_tier
"text", // billing_mode
"numeric", // account_stats_cost
"text", // session_id
"timestamptz", // created_at
}
@@ -274,6 +275,7 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor,
billing_tier,
billing_mode,
account_stats_cost,
session_id,
created_at
) VALUES (
$1, $2, $3, $4, $5, $6, $7,
@@ -281,7 +283,7 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor,
$10, $11, $12, $13,
$14, $15, $16, $17,
$18, $19, $20, $21, $22, $23,
$24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56
$24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56, $57
)
ON CONFLICT (request_id, api_key_id) DO NOTHING
RETURNING id, created_at
@@ -728,10 +730,13 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage
billing_tier,
billing_mode,
account_stats_cost,
session_id,
created_at
) AS (VALUES `)
args := make([]any, 0, len(keys)*56)
// Each batch row prepends the synthetic input_index before the 57
// usage-log column values.
args := make([]any, 0, len(keys)*58)
argPos := 1
for idx, key := range keys {
if idx > 0 {
@@ -815,6 +820,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage
billing_tier,
billing_mode,
account_stats_cost,
session_id,
created_at
)
SELECT
@@ -873,6 +879,7 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage
billing_tier,
billing_mode,
account_stats_cost,
session_id,
created_at
FROM input
ON CONFLICT (request_id, api_key_id) DO NOTHING
@@ -971,10 +978,11 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (
billing_tier,
billing_mode,
account_stats_cost,
session_id,
created_at
) AS (VALUES `)
args := make([]any, 0, len(preparedList)*56)
args := make([]any, 0, len(preparedList)*57)
argPos := 1
for idx, prepared := range preparedList {
if idx > 0 {
@@ -1055,6 +1063,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (
billing_tier,
billing_mode,
account_stats_cost,
session_id,
created_at
)
SELECT
@@ -1113,6 +1122,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (
billing_tier,
billing_mode,
account_stats_cost,
session_id,
created_at
FROM input
ON CONFLICT (request_id, api_key_id) DO NOTHING
@@ -1179,6 +1189,7 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared
billing_tier,
billing_mode,
account_stats_cost,
session_id,
created_at
) VALUES (
$1, $2, $3, $4, $5, $6, $7,
@@ -1186,7 +1197,7 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared
$10, $11, $12, $13,
$14, $15, $16, $17,
$18, $19, $20, $21, $22, $23,
$24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56
$24, $25, $26, $27, $28, $29, $30, $31, $32, $33, $34, $35, $36, $37, $38, $39, $40, $41, $42, $43, $44, $45, $46, $47, $48, $49, $50, $51, $52, $53, $54, $55, $56, $57
)
ON CONFLICT (request_id, api_key_id) DO NOTHING
`, prepared.args...)
@@ -1227,6 +1238,7 @@ func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared {
modelMappingChain := nullString(log.ModelMappingChain)
billingTier := nullString(log.BillingTier)
billingMode := nullString(log.BillingMode)
sessionID := nullString(log.SessionID)
requestedModel := strings.TrimSpace(log.RequestedModel)
if requestedModel == "" {
requestedModel = strings.TrimSpace(log.Model)
@@ -1299,6 +1311,7 @@ func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared {
billingTier,
billingMode,
log.AccountStatsCost, // account_stats_cost
sessionID, // session_id
createdAt,
},
}
@@ -19,7 +19,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/service"
)
const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, image_input_tokens, image_input_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, image_input_size, image_output_size, image_size_source, image_size_breakdown, video_count, video_resolution, video_duration_seconds, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, created_at"
const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, image_output_tokens, image_output_cost, image_input_tokens, image_input_cost, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, account_rate_multiplier, billing_type, request_type, stream, openai_ws_mode, duration_ms, first_token_ms, user_agent, ip_address, image_count, image_size, image_input_size, image_output_size, image_size_source, image_size_breakdown, video_count, video_resolution, video_duration_seconds, service_tier, reasoning_effort, inbound_endpoint, upstream_endpoint, cache_ttl_overridden, long_context_billing_applied, channel_id, model_mapping_chain, billing_tier, billing_mode, account_stats_cost, session_id, created_at"
func (r *usageLogRepository) GetByID(ctx context.Context, id int64) (log *service.UsageLog, err error) {
query := "SELECT " + usageLogSelectColumns + " FROM usage_logs WHERE id = $1"
@@ -481,6 +481,7 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
billingTier sql.NullString
billingMode sql.NullString
accountStatsCost sql.NullFloat64
sessionID sql.NullString
createdAt time.Time
)
@@ -541,6 +542,7 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
&billingTier,
&billingMode,
&accountStatsCost,
&sessionID,
&createdAt,
); err != nil {
return nil, err
@@ -661,6 +663,9 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
if accountStatsCost.Valid {
log.AccountStatsCost = &accountStatsCost.Float64
}
if sessionID.Valid {
log.SessionID = &sessionID.String
}
return log, nil
}
@@ -96,6 +96,7 @@ func TestUsageLogRepositoryCreateSyncRequestTypeAndLegacyFields(t *testing.T) {
sqlmock.AnyArg(), // billing_tier
sqlmock.AnyArg(), // billing_mode
sqlmock.AnyArg(), // account_stats_cost
sqlmock.AnyArg(), // session_id
createdAt,
).
WillReturnRows(sqlmock.NewRows([]string{"id", "created_at"}).AddRow(int64(99), createdAt))
@@ -185,6 +186,7 @@ func TestUsageLogRepositoryCreate_PersistsServiceTier(t *testing.T) {
sqlmock.AnyArg(), // billing_tier
sqlmock.AnyArg(), // billing_mode
sqlmock.AnyArg(), // account_stats_cost
sqlmock.AnyArg(), // session_id
createdAt,
).
WillReturnRows(sqlmock.NewRows([]string{"id", "created_at"}).AddRow(int64(100), createdAt))
@@ -826,6 +828,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
sql.NullString{},
sql.NullString{},
sql.NullFloat64{},
sql.NullString{},
now,
}})
require.NoError(t, err)
@@ -900,6 +903,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
sql.NullString{}, // billing_tier
sql.NullString{}, // billing_mode
sql.NullFloat64{}, // account_stats_cost
sql.NullString{}, // session_id
now,
}})
require.NoError(t, err)
@@ -957,6 +961,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
sql.NullString{}, // billing_tier
sql.NullString{}, // billing_mode
sql.NullFloat64{}, // account_stats_cost
sql.NullString{}, // session_id
now,
}})
require.NoError(t, err)
@@ -1014,6 +1019,7 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
sql.NullString{}, // billing_tier
sql.NullString{}, // billing_mode
sql.NullFloat64{}, // account_stats_cost
sql.NullString{}, // session_id
now,
}})
require.NoError(t, err)
@@ -0,0 +1,71 @@
//go:build integration
package repository
import (
"context"
"testing"
"time"
"github.com/google/uuid"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
// TestUsageLog_SessionIDPersistence proves session_id round-trips from insert to
// read and is omitted (NULL) when absent.
func TestUsageLog_SessionIDPersistence(t *testing.T) {
ctx := context.Background()
client := testEntClient(t)
repo := newUsageLogRepositoryWithSQL(client, integrationDB)
user := mustCreateUser(t, client, &service.User{Email: "session-id-" + uuid.NewString() + "@example.com"})
apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-session-" + uuid.NewString(), Name: "k"})
account := mustCreateAccount(t, client, &service.Account{Name: "acc-session-" + uuid.NewString()})
sessionID := "sess-" + uuid.NewString()
withSession := &service.UsageLog{
UserID: user.ID,
APIKeyID: apiKey.ID,
AccountID: account.ID,
RequestID: uuid.NewString(),
Model: "claude-3",
InputTokens: 10,
OutputTokens: 5,
TotalCost: 1.0,
ActualCost: 1.0,
SessionID: &sessionID,
CreatedAt: time.Now().UTC(),
}
_, err := repo.Create(ctx, withSession)
require.NoError(t, err)
require.NotZero(t, withSession.ID)
withoutSession := &service.UsageLog{
UserID: user.ID,
APIKeyID: apiKey.ID,
AccountID: account.ID,
RequestID: uuid.NewString(),
Model: "claude-3",
InputTokens: 7,
OutputTokens: 3,
TotalCost: 0.5,
ActualCost: 0.5,
CreatedAt: time.Now().UTC(),
}
_, err = repo.Create(ctx, withoutSession)
require.NoError(t, err)
// Round-trip: session id survives insert → read.
got, err := repo.GetByID(ctx, withSession.ID)
require.NoError(t, err)
require.NotNil(t, got.SessionID)
require.Equal(t, sessionID, *got.SessionID)
// Omission: absent session id reads back as nil (NULL), not empty string.
gotNone, err := repo.GetByID(ctx, withoutSession.ID)
require.NoError(t, err)
require.Nil(t, gotNone.SessionID)
}
@@ -0,0 +1,91 @@
//go:build unit
package repository
import (
"database/sql"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
func newSessionIDUsageLog(sessionID *string) *service.UsageLog {
return &service.UsageLog{
UserID: 1,
APIKeyID: 2,
AccountID: 3,
RequestID: "req-session-id",
Model: "claude-3",
InputTokens: 10,
OutputTokens: 5,
TotalCost: 1.0,
ActualCost: 1.0,
SessionID: sessionID,
CreatedAt: time.Now().UTC(),
}
}
// TestPrepareUsageLogInsert_SessionIDArgWiring pins the session_id column to the
// arg slice / arg-type table so the five INSERT column lists stay in sync. session_id
// is the penultimate arg (created_at is always last).
func TestPrepareUsageLogInsert_SessionIDArgWiring(t *testing.T) {
require.Len(t, usageLogInsertArgTypes, 57, "arg-type table must include session_id")
sessionID := "sess-persisted-123"
prepared := prepareUsageLogInsert(newSessionIDUsageLog(&sessionID))
require.Len(t, prepared.args, len(usageLogInsertArgTypes),
"prepared args must match the arg-type table length")
// created_at is last; session_id is the arg immediately before it.
sessionArg := prepared.args[len(prepared.args)-2]
ns, ok := sessionArg.(sql.NullString)
require.True(t, ok, "session_id arg should be a sql.NullString, got %T", sessionArg)
require.True(t, ns.Valid)
require.Equal(t, sessionID, ns.String)
require.Equal(t, "text", usageLogInsertArgTypes[len(usageLogInsertArgTypes)-2],
"session_id arg type must be text")
}
// TestPrepareUsageLogInsert_SessionIDNullWhenAbsent proves an absent session id is
// persisted as SQL NULL rather than an empty string.
func TestPrepareUsageLogInsert_SessionIDNullWhenAbsent(t *testing.T) {
prepared := prepareUsageLogInsert(newSessionIDUsageLog(nil))
sessionArg := prepared.args[len(prepared.args)-2]
ns, ok := sessionArg.(sql.NullString)
require.True(t, ok, "session_id arg should be a sql.NullString, got %T", sessionArg)
require.False(t, ns.Valid, "absent session id must be NULL, not empty string")
empty := ""
preparedEmpty := prepareUsageLogInsert(newSessionIDUsageLog(&empty))
nsEmpty := preparedEmpty.args[len(preparedEmpty.args)-2].(sql.NullString)
require.False(t, nsEmpty.Valid, "empty session id must also be NULL")
}
// TestUsageLogInsertQueries_IncludeSessionID guards that every generated INSERT path
// and the SELECT column list reference session_id.
func TestUsageLogInsertQueries_IncludeSessionID(t *testing.T) {
require.Contains(t, usageLogSelectColumns, "session_id",
"SELECT column list must include session_id")
sessionID := "sess-in-query"
log := newSessionIDUsageLog(&sessionID)
prepared := prepareUsageLogInsert(log)
key := usageLogBatchKey(log.RequestID, log.APIKeyID)
batchQuery, batchArgs := buildUsageLogBatchInsertQuery([]string{key},
map[string]usageLogInsertPrepared{key: prepared})
require.Contains(t, batchQuery, "session_id")
// Two column references (INSERT column list + SELECT ... FROM input) plus the CTE def.
require.GreaterOrEqual(t, strings.Count(batchQuery, "session_id"), 3)
require.Len(t, batchArgs, len(prepared.args)+1,
"batch args include the synthetic input_index before usage-log values")
bestEffortQuery, bestEffortArgs := buildUsageLogBestEffortInsertQuery([]usageLogInsertPrepared{prepared})
require.Contains(t, bestEffortQuery, "session_id")
require.Len(t, bestEffortArgs, len(prepared.args))
}
+118 -1
View File
@@ -43,6 +43,16 @@ func newUserRepositoryWithSQL(client *dbent.Client, sqlq sqlExecutor) *userRepos
}
func (r *userRepository) Create(ctx context.Context, userIn *service.User) error {
return r.create(ctx, userIn, false)
}
// CreateWithEmailAliasGuard 见 service.UserRepository:在邮箱唯一性锁内复查收件箱身份,
// 供注册路径使用。
func (r *userRepository) CreateWithEmailAliasGuard(ctx context.Context, userIn *service.User) error {
return r.create(ctx, userIn, true)
}
func (r *userRepository) create(ctx context.Context, userIn *service.User, guardEmailAlias bool) error {
if userIn == nil {
return nil
}
@@ -69,11 +79,16 @@ func (r *userRepository) Create(ctx context.Context, userIn *service.User) error
}
}
lockKeys := []string{normalizedEmailUniquenessLockKey(userIn.Email)}
if guardEmailAlias {
// 别名变体的字面量不同,唯一索引无法兜底;用收件箱身份锁把同一收件箱的并发注册串行化。
lockKeys = append(lockKeys, emailAliasUniquenessLockKey(userIn.Email))
}
releaseEmailLock, err := lockRepositoryScopedKeys(
txCtx,
txClient,
txAwareSQLExecutor(txCtx, r.sql, r.client),
normalizedEmailUniquenessLockKey(userIn.Email),
lockKeys...,
)
if err != nil {
return err
@@ -84,6 +99,16 @@ func (r *userRepository) Create(ctx context.Context, userIn *service.User) error
return err
}
if guardEmailAlias {
aliasExists, err := existsByEmailAliasWithClient(txCtx, txClient, userIn.Email)
if err != nil {
return err
}
if aliasExists {
return service.ErrEmailExists
}
}
created, err := txClient.User.Create().
SetEmail(userIn.Email).
SetUsername(userIn.Username).
@@ -904,6 +929,88 @@ func (r *userRepository) ExistsByEmail(ctx context.Context, email string) (bool,
return r.client.User.Query().Where(userEmailLookupPredicate(email)).Exist(ctx)
}
// emailAliasCandidateLimit 限制一次别名查重最多取回的候选行数。探针都以去点后的
// 本地部分为前缀锚定(见 dotStrippedEmailExpr),正常收件箱的变体只有个位数;
// 上限只是兜底,避免公开未鉴权的注册/发码端点把大表整张读进内存。
const emailAliasCandidateLimit = 50
// ExistsByEmailAlias 见 service.UserRepository。软删除过滤沿用 ExistsByEmail 的默认行为。
func (r *userRepository) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) {
return existsByEmailAliasWithClient(ctx, clientFromContext(ctx, r.client), email)
}
func existsByEmailAliasWithClient(ctx context.Context, client *dbent.Client, email string) (bool, error) {
if client == nil {
return false, nil
}
probes := service.EmailAliasDedupProbes(email)
if len(probes) == 0 {
return false, nil
}
preds := make([]predicate.User, 0, 2*len(probes))
for _, probe := range probes {
preds = append(preds,
dotStrippedEmailEQ(probe.Local+"@"+probe.Domain),
// "+后缀"的内容未知,只能按前缀匹配。
dotStrippedEmailLike(escapeLikeWildcards(probe.Local)+"+%@"+escapeLikeWildcards(probe.Domain)),
)
}
candidates, err := client.User.Query().
Where(dbuser.Or(preds...)).
Limit(emailAliasCandidateLimit).
Select(dbuser.FieldEmail).
Strings(ctx)
if err != nil {
return false, err
}
// 探针会有过度匹配(点号只在 Gmail 家族无意义),最终判定必须回到完整归一化规则。
identity := service.NormalizeEmailForAliasDedup(email)
for _, candidate := range candidates {
if service.NormalizeEmailForAliasDedup(candidate) == identity {
return true, nil
}
}
return false, nil
}
// dotStrippedEmailExpr 渲染下面的表达式:去掉存量邮箱的大小写、首尾空白(与
// userEmailLookupPredicate 的精确匹配口径一致,历史数据存在带空白的行)以及全部点号。
//
// REPLACE(LOWER(TRIM(email)), '.', '')
//
// 两侧都去点,因此一个域名探针即可同时覆盖 Gmail 点号变体与 FQDN 根点(user@gmail.com.)。
// migrations/190 为同一表达式建了索引。
func dotStrippedEmailExpr(b *entsql.Builder, s *entsql.Selector) *entsql.Builder {
return b.WriteString("REPLACE(LOWER(TRIM(").
Ident(s.C(dbuser.FieldEmail)).
WriteString(")), '.', '')")
}
func dotStrippedEmailEQ(value string) predicate.User {
return predicate.User(func(s *entsql.Selector) {
s.Where(entsql.P(func(b *entsql.Builder) {
dotStrippedEmailExpr(b, s).WriteString(" = ").Arg(value)
}))
})
}
func dotStrippedEmailLike(pattern string) predicate.User {
return predicate.User(func(s *entsql.Selector) {
s.Where(entsql.P(func(b *entsql.Builder) {
dotStrippedEmailExpr(b, s).WriteString(" LIKE ").Arg(pattern).WriteString(` ESCAPE '\'`)
}))
})
}
// escapeLikeWildcards 转义 LIKE 元字符:本地部分合法可含 % 与 _,不转义会扩大匹配面。
var likeWildcardEscaper = strings.NewReplacer(`\`, `\\`, "%", `\%`, "_", `\_`)
func escapeLikeWildcards(value string) string {
return likeWildcardEscaper.Replace(value)
}
func ensureNormalizedEmailAvailableWithClient(ctx context.Context, client *dbent.Client, userID int64, email string) error {
client = clientFromContext(ctx, client)
if client == nil {
@@ -951,6 +1058,16 @@ func normalizedEmailUniquenessLockKey(email string) string {
return "users:normalized-email:" + normalized
}
// emailAliasUniquenessLockKey 按收件箱身份(而非邮箱字面量)加锁,使同一收件箱的不同
// 别名变体在注册时互斥。
func emailAliasUniquenessLockKey(email string) string {
identity := service.NormalizeEmailForAliasDedup(email)
if identity == "" {
return ""
}
return "users:email-alias-identity:" + identity
}
func (r *userRepository) AddGroupToAllowedGroups(ctx context.Context, userID int64, groupID int64) error {
client := clientFromContext(ctx, r.client)
err := client.UserAllowedGroup.Create().
@@ -0,0 +1,100 @@
package repository
import (
"context"
"testing"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
func seedUserForAliasTest(t *testing.T, repo *userRepository, email string) {
t.Helper()
require.NoError(t, repo.Create(context.Background(), &service.User{
Email: email,
Username: email,
PasswordHash: "hash",
Role: service.RoleUser,
Status: service.StatusActive,
}))
}
func TestUserRepositoryExistsByEmailAlias(t *testing.T) {
cases := []struct {
name string
stored string
probe string
want bool
}{
{"same address", "someone@gmail.com", "someone@gmail.com", true},
{"gmail plus alias", "someone@gmail.com", "someone+bulk294@gmail.com", true},
{"gmail dot trick", "d.axis.2026@gmail.com", "daxis2026@gmail.com", true},
{"gmail dot trick both sides", "d.axis.2026@gmail.com", "da.xis.2026@gmail.com", true},
{"stored plus alias found by canonical form", "someone+tag@gmail.com", "someone@gmail.com", true},
{"googlemail is a gmail alias", "someone@googlemail.com", "some.one@gmail.com", true},
{"fqdn root dot on probe", "d.axis.2026@gmail.com", "da.xis.2026@gmail.com.", true},
{"fqdn root dot on stored row", "d.axis.2026@gmail.com.", "daxis2026@gmail.com", true},
{"legacy row with spacing and case", " D.Axis.2026@Gmail.com ", "daxis2026@gmail.com", true},
{"non-gmail plus alias", "first.last@qq.com", "first.last+tag@qq.com", true},
{"different gmail inbox", "someone@gmail.com", "someoneelse@gmail.com", false},
{"non-gmail dots are significant", "first.last@qq.com", "firstlast@qq.com", false},
{"different domain", "someone@gmail.com", "someone@qq.com", false},
{"distinct plus-prefixed locals", "+alice@gmail.com", "+bob@gmail.com", false},
{"underscore is not a wildcard", "user_x@qq.com", "userax@qq.com", false},
{"percent is not a wildcard", "a%b@qq.com", "axxb@qq.com", false},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
repo, _ := newUserEntRepo(t)
seedUserForAliasTest(t, repo, tc.stored)
got, err := repo.ExistsByEmailAlias(context.Background(), tc.probe)
require.NoError(t, err)
require.Equal(t, tc.want, got)
})
}
}
func TestUserRepositoryExistsByEmailAliasIgnoresMalformedInput(t *testing.T) {
repo, _ := newUserEntRepo(t)
seedUserForAliasTest(t, repo, "someone@gmail.com")
got, err := repo.ExistsByEmailAlias(context.Background(), "not-an-email")
require.NoError(t, err)
require.False(t, got)
}
func TestUserRepositoryCreateWithEmailAliasGuard(t *testing.T) {
repo, _ := newUserEntRepo(t)
ctx := context.Background()
seedUserForAliasTest(t, repo, "d.axis.2026@gmail.com")
// 注册路径:别名变体在唯一性锁内被拒绝。
err := repo.CreateWithEmailAliasGuard(ctx, &service.User{
Email: "da.xis.2026+free@googlemail.com",
Username: "alias-variant",
PasswordHash: "hash",
Role: service.RoleUser,
Status: service.StatusActive,
})
require.ErrorIs(t, err, service.ErrEmailExists)
// 不同收件箱仍可注册。
require.NoError(t, repo.CreateWithEmailAliasGuard(ctx, &service.User{
Email: "other.person@gmail.com",
Username: "other-person",
PasswordHash: "hash",
Role: service.RoleUser,
Status: service.StatusActive,
}))
// 管理员建号(Create)不受别名限制。
require.NoError(t, repo.Create(ctx, &service.User{
Email: "daxis2026+support@gmail.com",
Username: "admin-created",
PasswordHash: "hash",
Role: service.RoleUser,
Status: service.StatusActive,
}))
}
@@ -377,6 +377,7 @@ func TestAPIContracts(t *testing.T) {
"video_rate_multiplier": 0,
"claude_code_only": false,
"allow_messages_dispatch": false,
"allow_live": false,
"fallback_group_id": null,
"fallback_group_id_on_invalid_request": null,
"require_oauth_only": false,
@@ -1494,6 +1495,10 @@ func (r *stubUserRepo) Create(ctx context.Context, user *service.User) error {
return errors.New("not implemented")
}
func (r *stubUserRepo) CreateWithEmailAliasGuard(ctx context.Context, user *service.User) error {
return errors.New("not implemented")
}
func (r *stubUserRepo) GetByID(ctx context.Context, id int64) (*service.User, error) {
user, ok := r.users[id]
if !ok {
@@ -1573,6 +1578,10 @@ func (r *stubUserRepo) ExistsByEmail(ctx context.Context, email string) (bool, e
return false, errors.New("not implemented")
}
func (r *stubUserRepo) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) {
return false, errors.New("not implemented")
}
func (r *stubUserRepo) RemoveGroupFromAllowedGroups(ctx context.Context, groupID int64) (int64, error) {
return 0, errors.New("not implemented")
}
@@ -131,6 +131,10 @@ func (s *stubUserRepo) Create(ctx context.Context, user *service.User) error {
panic("unexpected Create call")
}
func (s *stubUserRepo) CreateWithEmailAliasGuard(ctx context.Context, user *service.User) error {
panic("unexpected CreateWithEmailAliasGuard call")
}
func (s *stubUserRepo) GetByID(ctx context.Context, id int64) (*service.User, error) {
if s.getByID == nil {
panic("GetByID not stubbed")
@@ -208,6 +212,10 @@ func (s *stubUserRepo) ExistsByEmail(ctx context.Context, email string) (bool, e
panic("unexpected ExistsByEmail call")
}
func (s *stubUserRepo) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) {
panic("unexpected ExistsByEmailAlias call")
}
func (s *stubUserRepo) RemoveGroupFromAllowedGroups(ctx context.Context, groupID int64) (int64, error) {
panic("unexpected RemoveGroupFromAllowedGroups call")
}
+1
View File
@@ -317,6 +317,7 @@ func registerGroupRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
groups.GET("/all", h.Admin.Group.GetAll)
groups.GET("/usage-summary", h.Admin.Group.GetUsageSummary)
groups.GET("/capacity-summary", h.Admin.Group.GetCapacitySummary)
groups.GET("/live-capability", h.Admin.Group.GetLiveCapability)
groups.PUT("/sort-order", h.Admin.Group.UpdateSortOrder)
groups.GET("/:id/models-list-candidates", h.Admin.Group.GetModelsListCandidates)
groups.GET("/:id/composite-routes", h.Admin.Group.ListCompositeRoutes)
@@ -180,6 +180,8 @@ func RegisterGatewayRoutes(
// Codex manifest format; other clients keep the OpenAI-style list.
gateway.GET("/models", modelsHandler)
gateway.GET("/usage", h.Gateway.Usage)
gateway.POST("/live", h.OpenAIGateway.Live)
gateway.GET("/live/:call_id", h.OpenAIGateway.LiveSideband)
// OpenAI Responses API: auto-route based on group platform
gateway.POST("/responses", func(c *gin.Context) {
if isOpenAIResponsesCompatibleGatewayPlatform(c) {
@@ -277,6 +279,8 @@ func RegisterGatewayRoutes(
codexDirect := r.Group("/backend-api/codex")
codexDirect.Use(bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic)
{
codexDirect.POST("/realtime/calls", h.OpenAIGateway.Live)
codexDirect.GET("/:call_id", h.OpenAIGateway.LiveSideband)
codexDirect.POST("/responses", responsesHandler)
codexDirect.POST("/responses/*subpath", responsesHandler)
codexDirect.POST("/alpha/search", textBodyLimit, h.OpenAIGateway.AlphaSearch)
@@ -34,6 +34,8 @@ func TestEveryGatewayPOSTRouteIsClassifiedForPromptAuditCoverage(t *testing.T) {
"/chat/completions": {"gateway_handler_chat_completions.go", "openai_chat_completions.go"},
"/embeddings": {"openai_embeddings.go"},
"/alpha/search": {"openai_alpha_search.go"},
"/live": {"openai_live.go"},
"/realtime/calls": {"openai_live.go"},
"/images/generations": {"openai_images.go", "grok_media.go"},
"/images/edits": {"openai_images.go", "grok_media.go"},
"/images/generations/async": {"image_task_handler.go"},
+6
View File
@@ -89,6 +89,7 @@ const (
OpenAIEndpointCapabilityChatCompletions OpenAIEndpointCapability = "chat_completions"
OpenAIEndpointCapabilityEmbeddings OpenAIEndpointCapability = "embeddings"
OpenAIEndpointCapabilityAlphaSearch OpenAIEndpointCapability = "alpha_search"
OpenAIEndpointCapabilityLive OpenAIEndpointCapability = "live"
// OpenAIEndpointCapabilityGrokMediaGeneration keeps image/video generation
// away from Grok accounts that are explicitly disabled or whose billing
// entitlement probe was forbidden. Video status lookups intentionally do not
@@ -1435,6 +1436,11 @@ func (a *Account) SupportsOpenAIEndpointCapability(capability OpenAIEndpointCapa
}
switch capability {
case OpenAIEndpointCapabilityChatCompletions:
case OpenAIEndpointCapabilityLive:
return a.Platform == PlatformOpenAI &&
a.Type == AccountTypeOAuth &&
!a.IsOpenAIPersonalAccessToken() &&
!a.IsOpenAIAgentIdentity()
case OpenAIEndpointCapabilityResponses:
// Responses 支持状态由 accounts.extra 的自动探测标记决定,而非
// credentials 能力集。已探测确认不支持 /v1/responses 的 APIKey 上游
+10
View File
@@ -470,6 +470,7 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
MCPXMLInject: mcpXMLInject,
SupportedModelScopes: input.SupportedModelScopes,
AllowMessagesDispatch: input.AllowMessagesDispatch,
AllowLive: input.AllowLive,
RequireOAuthOnly: input.RequireOAuthOnly,
RequirePrivacySet: input.RequirePrivacySet,
DefaultMappedModel: input.DefaultMappedModel,
@@ -480,6 +481,9 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
ReasoningEffortMappings: reasoningEffortMappings,
}
sanitizeGroupMessagesDispatchFields(group)
if group.Platform != PlatformOpenAI {
group.AllowLive = false
}
sanitizeGroupReasoningEffortPolicy(group)
if err := s.groupRepo.Create(ctx, group); err != nil {
return nil, err
@@ -777,6 +781,9 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
if input.AllowMessagesDispatch != nil {
group.AllowMessagesDispatch = *input.AllowMessagesDispatch
}
if input.AllowLive != nil {
group.AllowLive = *input.AllowLive
}
if input.RequireOAuthOnly != nil {
group.RequireOAuthOnly = *input.RequireOAuthOnly
}
@@ -810,6 +817,9 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
group.ReasoningEffortMappings = reasoningEffortMappings
}
sanitizeGroupMessagesDispatchFields(group)
if group.Platform != PlatformOpenAI {
group.AllowLive = false
}
sanitizeGroupReasoningEffortPolicy(group)
if err := s.groupRepo.Update(ctx, group); err != nil {
@@ -120,6 +120,7 @@ func cloneGroupForDuplicate(source *Group, operationID string) *Group {
SupportedModelScopes: append([]string(nil), source.SupportedModelScopes...),
SortOrder: source.SortOrder,
AllowMessagesDispatch: source.AllowMessagesDispatch,
AllowLive: source.AllowLive,
RequireOAuthOnly: source.RequireOAuthOnly,
RequirePrivacySet: source.RequirePrivacySet,
DefaultMappedModel: source.DefaultMappedModel,
@@ -162,6 +162,7 @@ func TestDuplicateGroupCopiesConfigurationDeeplyAndResetsRuntimeState(t *testing
SupportedModelScopes: []string{"claude", "gemini_text"},
SortOrder: 9,
AllowMessagesDispatch: true,
AllowLive: true,
RequireOAuthOnly: true,
RequirePrivacySet: true,
DefaultMappedModel: "gpt-5.4",
@@ -252,6 +252,7 @@ type CreateGroupInput struct {
SupportedModelScopes []string
// OpenAI Messages 调度配置(仅 openai 平台使用)
AllowMessagesDispatch bool
AllowLive bool
DefaultMappedModel string
RequireOAuthOnly bool
RequirePrivacySet bool
@@ -312,6 +313,7 @@ type UpdateGroupInput struct {
SupportedModelScopes *[]string
// OpenAI Messages 调度配置(仅 openai 平台使用)
AllowMessagesDispatch *bool
AllowLive *bool
DefaultMappedModel *string
RequireOAuthOnly *bool
RequirePrivacySet *bool
@@ -33,6 +33,9 @@ func (s *userRepoStubForGroupUpdate) AddGroupToAllowedGroups(_ context.Context,
}
func (s *userRepoStubForGroupUpdate) Create(context.Context, *User) error { panic("unexpected") }
func (s *userRepoStubForGroupUpdate) CreateWithEmailAliasGuard(context.Context, *User) error {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) GetByID(context.Context, int64) (*User, error) {
panic("unexpected")
}
@@ -81,6 +84,9 @@ func (s *userRepoStubForGroupUpdate) BatchUpdateLimits(context.Context, []int64,
func (s *userRepoStubForGroupUpdate) ExistsByEmail(context.Context, string) (bool, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) ExistsByEmailAlias(context.Context, string) (bool, error) {
panic("unexpected")
}
func (s *userRepoStubForGroupUpdate) RemoveGroupFromAllowedGroups(context.Context, int64) (int64, error) {
panic("unexpected")
}
@@ -13,18 +13,21 @@ import (
)
type userRepoStub struct {
user *User
getErr error
createErr error
deleteErr error
exists bool
existsErr error
nextID int64
created []*User
updated []*User
deletedIDs []int64
usersByEmail map[string]*User
getByEmailErr error
user *User
getErr error
createErr error
deleteErr error
exists bool
existsErr error
aliasExists bool
aliasErr error
guardedCreates int
nextID int64
created []*User
updated []*User
deletedIDs []int64
usersByEmail map[string]*User
getByEmailErr error
}
func (s *userRepoStub) Create(ctx context.Context, user *User) error {
@@ -43,6 +46,17 @@ func (s *userRepoStub) Create(ctx context.Context, user *User) error {
return nil
}
func (s *userRepoStub) CreateWithEmailAliasGuard(ctx context.Context, user *User) error {
s.guardedCreates++
if s.aliasErr != nil {
return s.aliasErr
}
if s.aliasExists {
return ErrEmailExists
}
return s.Create(ctx, user)
}
func (s *userRepoStub) GetByID(ctx context.Context, id int64) (*User, error) {
if s.getErr != nil {
return nil, s.getErr
@@ -144,6 +158,13 @@ func (s *userRepoStub) ExistsByEmail(ctx context.Context, email string) (bool, e
return s.exists, nil
}
func (s *userRepoStub) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) {
if s.aliasErr != nil {
return false, s.aliasErr
}
return s.aliasExists, nil
}
func (s *userRepoStub) RemoveGroupFromAllowedGroups(ctx context.Context, groupID int64) (int64, error) {
panic("unexpected RemoveGroupFromAllowedGroups call")
}
@@ -35,6 +35,10 @@ type emailSyncRepoStub struct {
replaceErr error
}
func (s *emailSyncRepoStub) CreateWithEmailAliasGuard(ctx context.Context, user *User) error {
return s.Create(ctx, user)
}
func (s *emailSyncRepoStub) Create(_ context.Context, user *User) error {
if s.nextID != 0 && user.ID == 0 {
user.ID = s.nextID
@@ -109,6 +113,10 @@ func (s *emailSyncRepoStub) UpdateConcurrency(context.Context, int64, int) error
func (s *emailSyncRepoStub) ExistsByEmail(context.Context, string) (bool, error) { return false, nil }
func (s *emailSyncRepoStub) ExistsByEmailAlias(context.Context, string) (bool, error) {
return false, nil
}
func (s *emailSyncRepoStub) RemoveGroupFromAllowedGroups(context.Context, int64) (int64, error) {
return 0, nil
}
@@ -957,6 +957,7 @@ func TestAdminService_CreateGroup_ClearsMessagesDispatchFieldsForNonOpenAIPlatfo
Platform: PlatformAnthropic,
RateMultiplier: 1.0,
AllowMessagesDispatch: true,
AllowLive: true,
DefaultMappedModel: "gpt-5.4",
MessagesDispatchModelConfig: OpenAIMessagesDispatchModelConfig{
OpusMappedModel: "gpt-5.4",
@@ -966,6 +967,7 @@ func TestAdminService_CreateGroup_ClearsMessagesDispatchFieldsForNonOpenAIPlatfo
require.NotNil(t, group)
require.NotNil(t, repo.created)
require.False(t, repo.created.AllowMessagesDispatch)
require.False(t, repo.created.AllowLive)
require.Empty(t, repo.created.DefaultMappedModel)
require.Equal(t, OpenAIMessagesDispatchModelConfig{}, repo.created.MessagesDispatchModelConfig)
}
@@ -977,6 +979,7 @@ func TestAdminService_UpdateGroup_ClearsMessagesDispatchFieldsWhenPlatformChange
Platform: PlatformOpenAI,
Status: StatusActive,
AllowMessagesDispatch: true,
AllowLive: true,
DefaultMappedModel: "gpt-5.4",
MessagesDispatchModelConfig: OpenAIMessagesDispatchModelConfig{
SonnetMappedModel: "gpt-5.3-codex",
@@ -993,6 +996,7 @@ func TestAdminService_UpdateGroup_ClearsMessagesDispatchFieldsWhenPlatformChange
require.NotNil(t, repo.updated)
require.Equal(t, PlatformAnthropic, repo.updated.Platform)
require.False(t, repo.updated.AllowMessagesDispatch)
require.False(t, repo.updated.AllowLive)
require.Empty(t, repo.updated.DefaultMappedModel)
require.Equal(t, OpenAIMessagesDispatchModelConfig{}, repo.updated.MessagesDispatchModelConfig)
}
@@ -94,6 +94,7 @@ type APIKeyAuthGroupSnapshot struct {
// OpenAI Messages 调度配置(仅 openai 平台使用)
AllowMessagesDispatch bool `json:"allow_messages_dispatch"`
AllowLive bool `json:"allow_live"`
DefaultMappedModel string `json:"default_mapped_model,omitempty"`
MessagesDispatchModelConfig OpenAIMessagesDispatchModelConfig `json:"messages_dispatch_model_config,omitempty"`
ModelsListConfig GroupModelsListConfig `json:"models_list_config,omitempty"`
@@ -14,7 +14,7 @@ import (
"github.com/dgraph-io/ristretto"
)
const apiKeyAuthSnapshotVersion = 16 // v16: include group reasoning effort ceiling and mappings
const apiKeyAuthSnapshotVersion = 17 // v17: include the OpenAI group Live gate
type apiKeyAuthCacheConfig struct {
l1Size int
@@ -409,6 +409,7 @@ func (s *APIKeyService) snapshotFromAPIKey(ctx context.Context, apiKey *APIKey)
MCPXMLInject: apiKey.Group.MCPXMLInject,
SupportedModelScopes: apiKey.Group.SupportedModelScopes,
AllowMessagesDispatch: apiKey.Group.AllowMessagesDispatch,
AllowLive: apiKey.Group.AllowLive,
DefaultMappedModel: apiKey.Group.DefaultMappedModel,
MessagesDispatchModelConfig: apiKey.Group.MessagesDispatchModelConfig,
ModelsListConfig: apiKey.Group.ModelsListConfig,
@@ -495,6 +496,7 @@ func (s *APIKeyService) snapshotToAPIKey(key string, snapshot *APIKeyAuthSnapsho
MCPXMLInject: snapshot.Group.MCPXMLInject,
SupportedModelScopes: snapshot.Group.SupportedModelScopes,
AllowMessagesDispatch: snapshot.Group.AllowMessagesDispatch,
AllowLive: snapshot.Group.AllowLive,
DefaultMappedModel: snapshot.Group.DefaultMappedModel,
MessagesDispatchModelConfig: snapshot.Group.MessagesDispatchModelConfig,
ModelsListConfig: snapshot.Group.ModelsListConfig,
@@ -132,7 +132,8 @@ func (s *AuthService) RegisterOAuthEmailAccount(
return nil, nil, err
}
existsEmail, err := s.userRepo.ExistsByEmail(ctx, email)
// 含 +别名 / Gmail 点号 / FQDN 根点变体归一化:该路径同样发放注册赠额,不能被单个收件箱刷号。
existsEmail, err := s.existsByEmailOrAlias(ctx, email)
if err != nil {
slog.Error("oauth email register: ExistsByEmail failed", "email", email, "error", err.Error())
return nil, nil, ErrServiceUnavailable
@@ -159,7 +160,7 @@ func (s *AuthService) RegisterOAuthEmailAccount(
SignupSource: signupSource,
}
if err := s.userRepo.Create(ctx, user); err != nil {
if err := s.userRepo.CreateWithEmailAliasGuard(ctx, user); err != nil {
if errors.Is(err, ErrEmailExists) {
return nil, nil, ErrEmailExists
}
@@ -211,7 +212,8 @@ func (s *AuthService) RegisterVerifiedOAuthEmailAccount(
return nil, nil, err
}
existsEmail, err := s.userRepo.ExistsByEmail(ctx, email)
// 与本地注册同口径:同一收件箱的别名变体不能各自建号(该路径也发放注册赠额)。
existsEmail, err := s.existsByEmailOrAlias(ctx, email)
if err != nil {
return nil, nil, ErrServiceUnavailable
}
@@ -241,7 +243,7 @@ func (s *AuthService) RegisterVerifiedOAuthEmailAccount(
SignupSource: signupSource,
}
if err := s.userRepo.Create(ctx, user); err != nil {
if err := s.userRepo.CreateWithEmailAliasGuard(ctx, user); err != nil {
if errors.Is(err, ErrEmailExists) {
return nil, nil, ErrEmailExists
}
+7 -7
View File
@@ -189,8 +189,8 @@ func (s *AuthService) RegisterWithVerification(ctx context.Context, email, passw
}
}
// 检查邮箱是否已存在
existsEmail, err := s.userRepo.ExistsByEmail(ctx, email)
// 检查邮箱是否已存在(含 +别名 / Gmail 点号变体归一化,防止单个收件箱批量派生注册)
existsEmail, err := s.existsByEmailOrAlias(ctx, email)
if err != nil {
logger.LegacyPrintf("service.auth", "[Auth] Database error checking email exists: %v", err)
return "", nil, ErrServiceUnavailable
@@ -224,7 +224,7 @@ func (s *AuthService) RegisterWithVerification(ctx context.Context, email, passw
Status: StatusActive,
}
if err := s.userRepo.Create(ctx, user); err != nil {
if err := s.userRepo.CreateWithEmailAliasGuard(ctx, user); err != nil {
// 优先检查邮箱冲突错误(竞态条件下可能发生)
if errors.Is(err, ErrEmailExists) {
return "", nil, ErrEmailExists
@@ -296,8 +296,8 @@ func (s *AuthService) SendVerifyCode(ctx context.Context, email string, locale .
return err
}
// 检查邮箱是否已存在
existsEmail, err := s.userRepo.ExistsByEmail(ctx, email)
// 检查邮箱是否已存在(含 +别名 / Gmail 点号变体归一化,防止单个收件箱批量派生注册)
existsEmail, err := s.existsByEmailOrAlias(ctx, email)
if err != nil {
logger.LegacyPrintf("service.auth", "[Auth] Database error checking email exists: %v", err)
return ErrServiceUnavailable
@@ -337,8 +337,8 @@ func (s *AuthService) SendVerifyCodeAsync(ctx context.Context, email string, loc
return nil, err
}
// 检查邮箱是否已存在
existsEmail, err := s.userRepo.ExistsByEmail(ctx, email)
// 检查邮箱是否已存在(含 +别名 / Gmail 点号变体归一化;在发信前拦截,避免批量脚本消耗发信配额)
existsEmail, err := s.existsByEmailOrAlias(ctx, email)
if err != nil {
logger.LegacyPrintf("service.auth", "[Auth] Database error checking email exists: %v", err)
return nil, ErrServiceUnavailable
@@ -874,6 +874,10 @@ func newEmailBindUserRepoStub(user *service.User) *emailBindUserRepoStub {
func (s *emailBindUserRepoStub) Create(context.Context, *service.User) error { return nil }
func (s *emailBindUserRepoStub) CreateWithEmailAliasGuard(ctx context.Context, user *service.User) error {
return s.Create(ctx, user)
}
func (s *emailBindUserRepoStub) GetByID(_ context.Context, id int64) (*service.User, error) {
s.mu.Lock()
defer s.mu.Unlock()
@@ -957,6 +961,18 @@ func (s *emailBindUserRepoStub) ExistsByEmail(_ context.Context, email string) (
return ok, nil
}
func (s *emailBindUserRepoStub) ExistsByEmailAlias(_ context.Context, email string) (bool, error) {
s.mu.Lock()
defer s.mu.Unlock()
identity := service.NormalizeEmailForAliasDedup(email)
for stored := range s.usersByEmail {
if service.NormalizeEmailForAliasDedup(stored) == identity {
return true, nil
}
}
return false, nil
}
func (s *emailBindUserRepoStub) BatchSetConcurrency(context.Context, []int64, int) (int, error) {
return 0, nil
}
@@ -365,6 +365,30 @@ func TestAuthService_Register_EmailExists(t *testing.T) {
require.ErrorIs(t, err, ErrEmailExists)
}
func TestAuthService_Register_AliasDuplicateRejected(t *testing.T) {
repo := &userRepoStub{aliasExists: true}
service := newAuthService(repo, map[string]string{
SettingKeyRegistrationEnabled: "true",
}, nil, nil)
_, _, err := service.Register(context.Background(), "some.one+bulk294@gmail.com", "password")
require.ErrorIs(t, err, ErrEmailExists)
require.Empty(t, repo.created)
}
func TestAuthService_Register_UsesAliasGuardedCreate(t *testing.T) {
// 注册必须走带别名兜底的创建路径:服务层前置查重与写入之间存在竞态窗口。
repo := &userRepoStub{nextID: 91}
service := newAuthService(repo, map[string]string{
SettingKeyRegistrationEnabled: "true",
}, nil, nil)
_, user, err := service.Register(context.Background(), "newuser@gmail.com", "password")
require.NoError(t, err)
require.NotNil(t, user)
require.Equal(t, 1, repo.guardedCreates)
}
func TestAuthService_Register_CheckEmailError(t *testing.T) {
repo := &userRepoStub{existsErr: errors.New("db down")}
service := newAuthService(repo, map[string]string{
+2
View File
@@ -137,6 +137,7 @@ type BatchImageJob struct {
IdempotencyKey *string
RequestHash *string
ManifestHash *string
SessionID *string
RetryCount int
Version int
@@ -196,6 +197,7 @@ type CreateBatchImageJobParams struct {
IdempotencyKey *string
RequestHash *string
ManifestHash *string
SessionID *string
RetryCount int
@@ -425,6 +425,7 @@ func (r *fakeBatchImageRepository) CreateBatchImageJob(_ context.Context, params
Currency: params.Currency,
IdempotencyKey: params.IdempotencyKey,
RequestHash: params.RequestHash,
SessionID: params.SessionID,
CreatedAt: time.Now(),
}
r.jobs[job.BatchID] = job
@@ -57,6 +57,7 @@ type BatchImageSubmitRequest struct {
AspectRatio string `json:"aspect_ratio"`
ImageSize string `json:"image_size"`
Metadata map[string]string `json:"metadata"`
SessionID *string `json:"-"`
}
type BatchImageSubmitItem struct {
@@ -281,6 +282,7 @@ func (s *BatchImagePublicService) Submit(ctx context.Context, owner BatchImageOw
HoldID: &holdID,
IdempotencyKey: batchImageOptionalStringPtr(idempotencyKey),
RequestHash: batchImageStringPtr(requestHash),
SessionID: normalized.SessionID,
})
if err != nil {
return nil, err
@@ -27,8 +27,10 @@ func TestBatchImagePublicService_Submit(t *testing.T) {
t.Run("accepts valid request stores refs and enqueues once", func(t *testing.T) {
svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true)
req := validBatchImageSubmitRequest()
req.SessionID = batchImageStringPtr("batch-session-123")
got, err := svc.Submit(ctx, testBatchImageOwner(), validBatchImageSubmitRequest(), "")
got, err := svc.Submit(ctx, testBatchImageOwner(), req, "")
require.NoError(t, err)
require.Equal(t, "image.batch", got.Object)
require.Equal(t, "queued", got.Status)
@@ -61,6 +63,7 @@ func TestBatchImagePublicService_Submit(t *testing.T) {
require.InDelta(t, 0.6, job.HoldMultiplier, 1e-12)
require.InDelta(t, 0.125, job.BillableUnitPrice, 1e-12)
require.InDelta(t, 0.15, job.HoldUnitPrice, 1e-12)
require.Equal(t, "batch-session-123", batchImageDerefString(job.SessionID))
})
t.Run("combines user group image rate account rate discount and hold margin", func(t *testing.T) {
@@ -389,15 +392,18 @@ func TestBatchImagePublicService_Submit(t *testing.T) {
})
t.Run("idempotency returns same batch without provider resubmit", func(t *testing.T) {
svc, _, queue, gemini, _ := newTestBatchImagePublicService(true)
svc, repo, queue, gemini, _ := newTestBatchImagePublicService(true)
req := validBatchImageSubmitRequest()
req.SessionID = batchImageStringPtr("original-session")
first, err := svc.Submit(ctx, testBatchImageOwner(), req, "client-key")
require.NoError(t, err)
req.SessionID = batchImageStringPtr("retry-session")
second, err := svc.Submit(ctx, testBatchImageOwner(), req, "client-key")
require.NoError(t, err)
require.Equal(t, first.ID, second.ID)
require.Equal(t, "original-session", batchImageDerefString(repo.jobs[first.ID].SessionID))
require.Len(t, gemini.submits, 1)
require.Equal(t, []string{first.ID}, queue.enqueued)
})
@@ -277,6 +277,7 @@ func (s *BatchImageSettlementService) recordUsageLog(ctx context.Context, job *B
RequestType: RequestTypeSync,
BillingMode: &billingMode,
ImageSize: &imageSize,
SessionID: job.SessionID,
CreatedAt: createdAt,
}
writeUsageLogBestEffort(ctx, s.UsageLogRepo, usageLog, "service.batch_image_settlement")
@@ -18,9 +18,14 @@ func TestBatchImageSettlementService_SettlesAndChargesSuccessfulImagesOnly(t *te
job.SuccessCount = 3
job.FailCount = 2
job.ItemCount = 5
job.SessionID = batchImageStringPtr("batch-settlement-session")
repo.jobs[job.BatchID] = job
billing := &fakeBatchImageBillingRepo{}
svc := &BatchImageSettlementService{Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25}}
usageLogs := &openAIRecordUsageLogRepoStub{}
svc := &BatchImageSettlementService{
Repo: repo, BillingRepo: billing, Pricing: &fakeBatchImagePricingResolver{unitPrice: 0.25},
UsageLogRepo: usageLogs,
}
result, err := svc.Settle(context.Background(), job.BatchID)
require.NoError(t, err)
@@ -32,6 +37,7 @@ func TestBatchImageSettlementService_SettlesAndChargesSuccessfulImagesOnly(t *te
require.Equal(t, 0.75, *repo.jobs[job.BatchID].ActualCost)
require.NotEmpty(t, batchImageDerefString(repo.jobs[job.BatchID].ManifestHash))
require.NotNil(t, repo.jobs[job.BatchID].SettledAt)
require.Equal(t, "batch-settlement-session", batchImageDerefString(usageLogs.lastLog.SessionID))
require.Len(t, billing.captures, 1)
require.Equal(t, int64(321), billing.captures[0].APIKeyID)
require.Equal(t, job.UserID, billing.captures[0].UserID)
+7 -2
View File
@@ -347,8 +347,13 @@ func removeCustomFieldFromTools(body []byte) []byte {
}
// claudeVersionRe 匹配 Claude 模型 ID 中的版本号部分
// 支持 claude-{tier}-{major}-{minor} 和 claude-{tier}-{major}.{minor} 格式
var claudeVersionRe = regexp.MustCompile(`claude-(?:haiku|sonnet|opus)-(\d+)[-.](\d+)`)
// 支持 claude-{tier}-{major}-{minor} 和 claude-{tier}-{major}.{minor} 格式;
// minor 可选,用于覆盖只有主版本号的新模型 ID(claude-opus-5 / claude-sonnet-5),
// 此时 matches[2] 为空串,调用方 strconv.Atoi 得到 minor=0。
// 缺少可选 minor 时这类 ID 完全不匹配,会被当作"旧模型"降级处理:
// thinking.enabled 不转 adaptive(Opus 5 上游已移除 budget_tokens,直接 400)、
// cache_control.ttl 被剥离、tool search 被过滤。
var claudeVersionRe = regexp.MustCompile(`claude-(?:haiku|sonnet|opus)-(\d+)(?:[-.](\d+))?`)
// isBedrockClaude45OrNewer 判断 Bedrock 模型 ID 是否为 Claude 4.5 或更新版本
// Claude 4.5+ 支持 cache_control 中的 ttl 字段("5m" 和 "1h")
@@ -258,6 +258,11 @@ func (s *BillingService) initFallbackPricing() {
// Claude 4.7 Opus (暂与4.6同价,待官方定价更新)
s.fallbackPrices["claude-opus-4.7"] = s.fallbackPrices["claude-opus-4.6"]
// Claude 4.8 Opus / Claude Opus 5(官方同价:$5 输入 / $25 输出 per MTok)。
// 缺少这两条时 getFallbackPricing 会掉到 claude-3-opus($15/$75),造成 3 倍超收。
s.fallbackPrices["claude-opus-4.8"] = s.fallbackPrices["claude-opus-4.7"]
s.fallbackPrices["claude-opus-5"] = s.fallbackPrices["claude-opus-4.8"]
// Gemini 3.1 Pro
s.fallbackPrices["gemini-3.1-pro"] = &ModelPricing{
InputPricePerToken: 2e-6, // $2 per MTok
@@ -588,6 +593,13 @@ func (s *BillingService) getFallbackPricing(model string) *ModelPricing {
// 按模型系列匹配
if strings.Contains(modelLower, "opus") {
// "opus-5" 必须先判:不能用裸 "5" 匹配,否则 claude-opus-4-5 会被误判。
if strings.Contains(modelLower, "opus-5") || strings.Contains(modelLower, "opus5") {
return s.fallbackPrices["claude-opus-5"]
}
if strings.Contains(modelLower, "4.8") || strings.Contains(modelLower, "4-8") {
return s.fallbackPrices["claude-opus-4.8"]
}
if strings.Contains(modelLower, "4.7") || strings.Contains(modelLower, "4-7") {
return s.fallbackPrices["claude-opus-4.7"]
}
@@ -0,0 +1,139 @@
package service
import (
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/domain"
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// Claude Opus 5 官方定价(USD per token):$5 输入 / $25 输出 per MTok。
const (
opus5InputPricePerToken = 5e-6
opus5OutputPricePerToken = 25e-6
opus5CacheCreationPricePerToken = 6.25e-6
opus5CacheReadPricePerToken = 0.5e-6
)
// TestClaudeOpus5_FamilyFallbackDoesNotUseOpus4Rates 覆盖定价数据里还没有
// claude-opus-5 条目的场景(远端价格表滞后于模型发布)。
// 修复前 matchByModelFamily 会把 claude-opus-5 归到 "opus-4" 系列,
// 按 $15/$75 计费 —— 输入与输出双双 3 倍超收。
func TestClaudeOpus5_FamilyFallbackDoesNotUseOpus4Rates(t *testing.T) {
svc := NewBillingService(&config.Config{}, &PricingService{
pricingData: map[string]*LiteLLMModelPricing{
// 有 4.8(同价),故意不放 claude-opus-5
"claude-opus-4-8": {
InputCostPerToken: opus5InputPricePerToken,
OutputCostPerToken: opus5OutputPricePerToken,
CacheCreationInputTokenCost: opus5CacheCreationPricePerToken,
CacheReadInputTokenCost: opus5CacheReadPricePerToken,
},
// 修复前会被误命中的旧 Opus 条目($15/$75)
"claude-opus-4-1": {InputCostPerToken: 15e-6, OutputCostPerToken: 75e-6},
"claude-opus-4-20250514": {InputCostPerToken: 15e-6, OutputCostPerToken: 75e-6},
"claude-3-opus-20240229": {InputCostPerToken: 15e-6, OutputCostPerToken: 75e-6},
},
})
for _, model := range []string{"claude-opus-5", "us.anthropic.claude-opus-5-v1"} {
t.Run(model, func(t *testing.T) {
pricing, err := svc.GetModelPricing(model)
require.NoError(t, err)
require.NotNil(t, pricing)
assert.InDelta(t, opus5InputPricePerToken, pricing.InputPricePerToken, 1e-12)
assert.InDelta(t, opus5OutputPricePerToken, pricing.OutputPricePerToken, 1e-12)
})
}
}
// TestClaudeOpus5_HardcodedFallbackPricing 覆盖动态价格服务完全不可用时的
// 硬编码兜底表。同时锁定不能被 "opus-5" 子串误伤的相邻型号。
func TestClaudeOpus5_HardcodedFallbackPricing(t *testing.T) {
// pricingService 为 nil,强制走硬编码兜底表
svc := NewBillingService(&config.Config{}, nil)
tests := []struct {
model string
input float64
output float64
}{
{"claude-opus-5", opus5InputPricePerToken, opus5OutputPricePerToken},
{"us.anthropic.claude-opus-5-v1", opus5InputPricePerToken, opus5OutputPricePerToken},
// 4.8 与 5 同价;修复前兜底表缺失,会掉到 claude-3-opus 的 $15/$75
{"claude-opus-4-8", opus5InputPricePerToken, opus5OutputPricePerToken},
// 相邻型号不能被 "opus-5" 误匹配
{"claude-opus-4-5-20251101", 5e-6, 25e-6},
{"claude-opus-4-1-20250805", 15e-6, 75e-6},
{"claude-3-opus-20240229", 15e-6, 75e-6},
}
for _, tt := range tests {
t.Run(tt.model, func(t *testing.T) {
pricing, err := svc.GetModelPricing(tt.model)
require.NoError(t, err)
require.NotNil(t, pricing)
assert.InDelta(t, tt.input, pricing.InputPricePerToken, 1e-12)
assert.InDelta(t, tt.output, pricing.OutputPricePerToken, 1e-12)
})
}
opus5, err := svc.GetModelPricing("claude-opus-5")
require.NoError(t, err)
assert.InDelta(t, opus5CacheCreationPricePerToken, opus5.CacheCreationPricePerToken, 1e-12)
assert.InDelta(t, opus5CacheReadPricePerToken, opus5.CacheReadPricePerToken, 1e-12)
}
// TestClaudeOpus5_BedrockCapabilityGates 锁定只有主版本号的模型 ID
// (claude-opus-5 / claude-sonnet-5)能被版本闸门识别。
// 修复前 claudeVersionRe 强制要求 major-minor,这类 ID 完全不匹配,
// 会被当成旧模型降级。
func TestClaudeOpus5_BedrockCapabilityGates(t *testing.T) {
tests := []struct {
modelID string
claude45Newer bool
toolSearch bool
opus47Newer bool
}{
{"claude-opus-5", true, true, true},
{"us.anthropic.claude-opus-5-v1", true, true, true},
{"eu.anthropic.claude-opus-5-v1", true, true, true},
{"claude-sonnet-5", true, true, false},
{"us.anthropic.claude-sonnet-5-v1", true, true, false},
// 回归保护:旧模型不能因为 minor 可选而被误判为新版本
{"anthropic.claude-opus-4-1-v1", false, false, false},
{"anthropic.claude-sonnet-4-0-v1", false, false, false},
{"anthropic.claude-3-opus-20240229-v1:0", false, false, false},
{"us.anthropic.claude-opus-4-8-v1", true, true, true},
// Haiku 不支持 tool search
{"us.anthropic.claude-haiku-4-5-20251001-v1:0", true, false, false},
}
for _, tt := range tests {
t.Run(tt.modelID, func(t *testing.T) {
assert.Equal(t, tt.claude45Newer, isBedrockClaude45OrNewer(tt.modelID), "isBedrockClaude45OrNewer")
assert.Equal(t, tt.toolSearch, bedrockModelSupportsToolSearch(tt.modelID), "bedrockModelSupportsToolSearch")
assert.Equal(t, tt.opus47Newer, isBedrockOpus47OrNewer(tt.modelID), "isBedrockOpus47OrNewer")
})
}
}
// TestClaudeOpus5_BedrockThinkingConvertedToAdaptive 验证 Opus 5 在 Bedrock 路径上
// 会把 thinking.type=enabled 转成 adaptive 并移除 budget_tokens。
// 上游 Opus 5 已移除 budget_tokens,透传过去会直接 400。
func TestClaudeOpus5_BedrockThinkingConvertedToAdaptive(t *testing.T) {
body := []byte(`{"thinking":{"type":"enabled","budget_tokens":10000}}`)
got := sanitizeBedrockThinking(body, "us.anthropic.claude-opus-5-v1")
assert.JSONEq(t, `{"thinking":{"type":"adaptive"}}`, string(got))
}
// TestClaudeOpus5_CatalogAndBedrockMapping 锁定模型清单与 Bedrock 默认映射。
func TestClaudeOpus5_CatalogAndBedrockMapping(t *testing.T) {
assert.Contains(t, claude.DefaultModelIDs(), "claude-opus-5")
mapped, ok := domain.DefaultBedrockModelMapping["claude-opus-5"]
require.True(t, ok, "claude-opus-5 missing from DefaultBedrockModelMapping")
assert.Equal(t, "us.anthropic.claude-opus-5-v1", mapped)
}
@@ -168,6 +168,10 @@ func (r *contentModerationTestUserRepo) Create(ctx context.Context, user *User)
panic("unexpected Create call")
}
func (r *contentModerationTestUserRepo) CreateWithEmailAliasGuard(ctx context.Context, user *User) error {
panic("unexpected CreateWithEmailAliasGuard call")
}
func (r *contentModerationTestUserRepo) GetByID(ctx context.Context, id int64) (*User, error) {
if r.user == nil {
return nil, ErrUserNotFound
@@ -257,6 +261,10 @@ func (r *contentModerationTestUserRepo) ExistsByEmail(ctx context.Context, email
panic("unexpected ExistsByEmail call")
}
func (r *contentModerationTestUserRepo) ExistsByEmailAlias(ctx context.Context, email string) (bool, error) {
panic("unexpected ExistsByEmailAlias call")
}
func (r *contentModerationTestUserRepo) RemoveGroupFromAllowedGroups(ctx context.Context, groupID int64) (int64, error) {
panic("unexpected RemoveGroupFromAllowedGroups call")
}
@@ -1623,3 +1623,114 @@ func TestGatewayService_AnthropicAPIKeyPassthrough_StreamingUpstreamReadErrorAft
require.True(t, result.clientDisconnect)
require.Equal(t, 8, result.usage.InputTokens)
}
func TestGatewayService_AnthropicAPIKeyPassthrough_TransportErrorRecordsOllamaActivity(t *testing.T) {
gin.SetMode(gin.TestMode)
deferred := NewDeferredService(nil, nil, time.Second)
upstream := &anthropicHTTPUpstreamRecorder{err: errors.New("dial tcp timeout")}
svc := &GatewayService{
cfg: &config.Config{
Security: config.SecurityConfig{
URLAllowlist: config.URLAllowlistConfig{Enabled: false},
},
},
httpUpstream: upstream,
deferredService: deferred,
}
ollama := &Account{
ID: 601, Name: "ollama-anthropic", Platform: PlatformAnthropic, Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{"api_key": "k-ollama", "base_url": "https://ollama.com"},
Extra: map[string]any{"anthropic_passthrough": true},
Status: StatusActive, Schedulable: true,
}
other := newAnthropicAPIKeyAccountForTest()
other.ID = 602
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
_, err := svc.forwardAnthropicAPIKeyPassthrough(context.Background(), c, ollama, []byte(`{"model":"x"}`), "x", "x", false, time.Now())
require.Error(t, err)
rec2 := httptest.NewRecorder()
c2, _ := gin.CreateTestContext(rec2)
c2.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
_, err = svc.forwardAnthropicAPIKeyPassthrough(context.Background(), c2, other, []byte(`{"model":"x"}`), "x", "x", false, time.Now())
require.Error(t, err)
_, ok := deferred.lastUsedUpdates.Load(int64(601))
require.True(t, ok, "Anthropic passthrough transport error on Ollama account must record activity")
_, ok = deferred.lastUsedUpdates.Load(int64(602))
require.False(t, ok, "non-Ollama Anthropic passthrough transport error must not record Ollama activity")
}
func TestGatewayService_AnthropicAPIKeyPassthrough_ContextCanceledSkipsOllamaActivity(t *testing.T) {
gin.SetMode(gin.TestMode)
deferred := NewDeferredService(nil, nil, time.Second)
upstream := &anthropicHTTPUpstreamRecorder{err: context.Canceled}
svc := &GatewayService{
cfg: &config.Config{
Security: config.SecurityConfig{
URLAllowlist: config.URLAllowlistConfig{Enabled: false},
},
},
httpUpstream: upstream,
deferredService: deferred,
}
ollama := &Account{
ID: 603, Name: "ollama-canceled", Platform: PlatformAnthropic, Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{"api_key": "k-ollama", "base_url": "https://ollama.com"},
Extra: map[string]any{"anthropic_passthrough": true},
Status: StatusActive, Schedulable: true,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
_, err := svc.forwardAnthropicAPIKeyPassthrough(context.Background(), c, ollama, []byte(`{"model":"x"}`), "x", "x", false, time.Now())
require.Error(t, err)
_, ok := deferred.lastUsedUpdates.Load(int64(603))
require.False(t, ok, "context.Canceled on Anthropic passthrough must not count as Ollama activity")
}
func TestGatewayService_AnthropicAPIKeyPassthrough_Non2xxRecordsOllamaActivity(t *testing.T) {
gin.SetMode(gin.TestMode)
deferred := NewDeferredService(nil, nil, time.Second)
// 400 is non-retryable / non-failover for default API-key accounts, so it reaches handleErrorResponse.
upstream := &anthropicHTTPUpstreamRecorder{
resp: &http.Response{
StatusCode: http.StatusBadRequest,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"type":"error","error":{"type":"invalid_request_error","message":"bad"}}`)),
},
}
svc := &GatewayService{
cfg: &config.Config{
Security: config.SecurityConfig{
URLAllowlist: config.URLAllowlistConfig{Enabled: false},
},
},
httpUpstream: upstream,
deferredService: deferred,
rateLimitService: &RateLimitService{},
}
ollama := &Account{
ID: 604, Name: "ollama-400", Platform: PlatformAnthropic, Type: AccountTypeAPIKey,
Concurrency: 1,
Credentials: map[string]any{"api_key": "k-ollama", "base_url": "https://ollama.com"},
Extra: map[string]any{"anthropic_passthrough": true},
Status: StatusActive, Schedulable: true,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
_, _ = svc.forwardAnthropicAPIKeyPassthrough(context.Background(), c, ollama, []byte(`{"model":"x"}`), "x", "x", false, time.Now())
_, ok := deferred.lastUsedUpdates.Load(int64(604))
require.True(t, ok, "Anthropic passthrough non-2xx on Ollama account must record activity via handleErrorResponse")
}
@@ -114,6 +114,9 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
if resp != nil && resp.Body != nil {
_ = resp.Body.Close()
}
if !errors.Is(err, context.Canceled) {
scheduleOllamaCloudUsageActivity(s.deferredService, account)
}
safeErr := sanitizeUpstreamErrorMessage(err.Error())
setOpsUpstreamError(c, 0, safeErr, "")
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
@@ -369,6 +369,10 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
if resp != nil && resp.Body != nil {
_ = resp.Body.Close()
}
// Transport attempt left local validation; count Ollama Cloud activity.
if !errors.Is(err, context.Canceled) {
scheduleOllamaCloudUsageActivity(s.deferredService, account)
}
// Ensure the client receives an error response (handlers assume Forward writes on non-failover errors).
safeErr := sanitizeUpstreamErrorMessage(err.Error())
setOpsUpstreamError(c, 0, safeErr, "")
@@ -356,6 +356,8 @@ func (s *GatewayService) readUpstreamErrorBody(resp *http.Response) ([]byte, err
}
func (s *GatewayService) handleErrorResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, requestedModel ...string) (*ForwardResult, error) {
// Upstream returned a non-success HTTP status; count Ollama Cloud activity.
scheduleOllamaCloudUsageActivity(s.deferredService, account)
body, readErr := s.readUpstreamErrorBody(resp)
if readErr != nil {
// 读取失败时 body 可能被截断,错误分类会基于不完整数据;记录日志以便排查,
@@ -46,6 +46,7 @@ type RecordUsageInput struct {
UpstreamEndpoint string // 上游端点(标准化后的上游路径)
UserAgent string // 请求的 User-Agent
IPAddress string // 请求的客户端 IP 地址
SessionID string // 客户端显式会话标识(session_id / X-Session-Id 等请求头),仅用于用量行会话关联
RequestPayloadHash string // 请求体语义哈希,用于降低 request_id 误复用时的静默误去重风险
ForceCacheBilling bool // 强制缓存计费:将 input_tokens 转为 cache_read 计费(用于粘性会话切换)
APIKeyService APIKeyQuotaUpdater // 可选:用于更新API Key配额
@@ -572,6 +573,7 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu
UpstreamEndpoint: input.UpstreamEndpoint,
UserAgent: input.UserAgent,
IPAddress: input.IPAddress,
SessionID: input.SessionID,
RequestPayloadHash: input.RequestPayloadHash,
ForceCacheBilling: input.ForceCacheBilling,
APIKeyService: input.APIKeyService,
@@ -591,6 +593,7 @@ type RecordUsageLongContextInput struct {
UpstreamEndpoint string // 上游端点(标准化后的上游路径)
UserAgent string // 请求的 User-Agent
IPAddress string // 请求的客户端 IP 地址
SessionID string // 客户端显式会话标识(session_id / X-Session-Id 等请求头),仅用于用量行会话关联
RequestPayloadHash string // 请求体语义哈希,用于降低 request_id 误复用时的静默误去重风险
LongContextThreshold int // 长上下文阈值(如 200000)
LongContextMultiplier float64 // 超出阈值部分的倍率(如 2.0)
@@ -613,6 +616,7 @@ func (s *GatewayService) RecordUsageWithLongContext(ctx context.Context, input *
UpstreamEndpoint: input.UpstreamEndpoint,
UserAgent: input.UserAgent,
IPAddress: input.IPAddress,
SessionID: input.SessionID,
RequestPayloadHash: input.RequestPayloadHash,
ForceCacheBilling: input.ForceCacheBilling,
APIKeyService: input.APIKeyService,
@@ -635,6 +639,7 @@ type recordUsageCoreInput struct {
UpstreamEndpoint string
UserAgent string
IPAddress string
SessionID string
RequestPayloadHash string
ForceCacheBilling bool
APIKeyService APIKeyQuotaUpdater
@@ -1010,6 +1015,7 @@ func (s *GatewayService) buildRecordUsageLog(
ModelMappingChain: optionalTrimmedStringPtr(input.ModelMappingChain),
UserAgent: optionalTrimmedStringPtr(input.UserAgent),
IPAddress: optionalTrimmedStringPtr(input.IPAddress),
SessionID: optionalTrimmedStringPtr(input.SessionID),
GroupID: apiKey.GroupID,
SubscriptionID: optionalSubscriptionID(subscription),
CreatedAt: time.Now(),
@@ -473,7 +473,7 @@ func geminiResponseToChatCompletions(
rawData []byte,
usageOverride *ClaudeUsage,
) (*apicompat.ChatCompletionsResponse, *ClaudeUsage, error) {
claudeRespMap, usage := convertGeminiToClaudeMessage(geminiResp, originalModel, rawData)
claudeRespMap, usage := convertGeminiToClaudeMessage(geminiResp, originalModel, rawData, true)
if usageOverride != nil && (usageOverride.InputTokens > 0 || usageOverride.OutputTokens > 0 || usageOverride.CacheReadInputTokens > 0) {
usage = usageOverride
if usageMap, ok := claudeRespMap["usage"].(map[string]any); ok {
@@ -0,0 +1,163 @@
package service
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/require"
)
func TestGeminiResponseToChatCompletionsPreservesInlineData(t *testing.T) {
tests := []struct {
name string
parts []any
want string
}{
{
name: "image only",
parts: []any{
map[string]any{"inlineData": map[string]any{"mimeType": "image/png", "data": "aW1hZ2U="}},
},
want: "![image](data:image/png;base64,aW1hZ2U=)",
},
{
name: "text and image",
parts: []any{
map[string]any{"text": "rendered image:\n"},
map[string]any{"inlineData": map[string]any{"mimeType": "image/webp", "data": "d2VicA=="}},
},
want: "rendered image:\n![image](data:image/webp;base64,d2VicA==)",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
geminiResp := map[string]any{
"candidates": []any{map[string]any{
"content": map[string]any{"parts": tt.parts},
"finishReason": "STOP",
}},
}
rawData, err := json.Marshal(geminiResp)
require.NoError(t, err)
got, _, err := geminiResponseToChatCompletions(geminiResp, "gemini-test", rawData, nil)
require.NoError(t, err)
require.Len(t, got.Choices, 1)
var content string
require.NoError(t, json.Unmarshal(got.Choices[0].Message.Content, &content))
require.Equal(t, tt.want, content)
require.Equal(t, "stop", got.Choices[0].FinishReason)
})
}
}
func TestGeminiResponseToChatCompletionsOmitsInvalidInlineData(t *testing.T) {
tests := []struct {
name string
inlineData map[string]any
}{
{
name: "unsupported MIME type",
inlineData: map[string]any{"mimeType": "image/svg+xml", "data": "PHN2Zz48L3N2Zz4="},
},
{
name: "malformed MIME type",
inlineData: map[string]any{"mimeType": "image/png; charset=utf-8", "data": "aW1hZ2U="},
},
{
name: "malformed base64",
inlineData: map[string]any{"mimeType": "image/png", "data": "not-valid-base64!!!"},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
geminiResp := map[string]any{
"candidates": []any{map[string]any{
"content": map[string]any{"parts": []any{map[string]any{"text": "before"}, map[string]any{"inlineData": tt.inlineData}, map[string]any{"text": "after"}}},
"finishReason": "STOP",
}},
}
rawData, err := json.Marshal(geminiResp)
require.NoError(t, err)
got, _, err := geminiResponseToChatCompletions(geminiResp, "gemini-test", rawData, nil)
require.NoError(t, err)
var content string
require.NoError(t, json.Unmarshal(got.Choices[0].Message.Content, &content))
require.Equal(t, "beforeafter", content)
})
}
}
func TestConvertGeminiToClaudeMessageOmitsInlineDataForAnthropicMessages(t *testing.T) {
geminiResp := map[string]any{
"candidates": []any{map[string]any{
"content": map[string]any{"parts": []any{
map[string]any{"text": "before"},
map[string]any{"inlineData": map[string]any{"mimeType": "image/png", "data": "aW1hZ2U="}},
map[string]any{"functionCall": map[string]any{"name": "get_weather", "args": map[string]any{"city": "Paris"}}},
map[string]any{"text": "after"},
}},
"finishReason": "STOP",
}},
}
rawData, err := json.Marshal(geminiResp)
require.NoError(t, err)
withInlineData, _ := convertGeminiToClaudeMessage(geminiResp, "gemini-test", rawData, true)
contentWithInlineData, ok := withInlineData["content"].([]any)
require.True(t, ok)
require.Len(t, contentWithInlineData, 4)
require.Equal(t, map[string]any{"type": "text", "text": "before"}, contentWithInlineData[0])
require.Equal(t, map[string]any{"type": "text", "text": "![image](data:image/png;base64,aW1hZ2U=)"}, contentWithInlineData[1])
toolUse, ok := contentWithInlineData[2].(map[string]any)
require.True(t, ok)
require.Equal(t, "tool_use", toolUse["type"])
require.Equal(t, "get_weather", toolUse["name"])
require.Equal(t, map[string]any{"type": "text", "text": "after"}, contentWithInlineData[3])
withoutInlineData, _ := convertGeminiToClaudeMessage(geminiResp, "gemini-test", rawData, false)
contentWithoutInlineData, ok := withoutInlineData["content"].([]any)
require.True(t, ok)
require.Len(t, contentWithoutInlineData, 3)
require.Equal(t, map[string]any{"type": "text", "text": "before"}, contentWithoutInlineData[0])
toolUseWithoutInlineData, ok := contentWithoutInlineData[1].(map[string]any)
require.True(t, ok)
require.Equal(t, "tool_use", toolUseWithoutInlineData["type"])
require.Equal(t, "get_weather", toolUseWithoutInlineData["name"])
require.Equal(t, map[string]any{"type": "text", "text": "after"}, contentWithoutInlineData[2])
}
func TestGeminiResponseToChatCompletionsRetainsTextAndToolBehavior(t *testing.T) {
geminiResp := map[string]any{
"candidates": []any{map[string]any{
"content": map[string]any{"parts": []any{
map[string]any{"text": "checking"},
map[string]any{"functionCall": map[string]any{
"name": "get_weather",
"args": map[string]any{"city": "Paris"},
}},
}},
"finishReason": "STOP",
}},
}
rawData, err := json.Marshal(geminiResp)
require.NoError(t, err)
got, _, err := geminiResponseToChatCompletions(geminiResp, "gemini-test", rawData, nil)
require.NoError(t, err)
require.Len(t, got.Choices, 1)
choice := got.Choices[0]
var content string
require.NoError(t, json.Unmarshal(choice.Message.Content, &content))
require.Equal(t, "checking", content)
require.Equal(t, "tool_calls", choice.FinishReason)
require.Len(t, choice.Message.ToolCalls, 1)
require.Equal(t, "get_weather", choice.Message.ToolCalls[0].Function.Name)
require.JSONEq(t, `{"city":"Paris"}`, choice.Message.ToolCalls[0].Function.Arguments)
}
@@ -5,6 +5,7 @@ import (
"bytes"
"context"
"crypto/rand"
"encoding/base64"
"encoding/hex"
"encoding/json"
"errors"
@@ -1068,7 +1069,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to read upstream stream")
}
collectedBytes, _ := json.Marshal(collected)
claudeResp, usageObj2 := convertGeminiToClaudeMessage(collected, originalModel, collectedBytes)
claudeResp, usageObj2 := convertGeminiToClaudeMessage(collected, originalModel, collectedBytes, false)
c.JSON(http.StatusOK, claudeResp)
usage = usageObj2
if usageObj != nil && (usageObj.InputTokens > 0 || usageObj.OutputTokens > 0) {
@@ -1965,7 +1966,7 @@ func (s *GeminiMessagesCompatService) handleNonStreamingResponse(c *gin.Context,
return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response")
}
claudeResp, usage := convertGeminiToClaudeMessage(geminiResp, originalModel, unwrappedBody)
claudeResp, usage := convertGeminiToClaudeMessage(geminiResp, originalModel, unwrappedBody, false)
c.JSON(http.StatusOK, claudeResp)
return usage, nil
@@ -2717,7 +2718,7 @@ func unwrapGeminiResponse(raw []byte) ([]byte, error) {
return raw, nil
}
func convertGeminiToClaudeMessage(geminiResp map[string]any, originalModel string, rawData []byte) (map[string]any, *ClaudeUsage) {
func convertGeminiToClaudeMessage(geminiResp map[string]any, originalModel string, rawData []byte, includeInlineData bool) (map[string]any, *ClaudeUsage) {
usage := extractGeminiUsage(rawData)
if usage == nil {
usage = &ClaudeUsage{}
@@ -2740,6 +2741,16 @@ func convertGeminiToClaudeMessage(geminiResp map[string]any, originalModel strin
"text": text,
})
}
if inlineData, ok := pm["inlineData"].(map[string]any); includeInlineData && ok {
mimeType, _ := inlineData["mimeType"].(string)
data, _ := inlineData["data"].(string)
if isGeminiInlineImageMIMEType(mimeType) && isValidBase64(data) {
contentBlocks = append(contentBlocks, map[string]any{
"type": "text",
"text": fmt.Sprintf("![image](data:%s;base64,%s)", mimeType, data),
})
}
}
if fc, ok := pm["functionCall"].(map[string]any); ok {
name, _ := fc["name"].(string)
if strings.TrimSpace(name) == "" {
@@ -2782,6 +2793,23 @@ func convertGeminiToClaudeMessage(geminiResp map[string]any, originalModel strin
return resp, usage
}
func isGeminiInlineImageMIMEType(mimeType string) bool {
switch mimeType {
case "image/gif", "image/jpeg", "image/png", "image/webp":
return true
default:
return false
}
}
func isValidBase64(data string) bool {
if data == "" {
return false
}
_, err := base64.StdEncoding.DecodeString(data)
return err == nil
}
func extractGeminiUsage(data []byte) *ClaudeUsage {
usage := gjson.GetBytes(data, "usageMetadata")
if !usage.Exists() {
+1
View File
@@ -82,6 +82,7 @@ type Group struct {
// OpenAI Messages 调度配置(仅 openai 平台使用)
AllowMessagesDispatch bool
AllowLive bool
RequireOAuthOnly bool // 仅允许非 apikey 类型账号关联(OpenAI/Antigravity/Anthropic/Gemini)
RequirePrivacySet bool // 调度时仅允许 privacy 已成功设置的账号(OpenAI/Antigravity/Anthropic/Gemini)
DefaultMappedModel string
+198 -10
View File
@@ -32,10 +32,20 @@ const (
OllamaCloudUsageAutoRefreshExtraKey = "ollama_cloud_usage_auto_refresh"
OllamaCloudUsageSnapshotExtraKey = "ollama_cloud_usage_snapshot"
// OllamaCloudUsageMinFetchInterval is the hard floor between two successful
// fetches of the same group, mirroring the floor nextOllamaCloudUsageDelay
// applies to next_refresh_at. Activity may bring a refresh forward to this
// bound but never past it. Exported so the repository can apply the same
// floor inside the SQL due filter.
OllamaCloudUsageMinFetchInterval = ollamaCloudUsageMinIntervalMinutes * time.Minute
ollamaCloudUsageSettingsURL = "https://ollama.com/settings"
ollamaCloudUsageDefaultIntervalMinutes = 60
ollamaCloudUsageMinIntervalMinutes = 15
ollamaCloudUsageMaxIntervalMinutes = 24 * 60
ollamaCloudUsageDefaultDebounceMinutes = 1
ollamaCloudUsageMinDebounceMinutes = 1
ollamaCloudUsageMaxDebounceMinutes = 60
ollamaCloudUsageCycleInterval = time.Minute
ollamaCloudUsageManualRefreshInterval = 30 * time.Second
ollamaCloudUsageRequestTimeout = 15 * time.Second
@@ -76,10 +86,15 @@ const (
OllamaCloudUsageStatusFailed = "failed"
)
// OllamaCloudUsageSettings controls the opt-in periodic refresh runner.
// OllamaCloudUsageSettings controls the opt-in request-driven refresh runner.
//
// IntervalMinutes is the max-wait bound: when model requests keep arriving and
// the trailing debounce keeps sliding, a refresh is forced after this long.
// DebounceMinutes is the quiet period after the latest request in a group.
type OllamaCloudUsageSettings struct {
Enabled bool `json:"enabled"`
IntervalMinutes int `json:"interval_minutes"`
IntervalMinutes int `json:"interval_minutes"` // max wait while requests continue
DebounceMinutes int `json:"debounce_minutes"` // trailing quiet period after last request
}
// OllamaCloudUsageWindow is a narrow, sanitized view of one official usage window.
@@ -114,6 +129,12 @@ type OllamaCloudUsageData struct {
}
// OllamaCloudUsageSnapshot is the only usage observation persisted in account extra.
//
// NextRefreshAt remains a persisted compatibility field. For status=ok it is a
// max-wait horizon marker only; automatic success refreshes are driven by model
// request activity (group last_used_at + debounce/max-wait), not by this field
// alone. For failed/unauthorized snapshots it is the failure not-before time
// (Retry-After / exponential backoff) and is enforced as max(activityDue, NextRefreshAt).
type OllamaCloudUsageSnapshot struct {
Status string `json:"status"`
Data *OllamaCloudUsageData `json:"data,omitempty"`
@@ -142,7 +163,7 @@ type ollamaCloudUsageRepository interface {
SetOllamaCloudUsageAutoRefresh(context.Context, *Account, bool) error
UpdateOllamaCloudUsageSnapshot(context.Context, *Account, *OllamaCloudUsageSnapshot) error
DisableOllamaCloudUsageAutoRefresh(context.Context, *Account) error
ListDueOllamaCloudUsageAccounts(context.Context, time.Time, int) ([]Account, error)
ListDueOllamaCloudUsageAccounts(context.Context, time.Time, time.Duration, time.Duration, int) ([]Account, error)
}
// GetOllamaCloudUsageSettings returns fail-safe defaults when the setting is absent.
@@ -168,6 +189,9 @@ func (s *SettingService) GetOllamaCloudUsageSettings(ctx context.Context) (*Olla
if settings.IntervalMinutes == 0 {
settings.IntervalMinutes = defaults.IntervalMinutes
}
if settings.DebounceMinutes == 0 {
settings.DebounceMinutes = defaults.DebounceMinutes
}
normalizeOllamaCloudUsageSettings(&settings)
return &settings, nil
}
@@ -179,12 +203,31 @@ func (s *SettingService) SetOllamaCloudUsageSettings(ctx context.Context, settin
if settings == nil {
return infraerrors.BadRequest("INVALID_OLLAMA_CLOUD_USAGE_SETTINGS", "settings cannot be nil")
}
if settings.DebounceMinutes == 0 {
// Legacy clients that omit debounce_minutes keep the fail-safe default.
settings.DebounceMinutes = ollamaCloudUsageDefaultDebounceMinutes
}
if settings.IntervalMinutes < ollamaCloudUsageMinIntervalMinutes || settings.IntervalMinutes > ollamaCloudUsageMaxIntervalMinutes {
return infraerrors.BadRequest(
"INVALID_OLLAMA_CLOUD_USAGE_INTERVAL",
fmt.Sprintf("interval_minutes must be between %d and %d", ollamaCloudUsageMinIntervalMinutes, ollamaCloudUsageMaxIntervalMinutes),
)
}
if settings.DebounceMinutes < ollamaCloudUsageMinDebounceMinutes || settings.DebounceMinutes > ollamaCloudUsageMaxDebounceMinutes {
return infraerrors.BadRequest(
"INVALID_OLLAMA_CLOUD_USAGE_DEBOUNCE",
fmt.Sprintf("debounce_minutes must be between %d and %d", ollamaCloudUsageMinDebounceMinutes, ollamaCloudUsageMaxDebounceMinutes),
)
}
// The due time is min(lastUsed+debounce, fetchedAt+maxWait). Once the debounce
// reaches the max wait the debounce term can never win, so the knob would be
// silently inert instead of doing what the operator asked for.
if settings.DebounceMinutes >= settings.IntervalMinutes {
return infraerrors.BadRequest(
"INVALID_OLLAMA_CLOUD_USAGE_DEBOUNCE",
fmt.Sprintf("debounce_minutes (%d) must be less than interval_minutes (%d)", settings.DebounceMinutes, settings.IntervalMinutes),
)
}
normalizeOllamaCloudUsageSettings(settings)
data, err := json.Marshal(settings)
if err != nil {
@@ -194,7 +237,11 @@ func (s *SettingService) SetOllamaCloudUsageSettings(ctx context.Context, settin
}
func defaultOllamaCloudUsageSettings() *OllamaCloudUsageSettings {
return &OllamaCloudUsageSettings{Enabled: false, IntervalMinutes: ollamaCloudUsageDefaultIntervalMinutes}
return &OllamaCloudUsageSettings{
Enabled: false,
IntervalMinutes: ollamaCloudUsageDefaultIntervalMinutes,
DebounceMinutes: ollamaCloudUsageDefaultDebounceMinutes,
}
}
func normalizeOllamaCloudUsageSettings(settings *OllamaCloudUsageSettings) {
@@ -204,6 +251,125 @@ func normalizeOllamaCloudUsageSettings(settings *OllamaCloudUsageSettings) {
if settings.IntervalMinutes > ollamaCloudUsageMaxIntervalMinutes {
settings.IntervalMinutes = ollamaCloudUsageMaxIntervalMinutes
}
if settings.DebounceMinutes <= 0 {
settings.DebounceMinutes = ollamaCloudUsageDefaultDebounceMinutes
}
if settings.DebounceMinutes < ollamaCloudUsageMinDebounceMinutes {
settings.DebounceMinutes = ollamaCloudUsageMinDebounceMinutes
}
if settings.DebounceMinutes > ollamaCloudUsageMaxDebounceMinutes {
settings.DebounceMinutes = ollamaCloudUsageMaxDebounceMinutes
}
}
func ollamaCloudUsageDurations(settings *OllamaCloudUsageSettings) (debounce, maxWait time.Duration) {
normalized := defaultOllamaCloudUsageSettings()
if settings != nil {
*normalized = *settings
}
normalizeOllamaCloudUsageSettings(normalized)
return time.Duration(normalized.DebounceMinutes) * time.Minute,
time.Duration(normalized.IntervalMinutes) * time.Minute
}
// ollamaCloudUsageIsAutoRefreshDue decides whether a configured auto-refresh
// group should fetch now. groupLastUsedAt must be MAX(last_used_at) across the
// exact api_key group so shared multi-platform accounts do not miss activity.
//
// Success: a request must be newer than fetched_at; dueAt = min(lastUsed+debounce, fetchedAt+maxWait).
// Failure: a request must be newer than last_attempt_at; activity due uses the same min formula,
// then dueAt = max(activityDue, next_refresh_at) so Retry-After / exponential backoff win.
// Missing or invalid snapshots fail open to a first fetch.
func ollamaCloudUsageIsAutoRefreshDue(
snapshot *OllamaCloudUsageSnapshot,
groupLastUsedAt *time.Time,
now time.Time,
debounce, maxWait time.Duration,
) bool {
dueAt, ok := ollamaCloudUsageAutoRefreshDueAt(snapshot, groupLastUsedAt, debounce, maxWait)
if !ok {
return false
}
return !now.Before(dueAt)
}
func ollamaCloudUsageAutoRefreshDueAt(
snapshot *OllamaCloudUsageSnapshot,
groupLastUsedAt *time.Time,
debounce, maxWait time.Duration,
) (time.Time, bool) {
if debounce <= 0 {
debounce = time.Duration(ollamaCloudUsageDefaultDebounceMinutes) * time.Minute
}
if maxWait <= 0 {
maxWait = time.Duration(ollamaCloudUsageDefaultIntervalMinutes) * time.Minute
}
if snapshot == nil {
return time.Time{}, true
}
switch snapshot.Status {
case OllamaCloudUsageStatusOK:
if snapshot.FetchedAt == nil || snapshot.FetchedAt.IsZero() {
return time.Time{}, true
}
fetchedAt := snapshot.FetchedAt.UTC()
if groupLastUsedAt == nil || !groupLastUsedAt.After(fetchedAt) {
return time.Time{}, false
}
lastUsed := groupLastUsedAt.UTC()
dueAt := minTime(lastUsed.Add(debounce), fetchedAt.Add(maxWait))
// Keep the pre-existing hard floor between successful fetches. The success
// path no longer consults next_refresh_at, which is where
// nextOllamaCloudUsageDelay used to apply ollamaCloudUsageMinIntervalMinutes;
// without this, request traffic spaced slightly wider than the debounce
// drives the group's outbound rate far above the previous minimum.
if floor := fetchedAt.Add(OllamaCloudUsageMinFetchInterval); dueAt.Before(floor) {
return floor, true
}
return dueAt, true
case OllamaCloudUsageStatusFailed, OllamaCloudUsageStatusUnauthorized:
if snapshot.LastAttemptAt.IsZero() {
return time.Time{}, true
}
lastAttempt := snapshot.LastAttemptAt.UTC()
if groupLastUsedAt == nil || !groupLastUsedAt.After(lastAttempt) {
return time.Time{}, false
}
lastUsed := groupLastUsedAt.UTC()
activityDue := minTime(lastUsed.Add(debounce), lastAttempt.Add(maxWait))
if !snapshot.NextRefreshAt.IsZero() && snapshot.NextRefreshAt.UTC().After(activityDue) {
return snapshot.NextRefreshAt.UTC(), true
}
return activityDue, true
default:
return time.Time{}, true
}
}
// maxOllamaCloudUsageGroupLastUsed returns the newest last_used_at among group members.
func maxOllamaCloudUsageGroupLastUsed(accounts []Account) *time.Time {
var latest *time.Time
for i := range accounts {
candidate := accounts[i].LastUsedAt
if candidate == nil || candidate.IsZero() {
continue
}
if latest == nil || candidate.After(*latest) {
ts := candidate.UTC()
latest = &ts
}
}
return latest
}
// scheduleOllamaCloudUsageActivity records that an Ollama Cloud API-key account
// actually attempted an upstream model request (including 429/5xx/transport errors).
// Local auth/validation failures must not call this. DeferredService dedupes writes.
func scheduleOllamaCloudUsageActivity(deferred *DeferredService, account *Account) {
if deferred == nil || account == nil || !IsOllamaCloudUsageAccount(account) {
return
}
deferred.ScheduleLastUsedUpdate(account.ID)
}
// OllamaCloudUsageService refreshes the official settings HTML without affecting routing state.
@@ -531,7 +697,7 @@ func (s *OllamaCloudUsageService) Refresh(ctx context.Context, accountID int64)
if err != nil {
return nil, err
}
if _, err := s.refreshAccount(ctx, accountID, settings.IntervalMinutes, false); err != nil {
if _, err := s.refreshAccount(ctx, accountID, settings, false); err != nil {
return nil, err
}
return s.GetState(ctx, accountID)
@@ -561,7 +727,8 @@ func (s *OllamaCloudUsageService) RunDue(ctx context.Context) error {
return ErrOllamaCloudUsageUnavailable
}
now := s.currentTime()
accounts, err := writer.ListDueOllamaCloudUsageAccounts(ctx, now, ollamaCloudUsageMaxPerCycle)
debounce, maxWait := ollamaCloudUsageDurations(settings)
accounts, err := writer.ListDueOllamaCloudUsageAccounts(ctx, now, debounce, maxWait, ollamaCloudUsageMaxPerCycle)
if err != nil {
return fmt.Errorf("list due Ollama Cloud usage accounts: %w", err)
}
@@ -577,13 +744,15 @@ func (s *OllamaCloudUsageService) RunDue(ctx context.Context) error {
continue
}
seenGroups[fingerprint] = struct{}{}
if snapshot := decodeOllamaCloudUsageSnapshot(account.Extra); snapshot != nil && now.Before(snapshot.NextRefreshAt) {
snapshot := decodeOllamaCloudUsageSnapshot(account.Extra)
// ListDue stamps Account.LastUsedAt with the api_key group MAX(last_used_at).
if !ollamaCloudUsageIsAutoRefreshDue(snapshot, account.LastUsedAt, now, debounce, maxWait) {
continue
}
accountID := account.ID
expected := account
group.Go(func() error {
if _, refreshErr := s.refreshAccount(ctx, accountID, settings.IntervalMinutes, true); refreshErr != nil {
if _, refreshErr := s.refreshAccount(ctx, accountID, settings, true); refreshErr != nil {
if errors.Is(refreshErr, ErrOllamaCloudUsageIdentityChanged) {
if disableErr := writer.DisableOllamaCloudUsageAutoRefresh(ctx, &expected); disableErr != nil {
logger.LegacyPrintf("service.ollama_cloud_usage", "disable_auto_refresh_failed: account_id=%d err=%v", accountID, disableErr)
@@ -598,10 +767,15 @@ func (s *OllamaCloudUsageService) RunDue(ctx context.Context) error {
return group.Wait()
}
func (s *OllamaCloudUsageService) refreshAccount(ctx context.Context, accountID int64, intervalMinutes int, requireEnabled bool) (*OllamaCloudUsageSnapshot, error) {
func (s *OllamaCloudUsageService) refreshAccount(ctx context.Context, accountID int64, settings *OllamaCloudUsageSettings, requireEnabled bool) (*OllamaCloudUsageSnapshot, error) {
if s == nil || s.accountRepo == nil {
return nil, ErrOllamaCloudUsageUnavailable
}
if settings == nil {
settings = defaultOllamaCloudUsageSettings()
}
intervalMinutes := settings.IntervalMinutes
debounce, maxWait := ollamaCloudUsageDurations(settings)
anchor, err := s.accountRepo.GetByID(ctx, accountID)
if err != nil {
return nil, err
@@ -650,7 +824,21 @@ func (s *OllamaCloudUsageService) refreshAccount(ctx context.Context, accountID
if !account.IsActive() || !ollamaCloudUsageAutoRefreshEnabled(account) {
return nil, nil
}
if snapshot := decodeOllamaCloudUsageSnapshot(account.Extra); snapshot != nil && s.currentTime().Before(snapshot.NextRefreshAt) {
groupLastUsed := account.LastUsedAt
if writer, ok := s.accountRepo.(ollamaCloudUsageRepository); ok {
siblings, listErr := writer.ListOllamaCloudUsageGroupAccounts(ctx, []*Account{account})
if listErr != nil {
// Fall back to this account's own last_used_at. That is a narrower
// activity signal than the group maximum, so the due check may skip a
// refresh it would otherwise have run; surface it rather than
// silently changing the due semantics.
logger.LegacyPrintf("service.ollama_cloud_usage",
"group_last_used_lookup_failed: account_id=%d err=%v", account.ID, listErr)
} else {
groupLastUsed = maxOllamaCloudUsageGroupLastUsed(siblings)
}
}
if !ollamaCloudUsageIsAutoRefreshDue(decodeOllamaCloudUsageSnapshot(account.Extra), groupLastUsed, s.currentTime(), debounce, maxWait) {
return nil, nil
}
}
@@ -185,7 +185,7 @@ func applyOllamaUsageTestManagedExtra(account, source *Account) {
}
}
func (r *ollamaUsageTestRepo) ListDueOllamaCloudUsageAccounts(_ context.Context, _ time.Time, limit int) ([]Account, error) {
func (r *ollamaUsageTestRepo) ListDueOllamaCloudUsageAccounts(_ context.Context, _ time.Time, _, _ time.Duration, limit int) ([]Account, error) {
r.mu.Lock()
defer r.mu.Unlock()
if len(r.due) > 0 {
@@ -313,15 +313,141 @@ func TestOllamaCloudUsageSettingsDefaultOffAndValidation(t *testing.T) {
require.NoError(t, err)
require.False(t, settings.Enabled)
require.Equal(t, 60, settings.IntervalMinutes)
require.Equal(t, 1, settings.DebounceMinutes)
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 14})
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 14, DebounceMinutes: 1})
require.Error(t, err)
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 90})
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 90, DebounceMinutes: 61})
require.Error(t, err)
// DebounceMinutes=0 (legacy omit) defaults to 1 on write.
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 90, DebounceMinutes: 0})
require.NoError(t, err)
settings, err = settingsService.GetOllamaCloudUsageSettings(context.Background())
require.NoError(t, err)
require.Equal(t, 1, settings.DebounceMinutes)
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 90, DebounceMinutes: 2})
require.NoError(t, err)
settings, err = settingsService.GetOllamaCloudUsageSettings(context.Background())
require.NoError(t, err)
require.True(t, settings.Enabled)
require.Equal(t, 90, settings.IntervalMinutes)
require.Equal(t, 2, settings.DebounceMinutes)
// debounce >= interval would make the debounce term unreachable in
// min(lastUsed+debounce, fetchedAt+maxWait), silently ignoring the operator's
// setting, so it is rejected rather than accepted and dropped.
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 15, DebounceMinutes: 15})
require.Error(t, err, "debounce equal to interval must be rejected")
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 15, DebounceMinutes: 60})
require.Error(t, err, "debounce greater than interval must be rejected")
err = settingsService.SetOllamaCloudUsageSettings(context.Background(), &OllamaCloudUsageSettings{Enabled: true, IntervalMinutes: 16, DebounceMinutes: 15})
require.NoError(t, err, "debounce below interval stays valid")
// Legacy JSON without debounce_minutes defaults to 1.
repo.values[SettingKeyOllamaCloudUsageSettings] = `{"enabled":true,"interval_minutes":45}`
settings, err = settingsService.GetOllamaCloudUsageSettings(context.Background())
require.NoError(t, err)
require.Equal(t, 45, settings.IntervalMinutes)
require.Equal(t, 1, settings.DebounceMinutes)
}
func TestOllamaCloudUsageIsAutoRefreshDue(t *testing.T) {
debounce := time.Minute
maxWait := time.Hour
now := time.Date(2026, time.July, 25, 12, 0, 0, 0, time.UTC)
fetched := now.Add(-30 * time.Minute)
ptr := func(ts time.Time) *time.Time { return &ts }
require.True(t, ollamaCloudUsageIsAutoRefreshDue(nil, nil, now, debounce, maxWait), "missing snapshot first due")
require.True(t, ollamaCloudUsageIsAutoRefreshDue(&OllamaCloudUsageSnapshot{Status: "bogus"}, nil, now, debounce, maxWait), "invalid status first due")
okSnap := &OllamaCloudUsageSnapshot{
Status: OllamaCloudUsageStatusOK, FetchedAt: ptr(fetched),
LastAttemptAt: fetched, NextRefreshAt: fetched.Add(maxWait),
}
require.False(t, ollamaCloudUsageIsAutoRefreshDue(okSnap, nil, now, debounce, maxWait), "no request after success")
require.False(t, ollamaCloudUsageIsAutoRefreshDue(okSnap, ptr(fetched), now, debounce, maxWait), "request not after fetched_at")
require.False(t, ollamaCloudUsageIsAutoRefreshDue(okSnap, ptr(now.Add(-30*time.Second)), now, debounce, maxWait), "debounce not elapsed")
require.True(t, ollamaCloudUsageIsAutoRefreshDue(okSnap, ptr(now.Add(-time.Minute)), now, debounce, maxWait), "single request quiet for debounce")
// Continuous requests: last used is now, but max-wait from old fetch forces due.
oldFetched := now.Add(-2 * time.Hour)
oldSnap := &OllamaCloudUsageSnapshot{
Status: OllamaCloudUsageStatusOK, FetchedAt: ptr(oldFetched),
LastAttemptAt: oldFetched, NextRefreshAt: oldFetched.Add(maxWait),
}
require.True(t, ollamaCloudUsageIsAutoRefreshDue(oldSnap, ptr(now), now, debounce, maxWait), "max-wait forces due while requests continue")
// First request after a very old snapshot is immediately due because fetched+maxWait is past.
require.True(t, ollamaCloudUsageIsAutoRefreshDue(oldSnap, ptr(now.Add(-time.Second)), now, debounce, maxWait), "stale snapshot first request immediate")
failSnap := &OllamaCloudUsageSnapshot{
Status: OllamaCloudUsageStatusFailed, FetchedAt: ptr(fetched),
LastAttemptAt: now.Add(-10 * time.Minute), NextRefreshAt: now.Add(20 * time.Minute),
}
require.False(t, ollamaCloudUsageIsAutoRefreshDue(failSnap, nil, now, debounce, maxWait), "failure without new request")
require.False(t, ollamaCloudUsageIsAutoRefreshDue(failSnap, ptr(now.Add(-time.Minute)), now, debounce, maxWait), "failure blocked by backoff")
failSnap.NextRefreshAt = now.Add(-time.Second)
require.True(t, ollamaCloudUsageIsAutoRefreshDue(failSnap, ptr(now.Add(-time.Minute)), now, debounce, maxWait), "failure after backoff with new request")
require.True(t, ollamaCloudUsageIsAutoRefreshDue(&OllamaCloudUsageSnapshot{
Status: OllamaCloudUsageStatusOK, LastAttemptAt: now,
}, nil, now, debounce, maxWait), "ok without fetched_at fails open")
}
// The success path stopped consulting next_refresh_at, which is where
// nextOllamaCloudUsageDelay used to apply the minimum interval. Activity may pull
// a refresh forward only as far as that floor, otherwise request traffic spaced
// just wider than the debounce drives the group's outbound rate far above the
// pre-existing minimum.
func TestOllamaCloudUsageAutoRefreshDueAtHonoursMinFetchInterval(t *testing.T) {
debounce := time.Minute
maxWait := time.Hour
now := time.Date(2026, time.July, 25, 12, 0, 0, 0, time.UTC)
ptr := func(ts time.Time) *time.Time { return &ts }
// Debounce elapsed, but the last successful fetch is inside the floor.
recent := now.Add(-5 * time.Minute)
recentSnap := &OllamaCloudUsageSnapshot{
Status: OllamaCloudUsageStatusOK, FetchedAt: ptr(recent), LastAttemptAt: recent,
}
dueAt, ok := ollamaCloudUsageAutoRefreshDueAt(recentSnap, ptr(now.Add(-2*time.Minute)), debounce, maxWait)
require.True(t, ok)
require.Equal(t, recent.Add(OllamaCloudUsageMinFetchInterval), dueAt,
"due time must be clamped to fetched_at + min fetch interval")
require.False(t, ollamaCloudUsageIsAutoRefreshDue(recentSnap, ptr(now.Add(-2*time.Minute)), now, debounce, maxWait),
"debounce alone must not refresh within the min fetch interval")
// Once the floor has passed the debounce governs again.
atFloor := now.Add(-OllamaCloudUsageMinFetchInterval)
floorSnap := &OllamaCloudUsageSnapshot{
Status: OllamaCloudUsageStatusOK, FetchedAt: ptr(atFloor), LastAttemptAt: atFloor,
}
require.True(t, ollamaCloudUsageIsAutoRefreshDue(floorSnap, ptr(now.Add(-2*time.Minute)), now, debounce, maxWait),
"past the floor a quiet debounce window is due")
// The floor never delays a refresh that max-wait has already forced.
stale := now.Add(-2 * time.Hour)
staleSnap := &OllamaCloudUsageSnapshot{
Status: OllamaCloudUsageStatusOK, FetchedAt: ptr(stale), LastAttemptAt: stale,
}
require.True(t, ollamaCloudUsageIsAutoRefreshDue(staleSnap, ptr(now), now, debounce, maxWait),
"max-wait still forces due on a stale snapshot")
}
func TestScheduleOllamaCloudUsageActivityOnlyForOllama(t *testing.T) {
deferred := NewDeferredService(nil, nil, time.Second)
ollama := ollamaUsageAccount(1)
other := ollamaUsageAccount(2)
other.Credentials["base_url"] = "https://api.openai.com"
scheduleOllamaCloudUsageActivity(deferred, ollama)
scheduleOllamaCloudUsageActivity(deferred, other)
scheduleOllamaCloudUsageActivity(nil, ollama)
_, ok := deferred.lastUsedUpdates.Load(int64(1))
require.True(t, ok)
_, ok = deferred.lastUsedUpdates.Load(int64(2))
require.False(t, ok)
}
func TestIsOllamaCloudUsageAccountStrictOfficialHost(t *testing.T) {
@@ -51,6 +51,10 @@ func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Cont
if account != nil && account.Platform == PlatformGrok && isGrokContentPolicyRejection(statusCode, responseBody) {
return false
}
// Any non-2xx upstream HTTP response means the model request was actually sent.
if s != nil {
scheduleOllamaCloudUsageActivity(s.deferredService, account)
}
stateCtx, cancel := openAIAccountStateContext(ctx)
defer cancel()
@@ -91,7 +95,12 @@ func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Cont
if shouldDisable && !modelTempMatched {
s.BlockAccountScheduling(account, time.Time{}, "upstream_disable")
}
if !shouldDisable && account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey && shouldCooldownOpenAITransientUpstreamError(statusCode, responseBody) {
// Pool-mode retryable upstream errors are already bounded by the request-local
// same-account retry budget. Recording the generic account+model transient
// cooldown here would block the next approved retry before that budget is used.
poolModeRetryable := account.IsPoolMode() && account.IsPoolModeRetryableStatus(statusCode)
if !shouldDisable && account.Platform == PlatformOpenAI && account.Type == AccountTypeAPIKey &&
shouldCooldownOpenAITransientUpstreamError(statusCode, responseBody) && !poolModeRetryable {
model := ""
if len(canonicalModel) > 0 {
model = canonicalModel[0]
@@ -142,6 +142,89 @@ func TestOpenAIPoolModeTempRule_StopsSameAccountRetryAndIsolatesBlockToModel(t *
require.False(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.5"))
}
func TestOpenAIPoolModeRetryable5xx_DoesNotCreateModelTransientBlock(t *testing.T) {
repo := &errorPolicyRepoStub{}
rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
gateway := &OpenAIGatewayService{rateLimitService: rateLimitService}
account := &Account{
ID: 47,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
"pool_mode": true,
"pool_mode_retry_status_codes": []any{float64(524)},
},
}
for i := 0; i < 2; i++ {
shouldDisable := gateway.handleOpenAIAccountUpstreamError(
context.Background(),
account,
524,
http.Header{},
[]byte(`{"error":{"message":"upstream timeout"}}`),
"gpt-5.4",
)
require.False(t, shouldDisable)
}
require.False(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.4"))
}
func TestOpenAIPoolModeNonRetryable5xx_StillCreatesModelTransientBlock(t *testing.T) {
repo := &errorPolicyRepoStub{}
rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
gateway := &OpenAIGatewayService{rateLimitService: rateLimitService}
account := &Account{
ID: 48,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
"pool_mode": true,
"pool_mode_retry_status_codes": []any{float64(http.StatusGatewayTimeout)},
},
}
for i := 0; i < 2; i++ {
shouldDisable := gateway.handleOpenAIAccountUpstreamError(
context.Background(),
account,
http.StatusServiceUnavailable,
http.Header{},
[]byte(`{"error":{"message":"upstream unavailable"}}`),
"gpt-5.4",
)
require.False(t, shouldDisable)
}
require.True(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.4"))
}
func TestOpenAINonPoolAPIKey5xx_StillCreatesModelTransientBlock(t *testing.T) {
repo := &errorPolicyRepoStub{}
rateLimitService := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
gateway := &OpenAIGatewayService{rateLimitService: rateLimitService}
account := &Account{
ID: 49,
Platform: PlatformOpenAI,
Type: AccountTypeAPIKey,
}
for i := 0; i < 2; i++ {
shouldDisable := gateway.handleOpenAIAccountUpstreamError(
context.Background(),
account,
http.StatusGatewayTimeout,
http.Header{},
[]byte(`{"error":{"message":"upstream timeout"}}`),
"gpt-5.4",
)
require.False(t, shouldDisable)
}
require.True(t, gateway.isOpenAIAccountRequestRuntimeBlocked(account, "gpt-5.4"))
}
func TestOpenAIModelNotFound_DoesNotRuntimeBlockWholeAccount(t *testing.T) {
repo := &modelNotFoundAccountRepoStub{}
svc := &OpenAIGatewayService{
Loaded 100 of 168 files, more files were not shown because too many files have changed in this diff. Show more