Merge remote-tracking branch 'origin/main' into fix/composite-route-prefix-passthrough
This commit is contained in:
commit
1614ae9c99
168 files changed
+7169
-412
No files matched your search
@@ -1 +1 @@
|
||||
0.1.164
|
||||
0.1.165
|
||||
+12
-1
@@ -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(", ")
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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 类型账号关联到此分组"),
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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"},
|
||||
|
||||
@@ -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 上游
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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: "",
|
||||
},
|
||||
{
|
||||
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",
|
||||
},
|
||||
}
|
||||
|
||||
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": ""}, 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("", 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() {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
Reference in new issue
Block a user