Merge upstream/main into feat/grok-complete-integration
冲突集中在 chat completions / messages 两条 Responses 转发路径: upstream 给 OpenAIForwardResult 增加了 UpstreamResponseModel 与 UpstreamResponseModelConflict(配套 beginUpstreamResponseModelObservation 观测器),本分支在同样位置把返回值改成了具名变量以便挂 Grok 原生搜索计数。 两侧不互斥,合并结果同时保留上游的响应模型观测字段与 Grok SearchCount 逻辑。 frontend/pnpm-lock.yaml 取 upstream 版本:package.json 与 upstream 完全一致, 本地差异只是 pnpm install 的重解析噪音。
This commit is contained in:
commit
1f58e25ab3
141 files changed
+5942
-680
No files matched your search
@@ -1 +1 @@
|
||||
0.1.171
|
||||
0.1.172
|
||||
@@ -1626,6 +1626,8 @@ var (
|
||||
{Name: "model", Type: field.TypeString, Size: 100},
|
||||
{Name: "requested_model", Type: field.TypeString, Nullable: true, Size: 100},
|
||||
{Name: "upstream_model", Type: field.TypeString, Nullable: true, Size: 100},
|
||||
{Name: "upstream_response_model", Type: field.TypeString, Nullable: true, Size: 200},
|
||||
{Name: "upstream_model_mismatch", Type: field.TypeBool, Nullable: true},
|
||||
{Name: "channel_id", Type: field.TypeInt64, Nullable: true},
|
||||
{Name: "model_mapping_chain", Type: field.TypeString, Nullable: true, Size: 500},
|
||||
{Name: "billing_tier", Type: field.TypeString, Nullable: true, Size: 50},
|
||||
@@ -1676,31 +1678,31 @@ var (
|
||||
ForeignKeys: []*schema.ForeignKey{
|
||||
{
|
||||
Symbol: "usage_logs_api_keys_usage_logs",
|
||||
Columns: []*schema.Column{UsageLogsColumns[41]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[43]},
|
||||
RefColumns: []*schema.Column{APIKeysColumns[0]},
|
||||
OnDelete: schema.NoAction,
|
||||
},
|
||||
{
|
||||
Symbol: "usage_logs_accounts_usage_logs",
|
||||
Columns: []*schema.Column{UsageLogsColumns[42]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[44]},
|
||||
RefColumns: []*schema.Column{AccountsColumns[0]},
|
||||
OnDelete: schema.NoAction,
|
||||
},
|
||||
{
|
||||
Symbol: "usage_logs_groups_usage_logs",
|
||||
Columns: []*schema.Column{UsageLogsColumns[43]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[45]},
|
||||
RefColumns: []*schema.Column{GroupsColumns[0]},
|
||||
OnDelete: schema.SetNull,
|
||||
},
|
||||
{
|
||||
Symbol: "usage_logs_users_usage_logs",
|
||||
Columns: []*schema.Column{UsageLogsColumns[44]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[46]},
|
||||
RefColumns: []*schema.Column{UsersColumns[0]},
|
||||
OnDelete: schema.NoAction,
|
||||
},
|
||||
{
|
||||
Symbol: "usage_logs_user_subscriptions_usage_logs",
|
||||
Columns: []*schema.Column{UsageLogsColumns[45]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[47]},
|
||||
RefColumns: []*schema.Column{UserSubscriptionsColumns[0]},
|
||||
OnDelete: schema.SetNull,
|
||||
},
|
||||
@@ -1709,32 +1711,32 @@ var (
|
||||
{
|
||||
Name: "usagelog_user_id",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[44]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[46]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_api_key_id",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[41]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[43]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_account_id",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[42]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[44]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_group_id",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[43]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[45]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_subscription_id",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[45]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[47]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_created_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[40]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[42]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_model",
|
||||
@@ -1754,17 +1756,17 @@ var (
|
||||
{
|
||||
Name: "usagelog_user_id_created_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[44], UsageLogsColumns[40]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[46], UsageLogsColumns[42]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_api_key_id_created_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[41], UsageLogsColumns[40]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[43], UsageLogsColumns[42]},
|
||||
},
|
||||
{
|
||||
Name: "usagelog_group_id_created_at",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{UsageLogsColumns[43], UsageLogsColumns[40]},
|
||||
Columns: []*schema.Column{UsageLogsColumns[45], UsageLogsColumns[42]},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
+147
-1
@@ -43857,6 +43857,8 @@ type UsageLogMutation struct {
|
||||
model *string
|
||||
requested_model *string
|
||||
upstream_model *string
|
||||
upstream_response_model *string
|
||||
upstream_model_mismatch *bool
|
||||
channel_id *int64
|
||||
addchannel_id *int64
|
||||
model_mapping_chain *string
|
||||
@@ -44306,6 +44308,104 @@ func (m *UsageLogMutation) ResetUpstreamModel() {
|
||||
delete(m.clearedFields, usagelog.FieldUpstreamModel)
|
||||
}
|
||||
|
||||
// SetUpstreamResponseModel sets the "upstream_response_model" field.
|
||||
func (m *UsageLogMutation) SetUpstreamResponseModel(s string) {
|
||||
m.upstream_response_model = &s
|
||||
}
|
||||
|
||||
// UpstreamResponseModel returns the value of the "upstream_response_model" field in the mutation.
|
||||
func (m *UsageLogMutation) UpstreamResponseModel() (r string, exists bool) {
|
||||
v := m.upstream_response_model
|
||||
if v == nil {
|
||||
return
|
||||
}
|
||||
return *v, true
|
||||
}
|
||||
|
||||
// OldUpstreamResponseModel returns the old "upstream_response_model" field's value of the UsageLog entity.
|
||||
// If the UsageLog 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 *UsageLogMutation) OldUpstreamResponseModel(ctx context.Context) (v *string, err error) {
|
||||
if !m.op.Is(OpUpdateOne) {
|
||||
return v, errors.New("OldUpstreamResponseModel is only allowed on UpdateOne operations")
|
||||
}
|
||||
if m.id == nil || m.oldValue == nil {
|
||||
return v, errors.New("OldUpstreamResponseModel requires an ID field in the mutation")
|
||||
}
|
||||
oldValue, err := m.oldValue(ctx)
|
||||
if err != nil {
|
||||
return v, fmt.Errorf("querying old value for OldUpstreamResponseModel: %w", err)
|
||||
}
|
||||
return oldValue.UpstreamResponseModel, nil
|
||||
}
|
||||
|
||||
// ClearUpstreamResponseModel clears the value of the "upstream_response_model" field.
|
||||
func (m *UsageLogMutation) ClearUpstreamResponseModel() {
|
||||
m.upstream_response_model = nil
|
||||
m.clearedFields[usagelog.FieldUpstreamResponseModel] = struct{}{}
|
||||
}
|
||||
|
||||
// UpstreamResponseModelCleared returns if the "upstream_response_model" field was cleared in this mutation.
|
||||
func (m *UsageLogMutation) UpstreamResponseModelCleared() bool {
|
||||
_, ok := m.clearedFields[usagelog.FieldUpstreamResponseModel]
|
||||
return ok
|
||||
}
|
||||
|
||||
// ResetUpstreamResponseModel resets all changes to the "upstream_response_model" field.
|
||||
func (m *UsageLogMutation) ResetUpstreamResponseModel() {
|
||||
m.upstream_response_model = nil
|
||||
delete(m.clearedFields, usagelog.FieldUpstreamResponseModel)
|
||||
}
|
||||
|
||||
// SetUpstreamModelMismatch sets the "upstream_model_mismatch" field.
|
||||
func (m *UsageLogMutation) SetUpstreamModelMismatch(b bool) {
|
||||
m.upstream_model_mismatch = &b
|
||||
}
|
||||
|
||||
// UpstreamModelMismatch returns the value of the "upstream_model_mismatch" field in the mutation.
|
||||
func (m *UsageLogMutation) UpstreamModelMismatch() (r bool, exists bool) {
|
||||
v := m.upstream_model_mismatch
|
||||
if v == nil {
|
||||
return
|
||||
}
|
||||
return *v, true
|
||||
}
|
||||
|
||||
// OldUpstreamModelMismatch returns the old "upstream_model_mismatch" field's value of the UsageLog entity.
|
||||
// If the UsageLog 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 *UsageLogMutation) OldUpstreamModelMismatch(ctx context.Context) (v *bool, err error) {
|
||||
if !m.op.Is(OpUpdateOne) {
|
||||
return v, errors.New("OldUpstreamModelMismatch is only allowed on UpdateOne operations")
|
||||
}
|
||||
if m.id == nil || m.oldValue == nil {
|
||||
return v, errors.New("OldUpstreamModelMismatch requires an ID field in the mutation")
|
||||
}
|
||||
oldValue, err := m.oldValue(ctx)
|
||||
if err != nil {
|
||||
return v, fmt.Errorf("querying old value for OldUpstreamModelMismatch: %w", err)
|
||||
}
|
||||
return oldValue.UpstreamModelMismatch, nil
|
||||
}
|
||||
|
||||
// ClearUpstreamModelMismatch clears the value of the "upstream_model_mismatch" field.
|
||||
func (m *UsageLogMutation) ClearUpstreamModelMismatch() {
|
||||
m.upstream_model_mismatch = nil
|
||||
m.clearedFields[usagelog.FieldUpstreamModelMismatch] = struct{}{}
|
||||
}
|
||||
|
||||
// UpstreamModelMismatchCleared returns if the "upstream_model_mismatch" field was cleared in this mutation.
|
||||
func (m *UsageLogMutation) UpstreamModelMismatchCleared() bool {
|
||||
_, ok := m.clearedFields[usagelog.FieldUpstreamModelMismatch]
|
||||
return ok
|
||||
}
|
||||
|
||||
// ResetUpstreamModelMismatch resets all changes to the "upstream_model_mismatch" field.
|
||||
func (m *UsageLogMutation) ResetUpstreamModelMismatch() {
|
||||
m.upstream_model_mismatch = nil
|
||||
delete(m.clearedFields, usagelog.FieldUpstreamModelMismatch)
|
||||
}
|
||||
|
||||
// SetChannelID sets the "channel_id" field.
|
||||
func (m *UsageLogMutation) SetChannelID(i int64) {
|
||||
m.channel_id = &i
|
||||
@@ -46502,7 +46602,7 @@ func (m *UsageLogMutation) Type() string {
|
||||
// order to get all numeric fields that were incremented/decremented, call
|
||||
// AddedFields().
|
||||
func (m *UsageLogMutation) Fields() []string {
|
||||
fields := make([]string, 0, 45)
|
||||
fields := make([]string, 0, 47)
|
||||
if m.user != nil {
|
||||
fields = append(fields, usagelog.FieldUserID)
|
||||
}
|
||||
@@ -46524,6 +46624,12 @@ func (m *UsageLogMutation) Fields() []string {
|
||||
if m.upstream_model != nil {
|
||||
fields = append(fields, usagelog.FieldUpstreamModel)
|
||||
}
|
||||
if m.upstream_response_model != nil {
|
||||
fields = append(fields, usagelog.FieldUpstreamResponseModel)
|
||||
}
|
||||
if m.upstream_model_mismatch != nil {
|
||||
fields = append(fields, usagelog.FieldUpstreamModelMismatch)
|
||||
}
|
||||
if m.channel_id != nil {
|
||||
fields = append(fields, usagelog.FieldChannelID)
|
||||
}
|
||||
@@ -46660,6 +46766,10 @@ func (m *UsageLogMutation) Field(name string) (ent.Value, bool) {
|
||||
return m.RequestedModel()
|
||||
case usagelog.FieldUpstreamModel:
|
||||
return m.UpstreamModel()
|
||||
case usagelog.FieldUpstreamResponseModel:
|
||||
return m.UpstreamResponseModel()
|
||||
case usagelog.FieldUpstreamModelMismatch:
|
||||
return m.UpstreamModelMismatch()
|
||||
case usagelog.FieldChannelID:
|
||||
return m.ChannelID()
|
||||
case usagelog.FieldModelMappingChain:
|
||||
@@ -46759,6 +46869,10 @@ func (m *UsageLogMutation) OldField(ctx context.Context, name string) (ent.Value
|
||||
return m.OldRequestedModel(ctx)
|
||||
case usagelog.FieldUpstreamModel:
|
||||
return m.OldUpstreamModel(ctx)
|
||||
case usagelog.FieldUpstreamResponseModel:
|
||||
return m.OldUpstreamResponseModel(ctx)
|
||||
case usagelog.FieldUpstreamModelMismatch:
|
||||
return m.OldUpstreamModelMismatch(ctx)
|
||||
case usagelog.FieldChannelID:
|
||||
return m.OldChannelID(ctx)
|
||||
case usagelog.FieldModelMappingChain:
|
||||
@@ -46893,6 +47007,20 @@ func (m *UsageLogMutation) SetField(name string, value ent.Value) error {
|
||||
}
|
||||
m.SetUpstreamModel(v)
|
||||
return nil
|
||||
case usagelog.FieldUpstreamResponseModel:
|
||||
v, ok := value.(string)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T for field %s", value, name)
|
||||
}
|
||||
m.SetUpstreamResponseModel(v)
|
||||
return nil
|
||||
case usagelog.FieldUpstreamModelMismatch:
|
||||
v, ok := value.(bool)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T for field %s", value, name)
|
||||
}
|
||||
m.SetUpstreamModelMismatch(v)
|
||||
return nil
|
||||
case usagelog.FieldChannelID:
|
||||
v, ok := value.(int64)
|
||||
if !ok {
|
||||
@@ -47450,6 +47578,12 @@ func (m *UsageLogMutation) ClearedFields() []string {
|
||||
if m.FieldCleared(usagelog.FieldUpstreamModel) {
|
||||
fields = append(fields, usagelog.FieldUpstreamModel)
|
||||
}
|
||||
if m.FieldCleared(usagelog.FieldUpstreamResponseModel) {
|
||||
fields = append(fields, usagelog.FieldUpstreamResponseModel)
|
||||
}
|
||||
if m.FieldCleared(usagelog.FieldUpstreamModelMismatch) {
|
||||
fields = append(fields, usagelog.FieldUpstreamModelMismatch)
|
||||
}
|
||||
if m.FieldCleared(usagelog.FieldChannelID) {
|
||||
fields = append(fields, usagelog.FieldChannelID)
|
||||
}
|
||||
@@ -47524,6 +47658,12 @@ func (m *UsageLogMutation) ClearField(name string) error {
|
||||
case usagelog.FieldUpstreamModel:
|
||||
m.ClearUpstreamModel()
|
||||
return nil
|
||||
case usagelog.FieldUpstreamResponseModel:
|
||||
m.ClearUpstreamResponseModel()
|
||||
return nil
|
||||
case usagelog.FieldUpstreamModelMismatch:
|
||||
m.ClearUpstreamModelMismatch()
|
||||
return nil
|
||||
case usagelog.FieldChannelID:
|
||||
m.ClearChannelID()
|
||||
return nil
|
||||
@@ -47607,6 +47747,12 @@ func (m *UsageLogMutation) ResetField(name string) error {
|
||||
case usagelog.FieldUpstreamModel:
|
||||
m.ResetUpstreamModel()
|
||||
return nil
|
||||
case usagelog.FieldUpstreamResponseModel:
|
||||
m.ResetUpstreamResponseModel()
|
||||
return nil
|
||||
case usagelog.FieldUpstreamModelMismatch:
|
||||
m.ResetUpstreamModelMismatch()
|
||||
return nil
|
||||
case usagelog.FieldChannelID:
|
||||
m.ResetChannelID()
|
||||
return nil
|
||||
|
||||
@@ -1998,124 +1998,128 @@ func init() {
|
||||
usagelogDescUpstreamModel := usagelogFields[6].Descriptor()
|
||||
// usagelog.UpstreamModelValidator is a validator for the "upstream_model" field. It is called by the builders before save.
|
||||
usagelog.UpstreamModelValidator = usagelogDescUpstreamModel.Validators[0].(func(string) error)
|
||||
// usagelogDescUpstreamResponseModel is the schema descriptor for upstream_response_model field.
|
||||
usagelogDescUpstreamResponseModel := usagelogFields[7].Descriptor()
|
||||
// usagelog.UpstreamResponseModelValidator is a validator for the "upstream_response_model" field. It is called by the builders before save.
|
||||
usagelog.UpstreamResponseModelValidator = usagelogDescUpstreamResponseModel.Validators[0].(func(string) error)
|
||||
// usagelogDescModelMappingChain is the schema descriptor for model_mapping_chain field.
|
||||
usagelogDescModelMappingChain := usagelogFields[8].Descriptor()
|
||||
usagelogDescModelMappingChain := usagelogFields[10].Descriptor()
|
||||
// usagelog.ModelMappingChainValidator is a validator for the "model_mapping_chain" field. It is called by the builders before save.
|
||||
usagelog.ModelMappingChainValidator = usagelogDescModelMappingChain.Validators[0].(func(string) error)
|
||||
// usagelogDescBillingTier is the schema descriptor for billing_tier field.
|
||||
usagelogDescBillingTier := usagelogFields[9].Descriptor()
|
||||
usagelogDescBillingTier := usagelogFields[11].Descriptor()
|
||||
// usagelog.BillingTierValidator is a validator for the "billing_tier" field. It is called by the builders before save.
|
||||
usagelog.BillingTierValidator = usagelogDescBillingTier.Validators[0].(func(string) error)
|
||||
// usagelogDescBillingMode is the schema descriptor for billing_mode field.
|
||||
usagelogDescBillingMode := usagelogFields[10].Descriptor()
|
||||
usagelogDescBillingMode := usagelogFields[12].Descriptor()
|
||||
// usagelog.BillingModeValidator is a validator for the "billing_mode" field. It is called by the builders before save.
|
||||
usagelog.BillingModeValidator = usagelogDescBillingMode.Validators[0].(func(string) error)
|
||||
// usagelogDescInputTokens is the schema descriptor for input_tokens field.
|
||||
usagelogDescInputTokens := usagelogFields[13].Descriptor()
|
||||
usagelogDescInputTokens := usagelogFields[15].Descriptor()
|
||||
// usagelog.DefaultInputTokens holds the default value on creation for the input_tokens field.
|
||||
usagelog.DefaultInputTokens = usagelogDescInputTokens.Default.(int)
|
||||
// usagelogDescOutputTokens is the schema descriptor for output_tokens field.
|
||||
usagelogDescOutputTokens := usagelogFields[14].Descriptor()
|
||||
usagelogDescOutputTokens := usagelogFields[16].Descriptor()
|
||||
// usagelog.DefaultOutputTokens holds the default value on creation for the output_tokens field.
|
||||
usagelog.DefaultOutputTokens = usagelogDescOutputTokens.Default.(int)
|
||||
// usagelogDescCacheCreationTokens is the schema descriptor for cache_creation_tokens field.
|
||||
usagelogDescCacheCreationTokens := usagelogFields[15].Descriptor()
|
||||
usagelogDescCacheCreationTokens := usagelogFields[17].Descriptor()
|
||||
// usagelog.DefaultCacheCreationTokens holds the default value on creation for the cache_creation_tokens field.
|
||||
usagelog.DefaultCacheCreationTokens = usagelogDescCacheCreationTokens.Default.(int)
|
||||
// usagelogDescCacheReadTokens is the schema descriptor for cache_read_tokens field.
|
||||
usagelogDescCacheReadTokens := usagelogFields[16].Descriptor()
|
||||
usagelogDescCacheReadTokens := usagelogFields[18].Descriptor()
|
||||
// usagelog.DefaultCacheReadTokens holds the default value on creation for the cache_read_tokens field.
|
||||
usagelog.DefaultCacheReadTokens = usagelogDescCacheReadTokens.Default.(int)
|
||||
// usagelogDescCacheCreation5mTokens is the schema descriptor for cache_creation_5m_tokens field.
|
||||
usagelogDescCacheCreation5mTokens := usagelogFields[17].Descriptor()
|
||||
usagelogDescCacheCreation5mTokens := usagelogFields[19].Descriptor()
|
||||
// usagelog.DefaultCacheCreation5mTokens holds the default value on creation for the cache_creation_5m_tokens field.
|
||||
usagelog.DefaultCacheCreation5mTokens = usagelogDescCacheCreation5mTokens.Default.(int)
|
||||
// usagelogDescCacheCreation1hTokens is the schema descriptor for cache_creation_1h_tokens field.
|
||||
usagelogDescCacheCreation1hTokens := usagelogFields[18].Descriptor()
|
||||
usagelogDescCacheCreation1hTokens := usagelogFields[20].Descriptor()
|
||||
// usagelog.DefaultCacheCreation1hTokens holds the default value on creation for the cache_creation_1h_tokens field.
|
||||
usagelog.DefaultCacheCreation1hTokens = usagelogDescCacheCreation1hTokens.Default.(int)
|
||||
// usagelogDescInputCost is the schema descriptor for input_cost field.
|
||||
usagelogDescInputCost := usagelogFields[19].Descriptor()
|
||||
usagelogDescInputCost := usagelogFields[21].Descriptor()
|
||||
// usagelog.DefaultInputCost holds the default value on creation for the input_cost field.
|
||||
usagelog.DefaultInputCost = usagelogDescInputCost.Default.(float64)
|
||||
// usagelogDescOutputCost is the schema descriptor for output_cost field.
|
||||
usagelogDescOutputCost := usagelogFields[20].Descriptor()
|
||||
usagelogDescOutputCost := usagelogFields[22].Descriptor()
|
||||
// usagelog.DefaultOutputCost holds the default value on creation for the output_cost field.
|
||||
usagelog.DefaultOutputCost = usagelogDescOutputCost.Default.(float64)
|
||||
// usagelogDescCacheCreationCost is the schema descriptor for cache_creation_cost field.
|
||||
usagelogDescCacheCreationCost := usagelogFields[21].Descriptor()
|
||||
usagelogDescCacheCreationCost := usagelogFields[23].Descriptor()
|
||||
// usagelog.DefaultCacheCreationCost holds the default value on creation for the cache_creation_cost field.
|
||||
usagelog.DefaultCacheCreationCost = usagelogDescCacheCreationCost.Default.(float64)
|
||||
// usagelogDescCacheReadCost is the schema descriptor for cache_read_cost field.
|
||||
usagelogDescCacheReadCost := usagelogFields[22].Descriptor()
|
||||
usagelogDescCacheReadCost := usagelogFields[24].Descriptor()
|
||||
// usagelog.DefaultCacheReadCost holds the default value on creation for the cache_read_cost field.
|
||||
usagelog.DefaultCacheReadCost = usagelogDescCacheReadCost.Default.(float64)
|
||||
// usagelogDescTotalCost is the schema descriptor for total_cost field.
|
||||
usagelogDescTotalCost := usagelogFields[23].Descriptor()
|
||||
usagelogDescTotalCost := usagelogFields[25].Descriptor()
|
||||
// usagelog.DefaultTotalCost holds the default value on creation for the total_cost field.
|
||||
usagelog.DefaultTotalCost = usagelogDescTotalCost.Default.(float64)
|
||||
// usagelogDescActualCost is the schema descriptor for actual_cost field.
|
||||
usagelogDescActualCost := usagelogFields[24].Descriptor()
|
||||
usagelogDescActualCost := usagelogFields[26].Descriptor()
|
||||
// usagelog.DefaultActualCost holds the default value on creation for the actual_cost field.
|
||||
usagelog.DefaultActualCost = usagelogDescActualCost.Default.(float64)
|
||||
// usagelogDescRateMultiplier is the schema descriptor for rate_multiplier field.
|
||||
usagelogDescRateMultiplier := usagelogFields[25].Descriptor()
|
||||
usagelogDescRateMultiplier := usagelogFields[27].Descriptor()
|
||||
// usagelog.DefaultRateMultiplier holds the default value on creation for the rate_multiplier field.
|
||||
usagelog.DefaultRateMultiplier = usagelogDescRateMultiplier.Default.(float64)
|
||||
// usagelogDescLongContextBillingApplied is the schema descriptor for long_context_billing_applied field.
|
||||
usagelogDescLongContextBillingApplied := usagelogFields[26].Descriptor()
|
||||
usagelogDescLongContextBillingApplied := usagelogFields[28].Descriptor()
|
||||
// usagelog.DefaultLongContextBillingApplied holds the default value on creation for the long_context_billing_applied field.
|
||||
usagelog.DefaultLongContextBillingApplied = usagelogDescLongContextBillingApplied.Default.(bool)
|
||||
// usagelogDescBillingType is the schema descriptor for billing_type field.
|
||||
usagelogDescBillingType := usagelogFields[28].Descriptor()
|
||||
usagelogDescBillingType := usagelogFields[30].Descriptor()
|
||||
// usagelog.DefaultBillingType holds the default value on creation for the billing_type field.
|
||||
usagelog.DefaultBillingType = usagelogDescBillingType.Default.(int8)
|
||||
// usagelogDescStream is the schema descriptor for stream field.
|
||||
usagelogDescStream := usagelogFields[29].Descriptor()
|
||||
usagelogDescStream := usagelogFields[31].Descriptor()
|
||||
// usagelog.DefaultStream holds the default value on creation for the stream field.
|
||||
usagelog.DefaultStream = usagelogDescStream.Default.(bool)
|
||||
// usagelogDescUserAgent is the schema descriptor for user_agent field.
|
||||
usagelogDescUserAgent := usagelogFields[32].Descriptor()
|
||||
usagelogDescUserAgent := usagelogFields[34].Descriptor()
|
||||
// usagelog.UserAgentValidator is a validator for the "user_agent" field. It is called by the builders before save.
|
||||
usagelog.UserAgentValidator = usagelogDescUserAgent.Validators[0].(func(string) error)
|
||||
// usagelogDescIPAddress is the schema descriptor for ip_address field.
|
||||
usagelogDescIPAddress := usagelogFields[33].Descriptor()
|
||||
usagelogDescIPAddress := usagelogFields[35].Descriptor()
|
||||
// usagelog.IPAddressValidator is a validator for the "ip_address" field. It is called by the builders before save.
|
||||
usagelog.IPAddressValidator = usagelogDescIPAddress.Validators[0].(func(string) error)
|
||||
// usagelogDescImageCount is the schema descriptor for image_count field.
|
||||
usagelogDescImageCount := usagelogFields[34].Descriptor()
|
||||
usagelogDescImageCount := usagelogFields[36].Descriptor()
|
||||
// usagelog.DefaultImageCount holds the default value on creation for the image_count field.
|
||||
usagelog.DefaultImageCount = usagelogDescImageCount.Default.(int)
|
||||
// usagelogDescImageSize is the schema descriptor for image_size field.
|
||||
usagelogDescImageSize := usagelogFields[35].Descriptor()
|
||||
usagelogDescImageSize := usagelogFields[37].Descriptor()
|
||||
// usagelog.ImageSizeValidator is a validator for the "image_size" field. It is called by the builders before save.
|
||||
usagelog.ImageSizeValidator = usagelogDescImageSize.Validators[0].(func(string) error)
|
||||
// usagelogDescImageInputSize is the schema descriptor for image_input_size field.
|
||||
usagelogDescImageInputSize := usagelogFields[36].Descriptor()
|
||||
usagelogDescImageInputSize := usagelogFields[38].Descriptor()
|
||||
// usagelog.ImageInputSizeValidator is a validator for the "image_input_size" field. It is called by the builders before save.
|
||||
usagelog.ImageInputSizeValidator = usagelogDescImageInputSize.Validators[0].(func(string) error)
|
||||
// usagelogDescImageOutputSize is the schema descriptor for image_output_size field.
|
||||
usagelogDescImageOutputSize := usagelogFields[37].Descriptor()
|
||||
usagelogDescImageOutputSize := usagelogFields[39].Descriptor()
|
||||
// usagelog.ImageOutputSizeValidator is a validator for the "image_output_size" field. It is called by the builders before save.
|
||||
usagelog.ImageOutputSizeValidator = usagelogDescImageOutputSize.Validators[0].(func(string) error)
|
||||
// usagelogDescImageSizeSource is the schema descriptor for image_size_source field.
|
||||
usagelogDescImageSizeSource := usagelogFields[38].Descriptor()
|
||||
usagelogDescImageSizeSource := usagelogFields[40].Descriptor()
|
||||
// usagelog.ImageSizeSourceValidator is a validator for the "image_size_source" field. It is called by the builders before save.
|
||||
usagelog.ImageSizeSourceValidator = usagelogDescImageSizeSource.Validators[0].(func(string) error)
|
||||
// usagelogDescVideoCount is the schema descriptor for video_count field.
|
||||
usagelogDescVideoCount := usagelogFields[40].Descriptor()
|
||||
usagelogDescVideoCount := usagelogFields[42].Descriptor()
|
||||
// usagelog.DefaultVideoCount holds the default value on creation for the video_count field.
|
||||
usagelog.DefaultVideoCount = usagelogDescVideoCount.Default.(int)
|
||||
// usagelogDescVideoResolution is the schema descriptor for video_resolution field.
|
||||
usagelogDescVideoResolution := usagelogFields[41].Descriptor()
|
||||
usagelogDescVideoResolution := usagelogFields[43].Descriptor()
|
||||
// usagelog.VideoResolutionValidator is a validator for the "video_resolution" field. It is called by the builders before save.
|
||||
usagelog.VideoResolutionValidator = usagelogDescVideoResolution.Validators[0].(func(string) error)
|
||||
// usagelogDescCacheTTLOverridden is the schema descriptor for cache_ttl_overridden field.
|
||||
usagelogDescCacheTTLOverridden := usagelogFields[43].Descriptor()
|
||||
usagelogDescCacheTTLOverridden := usagelogFields[45].Descriptor()
|
||||
// usagelog.DefaultCacheTTLOverridden holds the default value on creation for the cache_ttl_overridden field.
|
||||
usagelog.DefaultCacheTTLOverridden = usagelogDescCacheTTLOverridden.Default.(bool)
|
||||
// usagelogDescCreatedAt is the schema descriptor for created_at field.
|
||||
usagelogDescCreatedAt := usagelogFields[44].Descriptor()
|
||||
usagelogDescCreatedAt := usagelogFields[46].Descriptor()
|
||||
// usagelog.DefaultCreatedAt holds the default value on creation for the created_at field.
|
||||
usagelog.DefaultCreatedAt = usagelogDescCreatedAt.Default.(func() time.Time)
|
||||
userMixin := schema.User{}.Mixin()
|
||||
|
||||
@@ -53,6 +53,17 @@ func (UsageLog) Fields() []ent.Field {
|
||||
MaxLen(100).
|
||||
Optional().
|
||||
Nillable(),
|
||||
// UpstreamResponseModel stores the model name declared by the upstream
|
||||
// response before any protocol conversion or client-facing rewrite.
|
||||
field.String("upstream_response_model").
|
||||
MaxLen(200).
|
||||
Optional().
|
||||
Nillable(),
|
||||
// UpstreamModelMismatch is tri-state: NULL means the upstream response did
|
||||
// not declare a model (or predates this field); false/true means observed.
|
||||
field.Bool("upstream_model_mismatch").
|
||||
Optional().
|
||||
Nillable(),
|
||||
field.Int64("channel_id").Optional().Nillable().Comment("渠道 ID"),
|
||||
field.String("model_mapping_chain").MaxLen(500).Optional().Nillable().Comment("模型映射链"),
|
||||
field.String("billing_tier").MaxLen(50).Optional().Nillable().Comment("计费层级标签"),
|
||||
|
||||
+30
-2
@@ -37,6 +37,10 @@ type UsageLog struct {
|
||||
RequestedModel *string `json:"requested_model,omitempty"`
|
||||
// UpstreamModel holds the value of the "upstream_model" field.
|
||||
UpstreamModel *string `json:"upstream_model,omitempty"`
|
||||
// UpstreamResponseModel holds the value of the "upstream_response_model" field.
|
||||
UpstreamResponseModel *string `json:"upstream_response_model,omitempty"`
|
||||
// UpstreamModelMismatch holds the value of the "upstream_model_mismatch" field.
|
||||
UpstreamModelMismatch *bool `json:"upstream_model_mismatch,omitempty"`
|
||||
// 渠道 ID
|
||||
ChannelID *int64 `json:"channel_id,omitempty"`
|
||||
// 模型映射链
|
||||
@@ -198,13 +202,13 @@ func (*UsageLog) scanValues(columns []string) ([]any, error) {
|
||||
switch columns[i] {
|
||||
case usagelog.FieldImageSizeBreakdown:
|
||||
values[i] = new([]byte)
|
||||
case usagelog.FieldLongContextBillingApplied, usagelog.FieldStream, usagelog.FieldCacheTTLOverridden:
|
||||
case usagelog.FieldUpstreamModelMismatch, usagelog.FieldLongContextBillingApplied, usagelog.FieldStream, usagelog.FieldCacheTTLOverridden:
|
||||
values[i] = new(sql.NullBool)
|
||||
case usagelog.FieldInputCost, usagelog.FieldOutputCost, usagelog.FieldCacheCreationCost, usagelog.FieldCacheReadCost, usagelog.FieldTotalCost, usagelog.FieldActualCost, usagelog.FieldRateMultiplier, usagelog.FieldAccountRateMultiplier:
|
||||
values[i] = new(sql.NullFloat64)
|
||||
case usagelog.FieldID, usagelog.FieldUserID, usagelog.FieldAPIKeyID, usagelog.FieldAccountID, usagelog.FieldChannelID, usagelog.FieldGroupID, usagelog.FieldSubscriptionID, usagelog.FieldInputTokens, usagelog.FieldOutputTokens, usagelog.FieldCacheCreationTokens, usagelog.FieldCacheReadTokens, usagelog.FieldCacheCreation5mTokens, usagelog.FieldCacheCreation1hTokens, usagelog.FieldBillingType, usagelog.FieldDurationMs, usagelog.FieldFirstTokenMs, usagelog.FieldImageCount, usagelog.FieldVideoCount, usagelog.FieldVideoDurationSeconds:
|
||||
values[i] = new(sql.NullInt64)
|
||||
case usagelog.FieldRequestID, usagelog.FieldModel, usagelog.FieldRequestedModel, usagelog.FieldUpstreamModel, usagelog.FieldModelMappingChain, usagelog.FieldBillingTier, usagelog.FieldBillingMode, usagelog.FieldUserAgent, usagelog.FieldIPAddress, usagelog.FieldImageSize, usagelog.FieldImageInputSize, usagelog.FieldImageOutputSize, usagelog.FieldImageSizeSource, usagelog.FieldVideoResolution:
|
||||
case usagelog.FieldRequestID, usagelog.FieldModel, usagelog.FieldRequestedModel, usagelog.FieldUpstreamModel, usagelog.FieldUpstreamResponseModel, usagelog.FieldModelMappingChain, usagelog.FieldBillingTier, usagelog.FieldBillingMode, usagelog.FieldUserAgent, usagelog.FieldIPAddress, usagelog.FieldImageSize, usagelog.FieldImageInputSize, usagelog.FieldImageOutputSize, usagelog.FieldImageSizeSource, usagelog.FieldVideoResolution:
|
||||
values[i] = new(sql.NullString)
|
||||
case usagelog.FieldCreatedAt:
|
||||
values[i] = new(sql.NullTime)
|
||||
@@ -273,6 +277,20 @@ func (_m *UsageLog) assignValues(columns []string, values []any) error {
|
||||
_m.UpstreamModel = new(string)
|
||||
*_m.UpstreamModel = value.String
|
||||
}
|
||||
case usagelog.FieldUpstreamResponseModel:
|
||||
if value, ok := values[i].(*sql.NullString); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field upstream_response_model", values[i])
|
||||
} else if value.Valid {
|
||||
_m.UpstreamResponseModel = new(string)
|
||||
*_m.UpstreamResponseModel = value.String
|
||||
}
|
||||
case usagelog.FieldUpstreamModelMismatch:
|
||||
if value, ok := values[i].(*sql.NullBool); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field upstream_model_mismatch", values[i])
|
||||
} else if value.Valid {
|
||||
_m.UpstreamModelMismatch = new(bool)
|
||||
*_m.UpstreamModelMismatch = value.Bool
|
||||
}
|
||||
case usagelog.FieldChannelID:
|
||||
if value, ok := values[i].(*sql.NullInt64); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field channel_id", values[i])
|
||||
@@ -606,6 +624,16 @@ func (_m *UsageLog) String() string {
|
||||
builder.WriteString(*v)
|
||||
}
|
||||
builder.WriteString(", ")
|
||||
if v := _m.UpstreamResponseModel; v != nil {
|
||||
builder.WriteString("upstream_response_model=")
|
||||
builder.WriteString(*v)
|
||||
}
|
||||
builder.WriteString(", ")
|
||||
if v := _m.UpstreamModelMismatch; v != nil {
|
||||
builder.WriteString("upstream_model_mismatch=")
|
||||
builder.WriteString(fmt.Sprintf("%v", *v))
|
||||
}
|
||||
builder.WriteString(", ")
|
||||
if v := _m.ChannelID; v != nil {
|
||||
builder.WriteString("channel_id=")
|
||||
builder.WriteString(fmt.Sprintf("%v", *v))
|
||||
|
||||
@@ -28,6 +28,10 @@ const (
|
||||
FieldRequestedModel = "requested_model"
|
||||
// FieldUpstreamModel holds the string denoting the upstream_model field in the database.
|
||||
FieldUpstreamModel = "upstream_model"
|
||||
// FieldUpstreamResponseModel holds the string denoting the upstream_response_model field in the database.
|
||||
FieldUpstreamResponseModel = "upstream_response_model"
|
||||
// FieldUpstreamModelMismatch holds the string denoting the upstream_model_mismatch field in the database.
|
||||
FieldUpstreamModelMismatch = "upstream_model_mismatch"
|
||||
// FieldChannelID holds the string denoting the channel_id field in the database.
|
||||
FieldChannelID = "channel_id"
|
||||
// FieldModelMappingChain holds the string denoting the model_mapping_chain field in the database.
|
||||
@@ -163,6 +167,8 @@ var Columns = []string{
|
||||
FieldModel,
|
||||
FieldRequestedModel,
|
||||
FieldUpstreamModel,
|
||||
FieldUpstreamResponseModel,
|
||||
FieldUpstreamModelMismatch,
|
||||
FieldChannelID,
|
||||
FieldModelMappingChain,
|
||||
FieldBillingTier,
|
||||
@@ -222,6 +228,8 @@ var (
|
||||
RequestedModelValidator func(string) error
|
||||
// UpstreamModelValidator is a validator for the "upstream_model" field. It is called by the builders before save.
|
||||
UpstreamModelValidator func(string) error
|
||||
// UpstreamResponseModelValidator is a validator for the "upstream_response_model" field. It is called by the builders before save.
|
||||
UpstreamResponseModelValidator func(string) error
|
||||
// ModelMappingChainValidator is a validator for the "model_mapping_chain" field. It is called by the builders before save.
|
||||
ModelMappingChainValidator func(string) error
|
||||
// BillingTierValidator is a validator for the "billing_tier" field. It is called by the builders before save.
|
||||
@@ -327,6 +335,16 @@ func ByUpstreamModel(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldUpstreamModel, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByUpstreamResponseModel orders the results by the upstream_response_model field.
|
||||
func ByUpstreamResponseModel(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldUpstreamResponseModel, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByUpstreamModelMismatch orders the results by the upstream_model_mismatch field.
|
||||
func ByUpstreamModelMismatch(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldUpstreamModelMismatch, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByChannelID orders the results by the channel_id field.
|
||||
func ByChannelID(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldChannelID, opts...).ToFunc()
|
||||
|
||||
@@ -90,6 +90,16 @@ func UpstreamModel(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldEQ(FieldUpstreamModel, v))
|
||||
}
|
||||
|
||||
// UpstreamResponseModel applies equality check predicate on the "upstream_response_model" field. It's identical to UpstreamResponseModelEQ.
|
||||
func UpstreamResponseModel(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldEQ(FieldUpstreamResponseModel, v))
|
||||
}
|
||||
|
||||
// UpstreamModelMismatch applies equality check predicate on the "upstream_model_mismatch" field. It's identical to UpstreamModelMismatchEQ.
|
||||
func UpstreamModelMismatch(v bool) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldEQ(FieldUpstreamModelMismatch, v))
|
||||
}
|
||||
|
||||
// ChannelID applies equality check predicate on the "channel_id" field. It's identical to ChannelIDEQ.
|
||||
func ChannelID(v int64) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldEQ(FieldChannelID, v))
|
||||
@@ -615,6 +625,101 @@ func UpstreamModelContainsFold(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldContainsFold(FieldUpstreamModel, v))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelEQ applies the EQ predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelEQ(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldEQ(FieldUpstreamResponseModel, v))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelNEQ applies the NEQ predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelNEQ(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldNEQ(FieldUpstreamResponseModel, v))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelIn applies the In predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelIn(vs ...string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldIn(FieldUpstreamResponseModel, vs...))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelNotIn applies the NotIn predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelNotIn(vs ...string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldNotIn(FieldUpstreamResponseModel, vs...))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelGT applies the GT predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelGT(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldGT(FieldUpstreamResponseModel, v))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelGTE applies the GTE predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelGTE(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldGTE(FieldUpstreamResponseModel, v))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelLT applies the LT predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelLT(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldLT(FieldUpstreamResponseModel, v))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelLTE applies the LTE predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelLTE(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldLTE(FieldUpstreamResponseModel, v))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelContains applies the Contains predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelContains(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldContains(FieldUpstreamResponseModel, v))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelHasPrefix applies the HasPrefix predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelHasPrefix(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldHasPrefix(FieldUpstreamResponseModel, v))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelHasSuffix applies the HasSuffix predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelHasSuffix(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldHasSuffix(FieldUpstreamResponseModel, v))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelIsNil applies the IsNil predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelIsNil() predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldIsNull(FieldUpstreamResponseModel))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelNotNil applies the NotNil predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelNotNil() predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldNotNull(FieldUpstreamResponseModel))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelEqualFold applies the EqualFold predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelEqualFold(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldEqualFold(FieldUpstreamResponseModel, v))
|
||||
}
|
||||
|
||||
// UpstreamResponseModelContainsFold applies the ContainsFold predicate on the "upstream_response_model" field.
|
||||
func UpstreamResponseModelContainsFold(v string) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldContainsFold(FieldUpstreamResponseModel, v))
|
||||
}
|
||||
|
||||
// UpstreamModelMismatchEQ applies the EQ predicate on the "upstream_model_mismatch" field.
|
||||
func UpstreamModelMismatchEQ(v bool) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldEQ(FieldUpstreamModelMismatch, v))
|
||||
}
|
||||
|
||||
// UpstreamModelMismatchNEQ applies the NEQ predicate on the "upstream_model_mismatch" field.
|
||||
func UpstreamModelMismatchNEQ(v bool) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldNEQ(FieldUpstreamModelMismatch, v))
|
||||
}
|
||||
|
||||
// UpstreamModelMismatchIsNil applies the IsNil predicate on the "upstream_model_mismatch" field.
|
||||
func UpstreamModelMismatchIsNil() predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldIsNull(FieldUpstreamModelMismatch))
|
||||
}
|
||||
|
||||
// UpstreamModelMismatchNotNil applies the NotNil predicate on the "upstream_model_mismatch" field.
|
||||
func UpstreamModelMismatchNotNil() predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldNotNull(FieldUpstreamModelMismatch))
|
||||
}
|
||||
|
||||
// ChannelIDEQ applies the EQ predicate on the "channel_id" field.
|
||||
func ChannelIDEQ(v int64) predicate.UsageLog {
|
||||
return predicate.UsageLog(sql.FieldEQ(FieldChannelID, v))
|
||||
|
||||
@@ -85,6 +85,34 @@ func (_c *UsageLogCreate) SetNillableUpstreamModel(v *string) *UsageLogCreate {
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetUpstreamResponseModel sets the "upstream_response_model" field.
|
||||
func (_c *UsageLogCreate) SetUpstreamResponseModel(v string) *UsageLogCreate {
|
||||
_c.mutation.SetUpstreamResponseModel(v)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetNillableUpstreamResponseModel sets the "upstream_response_model" field if the given value is not nil.
|
||||
func (_c *UsageLogCreate) SetNillableUpstreamResponseModel(v *string) *UsageLogCreate {
|
||||
if v != nil {
|
||||
_c.SetUpstreamResponseModel(*v)
|
||||
}
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetUpstreamModelMismatch sets the "upstream_model_mismatch" field.
|
||||
func (_c *UsageLogCreate) SetUpstreamModelMismatch(v bool) *UsageLogCreate {
|
||||
_c.mutation.SetUpstreamModelMismatch(v)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetNillableUpstreamModelMismatch sets the "upstream_model_mismatch" field if the given value is not nil.
|
||||
func (_c *UsageLogCreate) SetNillableUpstreamModelMismatch(v *bool) *UsageLogCreate {
|
||||
if v != nil {
|
||||
_c.SetUpstreamModelMismatch(*v)
|
||||
}
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetChannelID sets the "channel_id" field.
|
||||
func (_c *UsageLogCreate) SetChannelID(v int64) *UsageLogCreate {
|
||||
_c.mutation.SetChannelID(v)
|
||||
@@ -788,6 +816,11 @@ func (_c *UsageLogCreate) check() error {
|
||||
return &ValidationError{Name: "upstream_model", err: fmt.Errorf(`ent: validator failed for field "UsageLog.upstream_model": %w`, err)}
|
||||
}
|
||||
}
|
||||
if v, ok := _c.mutation.UpstreamResponseModel(); ok {
|
||||
if err := usagelog.UpstreamResponseModelValidator(v); err != nil {
|
||||
return &ValidationError{Name: "upstream_response_model", err: fmt.Errorf(`ent: validator failed for field "UsageLog.upstream_response_model": %w`, err)}
|
||||
}
|
||||
}
|
||||
if v, ok := _c.mutation.ModelMappingChain(); ok {
|
||||
if err := usagelog.ModelMappingChainValidator(v); err != nil {
|
||||
return &ValidationError{Name: "model_mapping_chain", err: fmt.Errorf(`ent: validator failed for field "UsageLog.model_mapping_chain": %w`, err)}
|
||||
@@ -950,6 +983,14 @@ func (_c *UsageLogCreate) createSpec() (*UsageLog, *sqlgraph.CreateSpec) {
|
||||
_spec.SetField(usagelog.FieldUpstreamModel, field.TypeString, value)
|
||||
_node.UpstreamModel = &value
|
||||
}
|
||||
if value, ok := _c.mutation.UpstreamResponseModel(); ok {
|
||||
_spec.SetField(usagelog.FieldUpstreamResponseModel, field.TypeString, value)
|
||||
_node.UpstreamResponseModel = &value
|
||||
}
|
||||
if value, ok := _c.mutation.UpstreamModelMismatch(); ok {
|
||||
_spec.SetField(usagelog.FieldUpstreamModelMismatch, field.TypeBool, value)
|
||||
_node.UpstreamModelMismatch = &value
|
||||
}
|
||||
if value, ok := _c.mutation.ChannelID(); ok {
|
||||
_spec.SetField(usagelog.FieldChannelID, field.TypeInt64, value)
|
||||
_node.ChannelID = &value
|
||||
@@ -1327,6 +1368,42 @@ func (u *UsageLogUpsert) ClearUpstreamModel() *UsageLogUpsert {
|
||||
return u
|
||||
}
|
||||
|
||||
// SetUpstreamResponseModel sets the "upstream_response_model" field.
|
||||
func (u *UsageLogUpsert) SetUpstreamResponseModel(v string) *UsageLogUpsert {
|
||||
u.Set(usagelog.FieldUpstreamResponseModel, v)
|
||||
return u
|
||||
}
|
||||
|
||||
// UpdateUpstreamResponseModel sets the "upstream_response_model" field to the value that was provided on create.
|
||||
func (u *UsageLogUpsert) UpdateUpstreamResponseModel() *UsageLogUpsert {
|
||||
u.SetExcluded(usagelog.FieldUpstreamResponseModel)
|
||||
return u
|
||||
}
|
||||
|
||||
// ClearUpstreamResponseModel clears the value of the "upstream_response_model" field.
|
||||
func (u *UsageLogUpsert) ClearUpstreamResponseModel() *UsageLogUpsert {
|
||||
u.SetNull(usagelog.FieldUpstreamResponseModel)
|
||||
return u
|
||||
}
|
||||
|
||||
// SetUpstreamModelMismatch sets the "upstream_model_mismatch" field.
|
||||
func (u *UsageLogUpsert) SetUpstreamModelMismatch(v bool) *UsageLogUpsert {
|
||||
u.Set(usagelog.FieldUpstreamModelMismatch, v)
|
||||
return u
|
||||
}
|
||||
|
||||
// UpdateUpstreamModelMismatch sets the "upstream_model_mismatch" field to the value that was provided on create.
|
||||
func (u *UsageLogUpsert) UpdateUpstreamModelMismatch() *UsageLogUpsert {
|
||||
u.SetExcluded(usagelog.FieldUpstreamModelMismatch)
|
||||
return u
|
||||
}
|
||||
|
||||
// ClearUpstreamModelMismatch clears the value of the "upstream_model_mismatch" field.
|
||||
func (u *UsageLogUpsert) ClearUpstreamModelMismatch() *UsageLogUpsert {
|
||||
u.SetNull(usagelog.FieldUpstreamModelMismatch)
|
||||
return u
|
||||
}
|
||||
|
||||
// SetChannelID sets the "channel_id" field.
|
||||
func (u *UsageLogUpsert) SetChannelID(v int64) *UsageLogUpsert {
|
||||
u.Set(usagelog.FieldChannelID, v)
|
||||
@@ -2162,6 +2239,48 @@ func (u *UsageLogUpsertOne) ClearUpstreamModel() *UsageLogUpsertOne {
|
||||
})
|
||||
}
|
||||
|
||||
// SetUpstreamResponseModel sets the "upstream_response_model" field.
|
||||
func (u *UsageLogUpsertOne) SetUpstreamResponseModel(v string) *UsageLogUpsertOne {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.SetUpstreamResponseModel(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateUpstreamResponseModel sets the "upstream_response_model" field to the value that was provided on create.
|
||||
func (u *UsageLogUpsertOne) UpdateUpstreamResponseModel() *UsageLogUpsertOne {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.UpdateUpstreamResponseModel()
|
||||
})
|
||||
}
|
||||
|
||||
// ClearUpstreamResponseModel clears the value of the "upstream_response_model" field.
|
||||
func (u *UsageLogUpsertOne) ClearUpstreamResponseModel() *UsageLogUpsertOne {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.ClearUpstreamResponseModel()
|
||||
})
|
||||
}
|
||||
|
||||
// SetUpstreamModelMismatch sets the "upstream_model_mismatch" field.
|
||||
func (u *UsageLogUpsertOne) SetUpstreamModelMismatch(v bool) *UsageLogUpsertOne {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.SetUpstreamModelMismatch(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateUpstreamModelMismatch sets the "upstream_model_mismatch" field to the value that was provided on create.
|
||||
func (u *UsageLogUpsertOne) UpdateUpstreamModelMismatch() *UsageLogUpsertOne {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.UpdateUpstreamModelMismatch()
|
||||
})
|
||||
}
|
||||
|
||||
// ClearUpstreamModelMismatch clears the value of the "upstream_model_mismatch" field.
|
||||
func (u *UsageLogUpsertOne) ClearUpstreamModelMismatch() *UsageLogUpsertOne {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.ClearUpstreamModelMismatch()
|
||||
})
|
||||
}
|
||||
|
||||
// SetChannelID sets the "channel_id" field.
|
||||
func (u *UsageLogUpsertOne) SetChannelID(v int64) *UsageLogUpsertOne {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
@@ -3276,6 +3395,48 @@ func (u *UsageLogUpsertBulk) ClearUpstreamModel() *UsageLogUpsertBulk {
|
||||
})
|
||||
}
|
||||
|
||||
// SetUpstreamResponseModel sets the "upstream_response_model" field.
|
||||
func (u *UsageLogUpsertBulk) SetUpstreamResponseModel(v string) *UsageLogUpsertBulk {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.SetUpstreamResponseModel(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateUpstreamResponseModel sets the "upstream_response_model" field to the value that was provided on create.
|
||||
func (u *UsageLogUpsertBulk) UpdateUpstreamResponseModel() *UsageLogUpsertBulk {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.UpdateUpstreamResponseModel()
|
||||
})
|
||||
}
|
||||
|
||||
// ClearUpstreamResponseModel clears the value of the "upstream_response_model" field.
|
||||
func (u *UsageLogUpsertBulk) ClearUpstreamResponseModel() *UsageLogUpsertBulk {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.ClearUpstreamResponseModel()
|
||||
})
|
||||
}
|
||||
|
||||
// SetUpstreamModelMismatch sets the "upstream_model_mismatch" field.
|
||||
func (u *UsageLogUpsertBulk) SetUpstreamModelMismatch(v bool) *UsageLogUpsertBulk {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.SetUpstreamModelMismatch(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateUpstreamModelMismatch sets the "upstream_model_mismatch" field to the value that was provided on create.
|
||||
func (u *UsageLogUpsertBulk) UpdateUpstreamModelMismatch() *UsageLogUpsertBulk {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.UpdateUpstreamModelMismatch()
|
||||
})
|
||||
}
|
||||
|
||||
// ClearUpstreamModelMismatch clears the value of the "upstream_model_mismatch" field.
|
||||
func (u *UsageLogUpsertBulk) ClearUpstreamModelMismatch() *UsageLogUpsertBulk {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
s.ClearUpstreamModelMismatch()
|
||||
})
|
||||
}
|
||||
|
||||
// SetChannelID sets the "channel_id" field.
|
||||
func (u *UsageLogUpsertBulk) SetChannelID(v int64) *UsageLogUpsertBulk {
|
||||
return u.Update(func(s *UsageLogUpsert) {
|
||||
|
||||
@@ -142,6 +142,46 @@ func (_u *UsageLogUpdate) ClearUpstreamModel() *UsageLogUpdate {
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetUpstreamResponseModel sets the "upstream_response_model" field.
|
||||
func (_u *UsageLogUpdate) SetUpstreamResponseModel(v string) *UsageLogUpdate {
|
||||
_u.mutation.SetUpstreamResponseModel(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableUpstreamResponseModel sets the "upstream_response_model" field if the given value is not nil.
|
||||
func (_u *UsageLogUpdate) SetNillableUpstreamResponseModel(v *string) *UsageLogUpdate {
|
||||
if v != nil {
|
||||
_u.SetUpstreamResponseModel(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
|
||||
// ClearUpstreamResponseModel clears the value of the "upstream_response_model" field.
|
||||
func (_u *UsageLogUpdate) ClearUpstreamResponseModel() *UsageLogUpdate {
|
||||
_u.mutation.ClearUpstreamResponseModel()
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetUpstreamModelMismatch sets the "upstream_model_mismatch" field.
|
||||
func (_u *UsageLogUpdate) SetUpstreamModelMismatch(v bool) *UsageLogUpdate {
|
||||
_u.mutation.SetUpstreamModelMismatch(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableUpstreamModelMismatch sets the "upstream_model_mismatch" field if the given value is not nil.
|
||||
func (_u *UsageLogUpdate) SetNillableUpstreamModelMismatch(v *bool) *UsageLogUpdate {
|
||||
if v != nil {
|
||||
_u.SetUpstreamModelMismatch(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
|
||||
// ClearUpstreamModelMismatch clears the value of the "upstream_model_mismatch" field.
|
||||
func (_u *UsageLogUpdate) ClearUpstreamModelMismatch() *UsageLogUpdate {
|
||||
_u.mutation.ClearUpstreamModelMismatch()
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetChannelID sets the "channel_id" field.
|
||||
func (_u *UsageLogUpdate) SetChannelID(v int64) *UsageLogUpdate {
|
||||
_u.mutation.ResetChannelID()
|
||||
@@ -1016,6 +1056,11 @@ func (_u *UsageLogUpdate) check() error {
|
||||
return &ValidationError{Name: "upstream_model", err: fmt.Errorf(`ent: validator failed for field "UsageLog.upstream_model": %w`, err)}
|
||||
}
|
||||
}
|
||||
if v, ok := _u.mutation.UpstreamResponseModel(); ok {
|
||||
if err := usagelog.UpstreamResponseModelValidator(v); err != nil {
|
||||
return &ValidationError{Name: "upstream_response_model", err: fmt.Errorf(`ent: validator failed for field "UsageLog.upstream_response_model": %w`, err)}
|
||||
}
|
||||
}
|
||||
if v, ok := _u.mutation.ModelMappingChain(); ok {
|
||||
if err := usagelog.ModelMappingChainValidator(v); err != nil {
|
||||
return &ValidationError{Name: "model_mapping_chain", err: fmt.Errorf(`ent: validator failed for field "UsageLog.model_mapping_chain": %w`, err)}
|
||||
@@ -1108,6 +1153,18 @@ func (_u *UsageLogUpdate) sqlSave(ctx context.Context) (_node int, err error) {
|
||||
if _u.mutation.UpstreamModelCleared() {
|
||||
_spec.ClearField(usagelog.FieldUpstreamModel, field.TypeString)
|
||||
}
|
||||
if value, ok := _u.mutation.UpstreamResponseModel(); ok {
|
||||
_spec.SetField(usagelog.FieldUpstreamResponseModel, field.TypeString, value)
|
||||
}
|
||||
if _u.mutation.UpstreamResponseModelCleared() {
|
||||
_spec.ClearField(usagelog.FieldUpstreamResponseModel, field.TypeString)
|
||||
}
|
||||
if value, ok := _u.mutation.UpstreamModelMismatch(); ok {
|
||||
_spec.SetField(usagelog.FieldUpstreamModelMismatch, field.TypeBool, value)
|
||||
}
|
||||
if _u.mutation.UpstreamModelMismatchCleared() {
|
||||
_spec.ClearField(usagelog.FieldUpstreamModelMismatch, field.TypeBool)
|
||||
}
|
||||
if value, ok := _u.mutation.ChannelID(); ok {
|
||||
_spec.SetField(usagelog.FieldChannelID, field.TypeInt64, value)
|
||||
}
|
||||
@@ -1599,6 +1656,46 @@ func (_u *UsageLogUpdateOne) ClearUpstreamModel() *UsageLogUpdateOne {
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetUpstreamResponseModel sets the "upstream_response_model" field.
|
||||
func (_u *UsageLogUpdateOne) SetUpstreamResponseModel(v string) *UsageLogUpdateOne {
|
||||
_u.mutation.SetUpstreamResponseModel(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableUpstreamResponseModel sets the "upstream_response_model" field if the given value is not nil.
|
||||
func (_u *UsageLogUpdateOne) SetNillableUpstreamResponseModel(v *string) *UsageLogUpdateOne {
|
||||
if v != nil {
|
||||
_u.SetUpstreamResponseModel(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
|
||||
// ClearUpstreamResponseModel clears the value of the "upstream_response_model" field.
|
||||
func (_u *UsageLogUpdateOne) ClearUpstreamResponseModel() *UsageLogUpdateOne {
|
||||
_u.mutation.ClearUpstreamResponseModel()
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetUpstreamModelMismatch sets the "upstream_model_mismatch" field.
|
||||
func (_u *UsageLogUpdateOne) SetUpstreamModelMismatch(v bool) *UsageLogUpdateOne {
|
||||
_u.mutation.SetUpstreamModelMismatch(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableUpstreamModelMismatch sets the "upstream_model_mismatch" field if the given value is not nil.
|
||||
func (_u *UsageLogUpdateOne) SetNillableUpstreamModelMismatch(v *bool) *UsageLogUpdateOne {
|
||||
if v != nil {
|
||||
_u.SetUpstreamModelMismatch(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
|
||||
// ClearUpstreamModelMismatch clears the value of the "upstream_model_mismatch" field.
|
||||
func (_u *UsageLogUpdateOne) ClearUpstreamModelMismatch() *UsageLogUpdateOne {
|
||||
_u.mutation.ClearUpstreamModelMismatch()
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetChannelID sets the "channel_id" field.
|
||||
func (_u *UsageLogUpdateOne) SetChannelID(v int64) *UsageLogUpdateOne {
|
||||
_u.mutation.ResetChannelID()
|
||||
@@ -2486,6 +2583,11 @@ func (_u *UsageLogUpdateOne) check() error {
|
||||
return &ValidationError{Name: "upstream_model", err: fmt.Errorf(`ent: validator failed for field "UsageLog.upstream_model": %w`, err)}
|
||||
}
|
||||
}
|
||||
if v, ok := _u.mutation.UpstreamResponseModel(); ok {
|
||||
if err := usagelog.UpstreamResponseModelValidator(v); err != nil {
|
||||
return &ValidationError{Name: "upstream_response_model", err: fmt.Errorf(`ent: validator failed for field "UsageLog.upstream_response_model": %w`, err)}
|
||||
}
|
||||
}
|
||||
if v, ok := _u.mutation.ModelMappingChain(); ok {
|
||||
if err := usagelog.ModelMappingChainValidator(v); err != nil {
|
||||
return &ValidationError{Name: "model_mapping_chain", err: fmt.Errorf(`ent: validator failed for field "UsageLog.model_mapping_chain": %w`, err)}
|
||||
@@ -2595,6 +2697,18 @@ func (_u *UsageLogUpdateOne) sqlSave(ctx context.Context) (_node *UsageLog, err
|
||||
if _u.mutation.UpstreamModelCleared() {
|
||||
_spec.ClearField(usagelog.FieldUpstreamModel, field.TypeString)
|
||||
}
|
||||
if value, ok := _u.mutation.UpstreamResponseModel(); ok {
|
||||
_spec.SetField(usagelog.FieldUpstreamResponseModel, field.TypeString, value)
|
||||
}
|
||||
if _u.mutation.UpstreamResponseModelCleared() {
|
||||
_spec.ClearField(usagelog.FieldUpstreamResponseModel, field.TypeString)
|
||||
}
|
||||
if value, ok := _u.mutation.UpstreamModelMismatch(); ok {
|
||||
_spec.SetField(usagelog.FieldUpstreamModelMismatch, field.TypeBool, value)
|
||||
}
|
||||
if _u.mutation.UpstreamModelMismatchCleared() {
|
||||
_spec.ClearField(usagelog.FieldUpstreamModelMismatch, field.TypeBool)
|
||||
}
|
||||
if value, ok := _u.mutation.ChannelID(); ok {
|
||||
_spec.SetField(usagelog.FieldChannelID, field.TypeInt64, value)
|
||||
}
|
||||
|
||||
@@ -117,6 +117,12 @@ var DefaultAntigravityModelMapping = map[string]string{
|
||||
"gemini-3.1-flash-image": "gemini-3.1-flash-image",
|
||||
// Gemini 3.1 image preview 映射
|
||||
"gemini-3.1-flash-image-preview": "gemini-3.1-flash-image",
|
||||
// Gemini 3.6 Flash tiered models
|
||||
"gemini-3.6-flash": "gemini-3.6-flash",
|
||||
"gemini-3.6-flash-high": "gemini-3.6-flash-high",
|
||||
"gemini-3.6-flash-low": "gemini-3.6-flash-low",
|
||||
"gemini-3.6-flash-medium": "gemini-3.6-flash-medium",
|
||||
"gemini-3.6-flash-tiered": "gemini-3.6-flash-tiered",
|
||||
// Gemini 3 image 兼容映射(向 3.1 image 迁移)
|
||||
"gemini-3-pro-image": "gemini-3.1-flash-image",
|
||||
"gemini-3-pro-image-preview": "gemini-3.1-flash-image",
|
||||
|
||||
@@ -65,6 +65,14 @@ func TestDefaultAntigravityModelMapping_Gemini31ProAliases(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultAntigravityModelMapping_Gemini36FlashModels(t *testing.T) {
|
||||
for _, model := range []string{"gemini-3.6-flash", "gemini-3.6-flash-high", "gemini-3.6-flash-low", "gemini-3.6-flash-medium", "gemini-3.6-flash-tiered"} {
|
||||
if got := DefaultAntigravityModelMapping[model]; got != model {
|
||||
t.Fatalf("expected %s to map to itself, got %q", model, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultBedrockModelMapping_ContainsNewClaudeModels(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
@@ -64,6 +64,18 @@ func parseTimeRange(c *gin.Context) (time.Time, time.Time) {
|
||||
return startTime, endTime
|
||||
}
|
||||
|
||||
func parseOptionalBoolDashboardFilter(c *gin.Context, name string) (*bool, error) {
|
||||
raw := strings.TrimSpace(c.Query(name))
|
||||
if raw == "" {
|
||||
return nil, nil
|
||||
}
|
||||
value, err := strconv.ParseBool(raw)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &value, nil
|
||||
}
|
||||
|
||||
// GetStats handles getting dashboard statistics
|
||||
// GET /api/v1/admin/dashboard/stats
|
||||
func (h *DashboardHandler) GetStats(c *gin.Context) {
|
||||
@@ -200,6 +212,7 @@ func (h *DashboardHandler) GetUsageTrend(c *gin.Context) {
|
||||
var requestType *int16
|
||||
var stream *bool
|
||||
var billingType *int8
|
||||
var upstreamModelMismatch *bool
|
||||
|
||||
if userIDStr := c.Query("user_id"); userIDStr != "" {
|
||||
if id, err := strconv.ParseInt(userIDStr, 10, 64); err == nil {
|
||||
@@ -249,8 +262,13 @@ func (h *DashboardHandler) GetUsageTrend(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
}
|
||||
upstreamModelMismatch, err := parseOptionalBoolDashboardFilter(c, "upstream_model_mismatch")
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid upstream_model_mismatch value, use true or false")
|
||||
return
|
||||
}
|
||||
|
||||
trend, hit, err := h.getUsageTrendCached(c.Request.Context(), startTime, endTime, granularity, userID, apiKeyID, accountID, groupID, model, requestType, stream, billingType)
|
||||
trend, hit, err := h.getUsageTrendCached(c.Request.Context(), startTime, endTime, granularity, userID, apiKeyID, accountID, groupID, model, requestType, stream, billingType, upstreamModelMismatch)
|
||||
if err != nil {
|
||||
response.Error(c, 500, "Failed to get usage trend")
|
||||
return
|
||||
@@ -277,6 +295,7 @@ func (h *DashboardHandler) GetModelStats(c *gin.Context) {
|
||||
var requestType *int16
|
||||
var stream *bool
|
||||
var billingType *int8
|
||||
var upstreamModelMismatch *bool
|
||||
|
||||
if userIDStr := c.Query("user_id"); userIDStr != "" {
|
||||
if id, err := strconv.ParseInt(userIDStr, 10, 64); err == nil {
|
||||
@@ -330,8 +349,13 @@ func (h *DashboardHandler) GetModelStats(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
}
|
||||
upstreamModelMismatch, err := parseOptionalBoolDashboardFilter(c, "upstream_model_mismatch")
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid upstream_model_mismatch value, use true or false")
|
||||
return
|
||||
}
|
||||
|
||||
stats, hit, err := h.getModelStatsCached(c.Request.Context(), startTime, endTime, userID, apiKeyID, accountID, groupID, modelSource, requestType, stream, billingType)
|
||||
stats, hit, err := h.getModelStatsCached(c.Request.Context(), startTime, endTime, userID, apiKeyID, accountID, groupID, modelSource, requestType, stream, billingType, upstreamModelMismatch)
|
||||
if err != nil {
|
||||
response.Error(c, 500, "Failed to get model statistics")
|
||||
return
|
||||
@@ -355,6 +379,7 @@ func (h *DashboardHandler) GetGroupStats(c *gin.Context) {
|
||||
var requestType *int16
|
||||
var stream *bool
|
||||
var billingType *int8
|
||||
var upstreamModelMismatch *bool
|
||||
|
||||
if userIDStr := c.Query("user_id"); userIDStr != "" {
|
||||
if id, err := strconv.ParseInt(userIDStr, 10, 64); err == nil {
|
||||
@@ -401,8 +426,13 @@ func (h *DashboardHandler) GetGroupStats(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
}
|
||||
upstreamModelMismatch, err := parseOptionalBoolDashboardFilter(c, "upstream_model_mismatch")
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid upstream_model_mismatch value, use true or false")
|
||||
return
|
||||
}
|
||||
|
||||
stats, hit, err := h.getGroupStatsCached(c.Request.Context(), startTime, endTime, userID, apiKeyID, accountID, groupID, requestType, stream, billingType)
|
||||
stats, hit, err := h.getGroupStatsCached(c.Request.Context(), startTime, endTime, userID, apiKeyID, accountID, groupID, requestType, stream, billingType, upstreamModelMismatch)
|
||||
if err != nil {
|
||||
response.Error(c, 500, "Failed to get group statistics")
|
||||
return
|
||||
|
||||
@@ -19,11 +19,26 @@ type dashboardUsageRepoCapture struct {
|
||||
trendStream *bool
|
||||
modelRequestType *int16
|
||||
modelStream *bool
|
||||
trendMismatch *bool
|
||||
modelMismatch *bool
|
||||
groupMismatch *bool
|
||||
rankingLimit int
|
||||
ranking []usagestats.UserSpendingRankingItem
|
||||
rankingTotal float64
|
||||
}
|
||||
|
||||
func (s *dashboardUsageRepoCapture) GetUsageTrendWithUsageFilters(
|
||||
ctx context.Context,
|
||||
startTime, endTime time.Time,
|
||||
granularity string,
|
||||
filters usagestats.UsageLogFilters,
|
||||
) ([]usagestats.TrendDataPoint, error) {
|
||||
s.trendRequestType = filters.RequestType
|
||||
s.trendStream = filters.Stream
|
||||
s.trendMismatch = filters.UpstreamModelMismatch
|
||||
return []usagestats.TrendDataPoint{}, nil
|
||||
}
|
||||
|
||||
func (s *dashboardUsageRepoCapture) GetUsageTrendWithFilters(
|
||||
ctx context.Context,
|
||||
startTime, endTime time.Time,
|
||||
@@ -39,6 +54,27 @@ func (s *dashboardUsageRepoCapture) GetUsageTrendWithFilters(
|
||||
return []usagestats.TrendDataPoint{}, nil
|
||||
}
|
||||
|
||||
func (s *dashboardUsageRepoCapture) GetModelStatsWithUsageFiltersBySource(
|
||||
ctx context.Context,
|
||||
startTime, endTime time.Time,
|
||||
filters usagestats.UsageLogFilters,
|
||||
source string,
|
||||
) ([]usagestats.ModelStat, error) {
|
||||
s.modelRequestType = filters.RequestType
|
||||
s.modelStream = filters.Stream
|
||||
s.modelMismatch = filters.UpstreamModelMismatch
|
||||
return []usagestats.ModelStat{}, nil
|
||||
}
|
||||
|
||||
func (s *dashboardUsageRepoCapture) GetGroupStatsWithUsageFilters(
|
||||
ctx context.Context,
|
||||
startTime, endTime time.Time,
|
||||
filters usagestats.UsageLogFilters,
|
||||
) ([]usagestats.GroupStat, error) {
|
||||
s.groupMismatch = filters.UpstreamModelMismatch
|
||||
return []usagestats.GroupStat{}, nil
|
||||
}
|
||||
|
||||
func (s *dashboardUsageRepoCapture) GetModelStatsWithFilters(
|
||||
ctx context.Context,
|
||||
startTime, endTime time.Time,
|
||||
@@ -73,6 +109,7 @@ func newDashboardRequestTypeTestRouter(repo *dashboardUsageRepoCapture) *gin.Eng
|
||||
router := gin.New()
|
||||
router.GET("/admin/dashboard/trend", handler.GetUsageTrend)
|
||||
router.GET("/admin/dashboard/models", handler.GetModelStats)
|
||||
router.GET("/admin/dashboard/groups", handler.GetGroupStats)
|
||||
router.GET("/admin/dashboard/users-ranking", handler.GetUserSpendingRanking)
|
||||
return router
|
||||
}
|
||||
@@ -171,6 +208,46 @@ func TestDashboardModelStatsValidModelSource(t *testing.T) {
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
}
|
||||
|
||||
func TestDashboardModelAuditFilterPropagatesToTrendModelAndGroupQueries(t *testing.T) {
|
||||
resetDashboardReadCachesForTest()
|
||||
repo := &dashboardUsageRepoCapture{}
|
||||
router := newDashboardRequestTypeTestRouter(repo)
|
||||
|
||||
for _, path := range []string{
|
||||
"/admin/dashboard/trend?upstream_model_mismatch=true",
|
||||
"/admin/dashboard/models?upstream_model_mismatch=true",
|
||||
"/admin/dashboard/groups?upstream_model_mismatch=true",
|
||||
} {
|
||||
req := httptest.NewRequest(http.MethodGet, path, nil)
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
require.Equal(t, http.StatusOK, rec.Code, path)
|
||||
}
|
||||
|
||||
require.NotNil(t, repo.trendMismatch)
|
||||
require.True(t, *repo.trendMismatch)
|
||||
require.NotNil(t, repo.modelMismatch)
|
||||
require.True(t, *repo.modelMismatch)
|
||||
require.NotNil(t, repo.groupMismatch)
|
||||
require.True(t, *repo.groupMismatch)
|
||||
}
|
||||
|
||||
func TestDashboardModelAuditFilterRejectsInvalidBoolean(t *testing.T) {
|
||||
repo := &dashboardUsageRepoCapture{}
|
||||
router := newDashboardRequestTypeTestRouter(repo)
|
||||
|
||||
for _, path := range []string{
|
||||
"/admin/dashboard/trend?upstream_model_mismatch=invalid",
|
||||
"/admin/dashboard/models?upstream_model_mismatch=invalid",
|
||||
"/admin/dashboard/groups?upstream_model_mismatch=invalid",
|
||||
} {
|
||||
req := httptest.NewRequest(http.MethodGet, path, nil)
|
||||
rec := httptest.NewRecorder()
|
||||
router.ServeHTTP(rec, req)
|
||||
require.Equal(t, http.StatusBadRequest, rec.Code, path)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDashboardUsersRankingLimitAndCache(t *testing.T) {
|
||||
dashboardUsersRankingCache = newSnapshotCache(5 * time.Minute)
|
||||
repo := &dashboardUsageRepoCapture{
|
||||
|
||||
@@ -18,30 +18,32 @@ var (
|
||||
)
|
||||
|
||||
type dashboardTrendCacheKey struct {
|
||||
StartTime string `json:"start_time"`
|
||||
EndTime string `json:"end_time"`
|
||||
Granularity string `json:"granularity"`
|
||||
UserID int64 `json:"user_id"`
|
||||
APIKeyID int64 `json:"api_key_id"`
|
||||
AccountID int64 `json:"account_id"`
|
||||
GroupID int64 `json:"group_id"`
|
||||
Model string `json:"model"`
|
||||
RequestType *int16 `json:"request_type"`
|
||||
Stream *bool `json:"stream"`
|
||||
BillingType *int8 `json:"billing_type"`
|
||||
StartTime string `json:"start_time"`
|
||||
EndTime string `json:"end_time"`
|
||||
Granularity string `json:"granularity"`
|
||||
UserID int64 `json:"user_id"`
|
||||
APIKeyID int64 `json:"api_key_id"`
|
||||
AccountID int64 `json:"account_id"`
|
||||
GroupID int64 `json:"group_id"`
|
||||
Model string `json:"model"`
|
||||
RequestType *int16 `json:"request_type"`
|
||||
Stream *bool `json:"stream"`
|
||||
BillingType *int8 `json:"billing_type"`
|
||||
UpstreamModelMismatch *bool `json:"upstream_model_mismatch"`
|
||||
}
|
||||
|
||||
type dashboardModelGroupCacheKey struct {
|
||||
StartTime string `json:"start_time"`
|
||||
EndTime string `json:"end_time"`
|
||||
UserID int64 `json:"user_id"`
|
||||
APIKeyID int64 `json:"api_key_id"`
|
||||
AccountID int64 `json:"account_id"`
|
||||
GroupID int64 `json:"group_id"`
|
||||
ModelSource string `json:"model_source,omitempty"`
|
||||
RequestType *int16 `json:"request_type"`
|
||||
Stream *bool `json:"stream"`
|
||||
BillingType *int8 `json:"billing_type"`
|
||||
StartTime string `json:"start_time"`
|
||||
EndTime string `json:"end_time"`
|
||||
UserID int64 `json:"user_id"`
|
||||
APIKeyID int64 `json:"api_key_id"`
|
||||
AccountID int64 `json:"account_id"`
|
||||
GroupID int64 `json:"group_id"`
|
||||
ModelSource string `json:"model_source,omitempty"`
|
||||
RequestType *int16 `json:"request_type"`
|
||||
Stream *bool `json:"stream"`
|
||||
BillingType *int8 `json:"billing_type"`
|
||||
UpstreamModelMismatch *bool `json:"upstream_model_mismatch"`
|
||||
}
|
||||
|
||||
type dashboardEntityTrendCacheKey struct {
|
||||
@@ -84,22 +86,28 @@ func (h *DashboardHandler) getUsageTrendCached(
|
||||
requestType *int16,
|
||||
stream *bool,
|
||||
billingType *int8,
|
||||
upstreamModelMismatch *bool,
|
||||
) ([]usagestats.TrendDataPoint, bool, error) {
|
||||
key := mustMarshalDashboardCacheKey(dashboardTrendCacheKey{
|
||||
StartTime: startTime.UTC().Format(time.RFC3339),
|
||||
EndTime: endTime.UTC().Format(time.RFC3339),
|
||||
Granularity: granularity,
|
||||
UserID: userID,
|
||||
APIKeyID: apiKeyID,
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
Model: model,
|
||||
RequestType: requestType,
|
||||
Stream: stream,
|
||||
BillingType: billingType,
|
||||
StartTime: startTime.UTC().Format(time.RFC3339),
|
||||
EndTime: endTime.UTC().Format(time.RFC3339),
|
||||
Granularity: granularity,
|
||||
UserID: userID,
|
||||
APIKeyID: apiKeyID,
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
Model: model,
|
||||
RequestType: requestType,
|
||||
Stream: stream,
|
||||
BillingType: billingType,
|
||||
UpstreamModelMismatch: upstreamModelMismatch,
|
||||
})
|
||||
entry, hit, err := dashboardTrendCache.GetOrLoad(key, func() (any, error) {
|
||||
return h.dashboardService.GetUsageTrendWithFilters(ctx, startTime, endTime, granularity, userID, apiKeyID, accountID, groupID, model, requestType, stream, billingType)
|
||||
return h.dashboardService.GetUsageTrendWithUsageFilters(ctx, startTime, endTime, granularity, usagestats.UsageLogFilters{
|
||||
UserID: userID, APIKeyID: apiKeyID, AccountID: accountID, GroupID: groupID,
|
||||
Model: model, RequestType: requestType, Stream: stream, BillingType: billingType,
|
||||
UpstreamModelMismatch: upstreamModelMismatch,
|
||||
})
|
||||
})
|
||||
if err != nil {
|
||||
return nil, hit, err
|
||||
@@ -116,21 +124,27 @@ func (h *DashboardHandler) getModelStatsCached(
|
||||
requestType *int16,
|
||||
stream *bool,
|
||||
billingType *int8,
|
||||
upstreamModelMismatch *bool,
|
||||
) ([]usagestats.ModelStat, bool, error) {
|
||||
key := mustMarshalDashboardCacheKey(dashboardModelGroupCacheKey{
|
||||
StartTime: startTime.UTC().Format(time.RFC3339),
|
||||
EndTime: endTime.UTC().Format(time.RFC3339),
|
||||
UserID: userID,
|
||||
APIKeyID: apiKeyID,
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
ModelSource: usagestats.NormalizeModelSource(modelSource),
|
||||
RequestType: requestType,
|
||||
Stream: stream,
|
||||
BillingType: billingType,
|
||||
StartTime: startTime.UTC().Format(time.RFC3339),
|
||||
EndTime: endTime.UTC().Format(time.RFC3339),
|
||||
UserID: userID,
|
||||
APIKeyID: apiKeyID,
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
ModelSource: usagestats.NormalizeModelSource(modelSource),
|
||||
RequestType: requestType,
|
||||
Stream: stream,
|
||||
BillingType: billingType,
|
||||
UpstreamModelMismatch: upstreamModelMismatch,
|
||||
})
|
||||
entry, hit, err := dashboardModelStatsCache.GetOrLoad(key, func() (any, error) {
|
||||
return h.dashboardService.GetModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, requestType, stream, billingType, modelSource)
|
||||
return h.dashboardService.GetModelStatsWithUsageFiltersBySource(ctx, startTime, endTime, usagestats.UsageLogFilters{
|
||||
UserID: userID, APIKeyID: apiKeyID, AccountID: accountID, GroupID: groupID,
|
||||
RequestType: requestType, Stream: stream, BillingType: billingType,
|
||||
UpstreamModelMismatch: upstreamModelMismatch,
|
||||
}, modelSource)
|
||||
})
|
||||
if err != nil {
|
||||
return nil, hit, err
|
||||
@@ -146,20 +160,26 @@ func (h *DashboardHandler) getGroupStatsCached(
|
||||
requestType *int16,
|
||||
stream *bool,
|
||||
billingType *int8,
|
||||
upstreamModelMismatch *bool,
|
||||
) ([]usagestats.GroupStat, bool, error) {
|
||||
key := mustMarshalDashboardCacheKey(dashboardModelGroupCacheKey{
|
||||
StartTime: startTime.UTC().Format(time.RFC3339),
|
||||
EndTime: endTime.UTC().Format(time.RFC3339),
|
||||
UserID: userID,
|
||||
APIKeyID: apiKeyID,
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
RequestType: requestType,
|
||||
Stream: stream,
|
||||
BillingType: billingType,
|
||||
StartTime: startTime.UTC().Format(time.RFC3339),
|
||||
EndTime: endTime.UTC().Format(time.RFC3339),
|
||||
UserID: userID,
|
||||
APIKeyID: apiKeyID,
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
RequestType: requestType,
|
||||
Stream: stream,
|
||||
BillingType: billingType,
|
||||
UpstreamModelMismatch: upstreamModelMismatch,
|
||||
})
|
||||
entry, hit, err := dashboardGroupStatsCache.GetOrLoad(key, func() (any, error) {
|
||||
return h.dashboardService.GetGroupStatsWithFilters(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, requestType, stream, billingType)
|
||||
return h.dashboardService.GetGroupStatsWithUsageFilters(ctx, startTime, endTime, usagestats.UsageLogFilters{
|
||||
UserID: userID, APIKeyID: apiKeyID, AccountID: accountID, GroupID: groupID,
|
||||
RequestType: requestType, Stream: stream, BillingType: billingType,
|
||||
UpstreamModelMismatch: upstreamModelMismatch,
|
||||
})
|
||||
})
|
||||
if err != nil {
|
||||
return nil, hit, err
|
||||
|
||||
@@ -37,34 +37,36 @@ type dashboardSnapshotV2Response struct {
|
||||
}
|
||||
|
||||
type dashboardSnapshotV2Filters struct {
|
||||
UserID int64
|
||||
APIKeyID int64
|
||||
AccountID int64
|
||||
GroupID int64
|
||||
Model string
|
||||
RequestType *int16
|
||||
Stream *bool
|
||||
BillingType *int8
|
||||
UserID int64
|
||||
APIKeyID int64
|
||||
AccountID int64
|
||||
GroupID int64
|
||||
Model string
|
||||
RequestType *int16
|
||||
Stream *bool
|
||||
BillingType *int8
|
||||
UpstreamModelMismatch *bool
|
||||
}
|
||||
|
||||
type dashboardSnapshotV2CacheKey struct {
|
||||
StartTime string `json:"start_time"`
|
||||
EndTime string `json:"end_time"`
|
||||
Granularity string `json:"granularity"`
|
||||
UserID int64 `json:"user_id"`
|
||||
APIKeyID int64 `json:"api_key_id"`
|
||||
AccountID int64 `json:"account_id"`
|
||||
GroupID int64 `json:"group_id"`
|
||||
Model string `json:"model"`
|
||||
RequestType *int16 `json:"request_type"`
|
||||
Stream *bool `json:"stream"`
|
||||
BillingType *int8 `json:"billing_type"`
|
||||
IncludeStats bool `json:"include_stats"`
|
||||
IncludeTrend bool `json:"include_trend"`
|
||||
IncludeModels bool `json:"include_models"`
|
||||
IncludeGroups bool `json:"include_groups"`
|
||||
IncludeUsersTrend bool `json:"include_users_trend"`
|
||||
UsersTrendLimit int `json:"users_trend_limit"`
|
||||
StartTime string `json:"start_time"`
|
||||
EndTime string `json:"end_time"`
|
||||
Granularity string `json:"granularity"`
|
||||
UserID int64 `json:"user_id"`
|
||||
APIKeyID int64 `json:"api_key_id"`
|
||||
AccountID int64 `json:"account_id"`
|
||||
GroupID int64 `json:"group_id"`
|
||||
Model string `json:"model"`
|
||||
RequestType *int16 `json:"request_type"`
|
||||
Stream *bool `json:"stream"`
|
||||
BillingType *int8 `json:"billing_type"`
|
||||
UpstreamModelMismatch *bool `json:"upstream_model_mismatch"`
|
||||
IncludeStats bool `json:"include_stats"`
|
||||
IncludeTrend bool `json:"include_trend"`
|
||||
IncludeModels bool `json:"include_models"`
|
||||
IncludeGroups bool `json:"include_groups"`
|
||||
IncludeUsersTrend bool `json:"include_users_trend"`
|
||||
UsersTrendLimit int `json:"users_trend_limit"`
|
||||
}
|
||||
|
||||
func (h *DashboardHandler) GetSnapshotV2(c *gin.Context) {
|
||||
@@ -93,23 +95,24 @@ func (h *DashboardHandler) GetSnapshotV2(c *gin.Context) {
|
||||
}
|
||||
|
||||
keyRaw, _ := json.Marshal(dashboardSnapshotV2CacheKey{
|
||||
StartTime: startTime.UTC().Format(time.RFC3339),
|
||||
EndTime: endTime.UTC().Format(time.RFC3339),
|
||||
Granularity: granularity,
|
||||
UserID: filters.UserID,
|
||||
APIKeyID: filters.APIKeyID,
|
||||
AccountID: filters.AccountID,
|
||||
GroupID: filters.GroupID,
|
||||
Model: filters.Model,
|
||||
RequestType: filters.RequestType,
|
||||
Stream: filters.Stream,
|
||||
BillingType: filters.BillingType,
|
||||
IncludeStats: includeStats,
|
||||
IncludeTrend: includeTrend,
|
||||
IncludeModels: includeModels,
|
||||
IncludeGroups: includeGroups,
|
||||
IncludeUsersTrend: includeUsersTrend,
|
||||
UsersTrendLimit: usersTrendLimit,
|
||||
StartTime: startTime.UTC().Format(time.RFC3339),
|
||||
EndTime: endTime.UTC().Format(time.RFC3339),
|
||||
Granularity: granularity,
|
||||
UserID: filters.UserID,
|
||||
APIKeyID: filters.APIKeyID,
|
||||
AccountID: filters.AccountID,
|
||||
GroupID: filters.GroupID,
|
||||
Model: filters.Model,
|
||||
RequestType: filters.RequestType,
|
||||
Stream: filters.Stream,
|
||||
BillingType: filters.BillingType,
|
||||
UpstreamModelMismatch: filters.UpstreamModelMismatch,
|
||||
IncludeStats: includeStats,
|
||||
IncludeTrend: includeTrend,
|
||||
IncludeModels: includeModels,
|
||||
IncludeGroups: includeGroups,
|
||||
IncludeUsersTrend: includeUsersTrend,
|
||||
UsersTrendLimit: usersTrendLimit,
|
||||
})
|
||||
cacheKey := string(keyRaw)
|
||||
|
||||
@@ -184,6 +187,7 @@ func (h *DashboardHandler) buildSnapshotV2Response(
|
||||
filters.RequestType,
|
||||
filters.Stream,
|
||||
filters.BillingType,
|
||||
filters.UpstreamModelMismatch,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to get usage trend")
|
||||
@@ -204,6 +208,7 @@ func (h *DashboardHandler) buildSnapshotV2Response(
|
||||
filters.RequestType,
|
||||
filters.Stream,
|
||||
filters.BillingType,
|
||||
filters.UpstreamModelMismatch,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to get model statistics")
|
||||
@@ -223,6 +228,7 @@ func (h *DashboardHandler) buildSnapshotV2Response(
|
||||
filters.RequestType,
|
||||
filters.Stream,
|
||||
filters.BillingType,
|
||||
filters.UpstreamModelMismatch,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, errors.New("failed to get group statistics")
|
||||
@@ -299,5 +305,13 @@ func parseDashboardSnapshotV2Filters(c *gin.Context) (*dashboardSnapshotV2Filter
|
||||
filters.BillingType = &bt
|
||||
}
|
||||
|
||||
if mismatchStr := strings.TrimSpace(c.Query("upstream_model_mismatch")); mismatchStr != "" {
|
||||
value, err := strconv.ParseBool(mismatchStr)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
filters.UpstreamModelMismatch = &value
|
||||
}
|
||||
|
||||
return filters, nil
|
||||
}
|
||||
@@ -143,6 +143,16 @@ func (h *UsageHandler) List(c *gin.Context) {
|
||||
billingType = &bt
|
||||
}
|
||||
|
||||
var upstreamModelMismatch *bool
|
||||
if raw := strings.TrimSpace(c.Query("upstream_model_mismatch")); raw != "" {
|
||||
value, err := strconv.ParseBool(raw)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid upstream_model_mismatch value, use true or false")
|
||||
return
|
||||
}
|
||||
upstreamModelMismatch = &value
|
||||
}
|
||||
|
||||
// Parse date range
|
||||
var startTime, endTime *time.Time
|
||||
userTZ := c.Query("timezone") // Get user's timezone from request
|
||||
@@ -173,20 +183,21 @@ func (h *UsageHandler) List(c *gin.Context) {
|
||||
SortOrder: c.DefaultQuery("sort_order", "desc"),
|
||||
}
|
||||
filters := usagestats.UsageLogFilters{
|
||||
UserID: userID,
|
||||
APIKeyID: apiKeyID,
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
RequestID: requestID,
|
||||
Model: model,
|
||||
ModelFilterSource: usagestats.ModelSourceRequested,
|
||||
RequestType: requestType,
|
||||
Stream: stream,
|
||||
BillingType: billingType,
|
||||
BillingMode: billingMode,
|
||||
StartTime: startTime,
|
||||
EndTime: endTime,
|
||||
ExactTotal: exactTotal,
|
||||
UserID: userID,
|
||||
APIKeyID: apiKeyID,
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
RequestID: requestID,
|
||||
Model: model,
|
||||
ModelFilterSource: usagestats.ModelSourceRequested,
|
||||
RequestType: requestType,
|
||||
Stream: stream,
|
||||
BillingType: billingType,
|
||||
BillingMode: billingMode,
|
||||
UpstreamModelMismatch: upstreamModelMismatch,
|
||||
StartTime: startTime,
|
||||
EndTime: endTime,
|
||||
ExactTotal: exactTotal,
|
||||
}
|
||||
|
||||
records, result, err := h.usageService.ListWithFilters(c.Request.Context(), params, filters)
|
||||
@@ -276,6 +287,16 @@ func (h *UsageHandler) Stats(c *gin.Context) {
|
||||
billingType = &bt
|
||||
}
|
||||
|
||||
var upstreamModelMismatch *bool
|
||||
if raw := strings.TrimSpace(c.Query("upstream_model_mismatch")); raw != "" {
|
||||
value, err := strconv.ParseBool(raw)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid upstream_model_mismatch value, use true or false")
|
||||
return
|
||||
}
|
||||
upstreamModelMismatch = &value
|
||||
}
|
||||
|
||||
// Parse date range
|
||||
userTZ := c.Query("timezone")
|
||||
now := timezone.NowInUserLocation(userTZ)
|
||||
@@ -315,18 +336,19 @@ func (h *UsageHandler) Stats(c *gin.Context) {
|
||||
|
||||
// Build filters and call GetStatsWithFilters
|
||||
filters := usagestats.UsageLogFilters{
|
||||
UserID: userID,
|
||||
APIKeyID: apiKeyID,
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
Model: model,
|
||||
ModelFilterSource: usagestats.ModelSourceRequested,
|
||||
RequestType: requestType,
|
||||
Stream: stream,
|
||||
BillingType: billingType,
|
||||
BillingMode: billingMode,
|
||||
StartTime: &startTime,
|
||||
EndTime: &endTime,
|
||||
UserID: userID,
|
||||
APIKeyID: apiKeyID,
|
||||
AccountID: accountID,
|
||||
GroupID: groupID,
|
||||
Model: model,
|
||||
ModelFilterSource: usagestats.ModelSourceRequested,
|
||||
RequestType: requestType,
|
||||
Stream: stream,
|
||||
BillingType: billingType,
|
||||
BillingMode: billingMode,
|
||||
UpstreamModelMismatch: upstreamModelMismatch,
|
||||
StartTime: &startTime,
|
||||
EndTime: &endTime,
|
||||
}
|
||||
|
||||
var stats *usagestats.UsageStats
|
||||
|
||||
@@ -11,17 +11,18 @@ import (
|
||||
var usageStatsCache = newSnapshotCache(30 * time.Second)
|
||||
|
||||
type usageStatsCacheKeyData struct {
|
||||
StartTime string `json:"start_time"`
|
||||
EndTime string `json:"end_time"`
|
||||
UserID int64 `json:"user_id"`
|
||||
APIKeyID int64 `json:"api_key_id"`
|
||||
AccountID int64 `json:"account_id"`
|
||||
GroupID int64 `json:"group_id"`
|
||||
Model string `json:"model"`
|
||||
BillingMode string `json:"billing_mode"`
|
||||
RequestType *int16 `json:"request_type"`
|
||||
Stream *bool `json:"stream"`
|
||||
BillingType *int8 `json:"billing_type"`
|
||||
StartTime string `json:"start_time"`
|
||||
EndTime string `json:"end_time"`
|
||||
UserID int64 `json:"user_id"`
|
||||
APIKeyID int64 `json:"api_key_id"`
|
||||
AccountID int64 `json:"account_id"`
|
||||
GroupID int64 `json:"group_id"`
|
||||
Model string `json:"model"`
|
||||
BillingMode string `json:"billing_mode"`
|
||||
RequestType *int16 `json:"request_type"`
|
||||
Stream *bool `json:"stream"`
|
||||
BillingType *int8 `json:"billing_type"`
|
||||
UpstreamModelMismatch *bool `json:"upstream_model_mismatch"`
|
||||
}
|
||||
|
||||
func usageStatsCacheKey(filters usagestats.UsageLogFilters) string {
|
||||
@@ -34,17 +35,18 @@ func usageStatsCacheKey(filters usagestats.UsageLogFilters) string {
|
||||
end = filters.EndTime.UTC().Format(time.RFC3339)
|
||||
}
|
||||
return mustMarshalDashboardCacheKey(usageStatsCacheKeyData{
|
||||
StartTime: start,
|
||||
EndTime: end,
|
||||
UserID: filters.UserID,
|
||||
APIKeyID: filters.APIKeyID,
|
||||
AccountID: filters.AccountID,
|
||||
GroupID: filters.GroupID,
|
||||
Model: filters.Model,
|
||||
BillingMode: filters.BillingMode,
|
||||
RequestType: filters.RequestType,
|
||||
Stream: filters.Stream,
|
||||
BillingType: filters.BillingType,
|
||||
StartTime: start,
|
||||
EndTime: end,
|
||||
UserID: filters.UserID,
|
||||
APIKeyID: filters.APIKeyID,
|
||||
AccountID: filters.AccountID,
|
||||
GroupID: filters.GroupID,
|
||||
Model: filters.Model,
|
||||
BillingMode: filters.BillingMode,
|
||||
RequestType: filters.RequestType,
|
||||
Stream: filters.Stream,
|
||||
BillingType: filters.BillingType,
|
||||
UpstreamModelMismatch: filters.UpstreamModelMismatch,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -1998,6 +1998,21 @@ func (h *AuthHandler) ExchangePendingOAuthCompletion(c *gin.Context) {
|
||||
response.Success(c, payload)
|
||||
return
|
||||
}
|
||||
// ─── 安全修复(账号接管 0day)────────────────────────────────────────────
|
||||
// 非终态 session(如 choose_account_action_required)的 TargetUserID 可能来自
|
||||
// 攻击者提交的他人邮箱:createPendingOAuthAccount / SendPendingOAuthVerifyCode
|
||||
// 发现邮箱已存在时会把本 session 指向该邮箱用户,全程无密码、无邮箱验证码、
|
||||
// 无账号所有权证明。若此时带着 adoption decision 继续执行,下方的
|
||||
// applyPendingOAuthAdoption 会把本 OAuth identity 直接绑定到 TargetUserID,
|
||||
// 攻击者随后再次 OAuth 登录即被系统识别为受害者本人(完整账号接管)。
|
||||
// 只有两类 session 允许在此处执行 adoption/binding:
|
||||
// 1. canIssueTokenPair == true —— 登录终态,identity 已安全绑定该用户;
|
||||
// 2. intent == bind_current_user —— 已登录用户主动发起绑定(绑定目标来自登录态 cookie)。
|
||||
// 其余状态一律只返回 payload,不绑定、不消费 session。
|
||||
if !canIssueTokenPair && !strings.EqualFold(strings.TrimSpace(session.Intent), oauthIntentBindCurrentUser) {
|
||||
response.Success(c, payload)
|
||||
return
|
||||
}
|
||||
if !adoptionDecision.hasDecision() {
|
||||
adoptionRequired, _ := payload["adoption_required"].(bool)
|
||||
if adoptionRequired {
|
||||
|
||||
@@ -910,6 +910,92 @@ func TestExchangePendingOAuthCompletionRejectsDisabledTargetUser(t *testing.T) {
|
||||
require.Nil(t, storedSession.ConsumedAt)
|
||||
}
|
||||
|
||||
func TestExchangePendingOAuthCompletionChoiceStateDoesNotBindIdentity(t *testing.T) {
|
||||
// 回归测试:复刻"补邮箱/创建账户"路径的账号接管 0day。
|
||||
// 攻击者用自己的 OAuth 账号登录后,在 create-account 步骤提交受害者邮箱,
|
||||
// 后端发现邮箱已存在会把 pending session 转入 choice 状态并指向受害者
|
||||
// (TargetUserID=受害者、无密码/验证码证明)。此时带 adoption decision 调
|
||||
// exchange 绝不能把 OAuth identity 绑定到受害者账号。
|
||||
handler, client := newOAuthPendingFlowTestHandler(t, false)
|
||||
ctx := context.Background()
|
||||
|
||||
victim, err := client.User.Create().
|
||||
SetEmail("victim@example.com").
|
||||
SetUsername("victim-user").
|
||||
SetPasswordHash("hash").
|
||||
SetRole(service.RoleUser).
|
||||
SetStatus(service.StatusActive).
|
||||
Save(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
session, err := client.PendingAuthSession.Create().
|
||||
SetSessionToken("choice-state-attack-session-token").
|
||||
SetIntent("login").
|
||||
SetProviderType("linuxdo").
|
||||
SetProviderKey("linuxdo").
|
||||
SetProviderSubject("attacker-subject-123").
|
||||
SetTargetUserID(victim.ID).
|
||||
SetResolvedEmail(victim.Email).
|
||||
SetBrowserSessionKey("choice-state-attack-browser-session-key").
|
||||
SetUpstreamIdentityClaims(map[string]any{
|
||||
"username": "attacker_linuxdo_user",
|
||||
"suggested_display_name": "Attacker Display Name",
|
||||
"suggested_avatar_url": "https://cdn.example/attacker.png",
|
||||
}).
|
||||
SetLocalFlowState(map[string]any{
|
||||
oauthCompletionResponseKey: map[string]any{
|
||||
"step": oauthPendingChoiceStep,
|
||||
"adoption_required": true,
|
||||
"force_email_on_signup": true,
|
||||
"email_binding_required": true,
|
||||
"existing_account_bindable": true,
|
||||
"email": victim.Email,
|
||||
"resolved_email": victim.Email,
|
||||
"redirect": "/dashboard",
|
||||
},
|
||||
}).
|
||||
SetExpiresAt(time.Now().UTC().Add(10 * time.Minute)).
|
||||
Save(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
body := bytes.NewBufferString(`{"adopt_display_name":true,"adopt_avatar":true}`)
|
||||
recorder := httptest.NewRecorder()
|
||||
ginCtx, _ := gin.CreateTestContext(recorder)
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/auth/oauth/pending/exchange", body)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.AddCookie(&http.Cookie{Name: oauthPendingSessionCookieName, Value: encodeCookieValue(session.SessionToken)})
|
||||
req.AddCookie(&http.Cookie{Name: oauthPendingBrowserCookieName, Value: encodeCookieValue("choice-state-attack-browser-session-key")})
|
||||
ginCtx.Request = req
|
||||
|
||||
handler.ExchangePendingOAuthCompletion(ginCtx)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
data := decodeJSONResponseData(t, recorder)
|
||||
require.NotContains(t, data, "access_token")
|
||||
require.Equal(t, oauthPendingChoiceStep, data["step"])
|
||||
|
||||
// 攻击者的 OAuth identity 绝不能绑定到受害者账号
|
||||
identityCount, err := client.AuthIdentity.Query().
|
||||
Where(
|
||||
authidentity.ProviderTypeEQ("linuxdo"),
|
||||
authidentity.ProviderKeyEQ("linuxdo"),
|
||||
authidentity.ProviderSubjectEQ("attacker-subject-123"),
|
||||
).
|
||||
Count(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, identityCount)
|
||||
|
||||
// 受害者资料不得被 adoption 篡改
|
||||
storedVictim, err := client.User.Get(ctx, victim.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "victim-user", storedVictim.Username)
|
||||
|
||||
// session 不得被消费(攻击者无法进入下一环)
|
||||
storedSession, err := client.PendingAuthSession.Get(ctx, session.ID)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, storedSession.ConsumedAt)
|
||||
}
|
||||
|
||||
func TestNormalizePendingOAuthCompletionResponseScrubsLegacyTokenPayload(t *testing.T) {
|
||||
payload := normalizePendingOAuthCompletionResponse(map[string]any{
|
||||
"access_token": "legacy-access-token",
|
||||
|
||||
@@ -710,6 +710,8 @@ func UsageLogFromServiceAdmin(l *service.UsageLog) *AdminUsageLog {
|
||||
return &AdminUsageLog{
|
||||
UsageLog: usageLog,
|
||||
UpstreamModel: l.UpstreamModel,
|
||||
UpstreamResponseModel: l.UpstreamResponseModel,
|
||||
UpstreamModelMismatch: l.UpstreamModelMismatch,
|
||||
ChannelID: l.ChannelID,
|
||||
ModelMappingChain: l.ModelMappingChain,
|
||||
BillingTier: l.BillingTier,
|
||||
|
||||
@@ -110,11 +110,15 @@ func TestUsageLogFromService_UsesRequestedModelAndKeepsUpstreamAdminOnly(t *test
|
||||
t.Parallel()
|
||||
|
||||
upstreamModel := "claude-sonnet-4-20250514"
|
||||
upstreamResponseModel := "claude-sonnet-4-20250513"
|
||||
upstreamModelMismatch := true
|
||||
log := &service.UsageLog{
|
||||
RequestID: "req_4",
|
||||
Model: upstreamModel,
|
||||
RequestedModel: "claude-sonnet-4",
|
||||
UpstreamModel: &upstreamModel,
|
||||
RequestID: "req_4",
|
||||
Model: upstreamModel,
|
||||
RequestedModel: "claude-sonnet-4",
|
||||
UpstreamModel: &upstreamModel,
|
||||
UpstreamResponseModel: &upstreamResponseModel,
|
||||
UpstreamModelMismatch: &upstreamModelMismatch,
|
||||
}
|
||||
|
||||
userDTO := UsageLogFromService(log)
|
||||
@@ -126,10 +130,14 @@ func TestUsageLogFromService_UsesRequestedModelAndKeepsUpstreamAdminOnly(t *test
|
||||
userJSON, err := json.Marshal(userDTO)
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, string(userJSON), "upstream_model")
|
||||
require.NotContains(t, string(userJSON), "upstream_response_model")
|
||||
require.NotContains(t, string(userJSON), "upstream_model_mismatch")
|
||||
|
||||
adminJSON, err := json.Marshal(adminDTO)
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, string(adminJSON), `"upstream_model":"claude-sonnet-4-20250514"`)
|
||||
require.Contains(t, string(adminJSON), `"upstream_response_model":"claude-sonnet-4-20250513"`)
|
||||
require.Contains(t, string(adminJSON), `"upstream_model_mismatch":true`)
|
||||
}
|
||||
|
||||
func TestUsageLogFromService_KeepsUserBillingAndIPWithoutAdminCostFields(t *testing.T) {
|
||||
|
||||
@@ -561,6 +561,10 @@ type AdminUsageLog struct {
|
||||
// UpstreamModel is the actual model sent to the upstream provider after mapping.
|
||||
// Omitted when no mapping was applied (requested model was used as-is).
|
||||
UpstreamModel *string `json:"upstream_model,omitempty"`
|
||||
// UpstreamResponseModel is the raw model declared by the upstream response.
|
||||
UpstreamResponseModel *string `json:"upstream_response_model,omitempty"`
|
||||
// UpstreamModelMismatch is nil when the upstream did not declare a model.
|
||||
UpstreamModelMismatch *bool `json:"upstream_model_mismatch,omitempty"`
|
||||
|
||||
// ChannelID 渠道 ID
|
||||
ChannelID *int64 `json:"channel_id,omitempty"`
|
||||
|
||||
@@ -144,7 +144,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
}
|
||||
|
||||
sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, body)
|
||||
requestCtx := service.WithOpenAIImageGenerationIntent(c.Request.Context())
|
||||
requestCtx := service.WithOpenAIImagesEndpoint(service.WithOpenAIImageGenerationIntent(c.Request.Context()))
|
||||
|
||||
maxAccountSwitches := h.maxAccountSwitches
|
||||
switchCount := 0
|
||||
|
||||
@@ -224,6 +224,8 @@ func TestUserUsageListKeepsUserBillingAndIPWithoutAdminCostFields(t *testing.T)
|
||||
require.NotContains(t, body, "account_rate_multiplier")
|
||||
require.NotContains(t, body, "account_stats_cost")
|
||||
require.NotContains(t, body, "upstream_model")
|
||||
require.NotContains(t, body, "upstream_response_model")
|
||||
require.NotContains(t, body, "upstream_model_mismatch")
|
||||
require.NotContains(t, body, "billing_tier")
|
||||
require.NotContains(t, body, "channel_id")
|
||||
require.NotContains(t, body, `"account":`)
|
||||
|
||||
@@ -178,6 +178,11 @@ var geminiModels = []modelDef{
|
||||
{ID: "gemini-3.1-pro-high", DisplayName: "Gemini 3.1 Pro High", CreatedAt: "2026-02-19T00:00:00Z", IsReasoning: true},
|
||||
{ID: "gemini-3.1-flash-image", DisplayName: "Gemini 3.1 Flash Image", CreatedAt: "2026-02-19T00:00:00Z"},
|
||||
{ID: "gemini-3.1-flash-image-preview", DisplayName: "Gemini 3.1 Flash Image Preview", CreatedAt: "2026-02-19T00:00:00Z"},
|
||||
{ID: "gemini-3.6-flash", DisplayName: "Gemini 3.6 Flash", CreatedAt: "2026-07-21T00:00:00Z"},
|
||||
{ID: "gemini-3.6-flash-high", DisplayName: "Gemini 3.6 Flash High", CreatedAt: "2026-07-21T00:00:00Z", IsReasoning: true},
|
||||
{ID: "gemini-3.6-flash-low", DisplayName: "Gemini 3.6 Flash Low", CreatedAt: "2026-07-21T00:00:00Z", IsReasoning: true},
|
||||
{ID: "gemini-3.6-flash-medium", DisplayName: "Gemini 3.6 Flash Medium", CreatedAt: "2026-07-21T00:00:00Z", IsReasoning: true},
|
||||
{ID: "gemini-3.6-flash-tiered", DisplayName: "Gemini 3.6 Flash", CreatedAt: "2026-07-21T00:00:00Z", IsReasoning: true},
|
||||
{ID: "gemini-3-pro-preview", DisplayName: "Gemini 3 Pro Preview", CreatedAt: "2025-06-01T00:00:00Z", IsReasoning: true},
|
||||
{ID: "gemini-3-pro-image", DisplayName: "Gemini 3 Pro Image", CreatedAt: "2025-06-01T00:00:00Z"},
|
||||
}
|
||||
|
||||
@@ -20,6 +20,11 @@ func TestDefaultModels_ContainsNewAndLegacyImageModels(t *testing.T) {
|
||||
"gemini-3.1-flash-image",
|
||||
"gemini-3.1-flash-image-preview",
|
||||
"gemini-3-pro-image", // legacy compatibility
|
||||
"gemini-3.6-flash",
|
||||
"gemini-3.6-flash-high",
|
||||
"gemini-3.6-flash-low",
|
||||
"gemini-3.6-flash-medium",
|
||||
"gemini-3.6-flash-tiered",
|
||||
}
|
||||
|
||||
for _, id := range requiredIDs {
|
||||
|
||||
@@ -0,0 +1,207 @@
|
||||
package apicompat
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// anthropicInboundBlockTypes 是 Anthropic Messages 请求体里合法的 content block
|
||||
// 类型。转换结果只能落在这个集合内——发出集合外的类型,上游一律回
|
||||
// 400 "Request body format invalid"(见 issue #5329)。
|
||||
var anthropicInboundBlockTypes = map[string]bool{
|
||||
"text": true,
|
||||
"image": true,
|
||||
"document": true,
|
||||
"tool_use": true,
|
||||
"tool_result": true,
|
||||
"thinking": true,
|
||||
"redacted_thinking": true,
|
||||
}
|
||||
|
||||
func responsesToAnthropicMessages(t *testing.T, input string) []AnthropicMessage {
|
||||
t.Helper()
|
||||
var req ResponsesRequest
|
||||
require.NoError(t, json.Unmarshal([]byte(`{"model":"glm-5.2","input":`+input+`}`), &req))
|
||||
out, err := ResponsesToAnthropicRequest(&req)
|
||||
require.NoError(t, err)
|
||||
return out.Messages
|
||||
}
|
||||
|
||||
// requireAnthropicMessagesAreSendable 断言消息序列不含 Anthropic 会拒收的形态:
|
||||
// 未知 block 类型、空内容消息、纯空白 text 块。
|
||||
func requireAnthropicMessagesAreSendable(t *testing.T, messages []AnthropicMessage) {
|
||||
t.Helper()
|
||||
for i, m := range messages {
|
||||
raw := strings.TrimSpace(string(m.Content))
|
||||
require.NotContains(t, []string{"", "null", `""`, "[]"}, raw,
|
||||
"messages[%d] 内容为空,Anthropic 拒收空内容消息", i)
|
||||
|
||||
var s string
|
||||
if err := json.Unmarshal(m.Content, &s); err == nil {
|
||||
require.NotEmpty(t, strings.TrimSpace(s), "messages[%d] 字符串内容不能全为空白", i)
|
||||
continue
|
||||
}
|
||||
blocks := parseContentBlocks(m.Content)
|
||||
require.NotEmpty(t, blocks, "messages[%d] 解析不出任何 block", i)
|
||||
for j, b := range blocks {
|
||||
require.True(t, anthropicInboundBlockTypes[b.Type],
|
||||
"messages[%d].content[%d] 是 Anthropic 不认识的 block 类型 %q", i, j, b.Type)
|
||||
if b.Type == "text" {
|
||||
require.NotEmpty(t, strings.TrimSpace(b.Text),
|
||||
"messages[%d].content[%d] 是空白 text 块,Anthropic 拒收", i, j)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// issue #5329:工具执行后的下一轮,Codex 会把 reasoning item 一起回放。
|
||||
// 该 item 带 content 数组时,reasoning_text 块以前会被原样塞进 Anthropic 请求体。
|
||||
func TestResponsesToAnthropic_ReasoningItemWithContentIsDropped(t *testing.T) {
|
||||
messages := responsesToAnthropicMessages(t, `[
|
||||
{"type":"message","role":"user","content":[{"type":"input_text","text":"run a shell command"}]},
|
||||
{"type":"reasoning","id":"rs_1","summary":[],"content":[{"type":"reasoning_text","text":"let me think"}]}
|
||||
]`)
|
||||
|
||||
requireAnthropicMessagesAreSendable(t, messages)
|
||||
require.Len(t, messages, 1)
|
||||
require.NotContains(t, string(messages[0].Content), "reasoning_text")
|
||||
require.NotContains(t, string(messages[0].Content), "let me think")
|
||||
}
|
||||
|
||||
// Codex 的常见 reasoning 形态(只有 summary + encrypted_content)本来就会被丢弃,
|
||||
// 这条守卫确保行为没有被改变。
|
||||
func TestResponsesToAnthropic_ReasoningItemSummaryOnlyStillDropped(t *testing.T) {
|
||||
messages := responsesToAnthropicMessages(t, `[
|
||||
{"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]},
|
||||
{"type":"reasoning","id":"rs_1","summary":[{"type":"summary_text","text":"s"}],"encrypted_content":"gAAAA"}
|
||||
]`)
|
||||
|
||||
requireAnthropicMessagesAreSendable(t, messages)
|
||||
require.Len(t, messages, 1)
|
||||
require.NotContains(t, string(messages[0].Content), "gAAAA")
|
||||
}
|
||||
|
||||
// 未知 item type 的 content 以前会被逐字透传,把 Responses 专有分片带进上游请求。
|
||||
func TestResponsesToAnthropic_UnknownItemTypeContentIsSanitized(t *testing.T) {
|
||||
messages := responsesToAnthropicMessages(t, `[
|
||||
{"type":"web_search_call","id":"ws_1","content":[{"type":"web_search_result","text":"payload"}]}
|
||||
]`)
|
||||
|
||||
requireAnthropicMessagesAreSendable(t, messages)
|
||||
require.Empty(t, messages, "整条内容都无法映射时不应发出消息")
|
||||
}
|
||||
|
||||
// 未知 item type 里夹带的可识别文本仍然保留,不做无谓丢弃。
|
||||
func TestResponsesToAnthropic_UnknownItemTypeKeepsRecognizableText(t *testing.T) {
|
||||
messages := responsesToAnthropicMessages(t, `[
|
||||
{"type":"some_future_item","content":[
|
||||
{"type":"input_text","text":"keep me"},
|
||||
{"type":"reasoning_text","text":"drop me"}
|
||||
]}
|
||||
]`)
|
||||
|
||||
requireAnthropicMessagesAreSendable(t, messages)
|
||||
require.Len(t, messages, 1)
|
||||
require.Contains(t, string(messages[0].Content), "keep me")
|
||||
require.NotContains(t, string(messages[0].Content), "drop me")
|
||||
}
|
||||
|
||||
// user 消息的分片全部不可识别时,以前会退化成 content:"",Anthropic 拒收空内容消息。
|
||||
func TestResponsesToAnthropic_UserMessageWithOnlyUnknownPartsIsDropped(t *testing.T) {
|
||||
messages := responsesToAnthropicMessages(t, `[
|
||||
{"type":"message","role":"user","content":[{"type":"input_file","file_id":"file_1"}]}
|
||||
]`)
|
||||
|
||||
requireAnthropicMessagesAreSendable(t, messages)
|
||||
require.Empty(t, messages)
|
||||
}
|
||||
|
||||
// assistant 侧同理:以前会退化成单个空 text 块,Anthropic 同样拒收。
|
||||
func TestResponsesToAnthropic_AssistantMessageWithOnlyUnknownPartsIsDropped(t *testing.T) {
|
||||
messages := responsesToAnthropicMessages(t, `[
|
||||
{"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]},
|
||||
{"type":"message","role":"assistant","content":[{"type":"refusal","refusal":"no"}]}
|
||||
]`)
|
||||
|
||||
requireAnthropicMessagesAreSendable(t, messages)
|
||||
require.Len(t, messages, 1)
|
||||
require.Equal(t, "user", messages[0].Role)
|
||||
}
|
||||
|
||||
// 完整的 Codex 工具续接回放:tool_use / tool_result 配对必须保持不变,
|
||||
// 同时整个序列满足可发送不变式。
|
||||
func TestResponsesToAnthropic_CodexToolRoundStaysIntactAndSendable(t *testing.T) {
|
||||
messages := responsesToAnthropicMessages(t, `[
|
||||
{"type":"message","role":"user","content":[{"type":"input_text","text":"run ls"}]},
|
||||
{"type":"reasoning","id":"rs_1","summary":[],"content":[{"type":"reasoning_text","text":"plan"}],"encrypted_content":"gAAAA"},
|
||||
{"type":"function_call","id":"fc_1","call_id":"call_1","name":"shell","arguments":"{\"cmd\":\"ls\"}"},
|
||||
{"type":"function_call_output","call_id":"call_1","output":"file1"},
|
||||
{"type":"message","role":"user","content":[{"type":"input_text","text":"continue"}]}
|
||||
]`)
|
||||
|
||||
requireAnthropicMessagesAreSendable(t, messages)
|
||||
|
||||
var sawToolUse, sawToolResult bool
|
||||
for _, m := range messages {
|
||||
for _, b := range parseContentBlocks(m.Content) {
|
||||
switch b.Type {
|
||||
case "tool_use":
|
||||
sawToolUse = true
|
||||
require.Equal(t, "call_1", b.ID)
|
||||
require.Equal(t, "shell", b.Name)
|
||||
case "tool_result":
|
||||
sawToolResult = true
|
||||
require.Equal(t, "call_1", b.ToolUseID)
|
||||
}
|
||||
}
|
||||
}
|
||||
require.True(t, sawToolUse, "function_call 必须转成 tool_use")
|
||||
require.True(t, sawToolResult, "function_call_output 必须转成 tool_result")
|
||||
require.NotContains(t, string(mustMarshal(t, messages)), "reasoning_text")
|
||||
require.NotContains(t, string(mustMarshal(t, messages)), "gAAAA")
|
||||
}
|
||||
|
||||
func mustMarshal(t *testing.T, v any) []byte {
|
||||
t.Helper()
|
||||
b, err := json.Marshal(v)
|
||||
require.NoError(t, err)
|
||||
return b
|
||||
}
|
||||
|
||||
func TestAnthropicContentIsEmpty(t *testing.T) {
|
||||
cases := []struct {
|
||||
raw string
|
||||
want bool
|
||||
}{
|
||||
{``, true},
|
||||
{`""`, true},
|
||||
{`null`, true},
|
||||
{`[]`, true},
|
||||
{` [] `, true},
|
||||
{`"hi"`, false},
|
||||
{`[{"type":"text","text":"hi"}]`, false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
require.Equal(t, tc.want, anthropicContentIsEmpty(json.RawMessage(tc.raw)), "raw=%q", tc.raw)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicContentIsOnlyBlankText(t *testing.T) {
|
||||
cases := []struct {
|
||||
raw string
|
||||
want bool
|
||||
}{
|
||||
{`[{"type":"text","text":""}]`, true},
|
||||
{`[{"type":"text","text":" "}]`, true},
|
||||
{`[{"type":"text","text":""},{"type":"text","text":" "}]`, true},
|
||||
{`[{"type":"text","text":"hi"}]`, false},
|
||||
{`[{"type":"text","text":""},{"type":"image","source":{}}]`, false},
|
||||
{`[]`, false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
require.Equal(t, tc.want, anthropicContentIsOnlyBlankText(json.RawMessage(tc.raw)), "raw=%q", tc.raw)
|
||||
}
|
||||
}
|
||||
@@ -164,11 +164,24 @@ func convertResponsesInputToAnthropic(instructions string, inputRaw json.RawMess
|
||||
Content: blockJSON,
|
||||
})
|
||||
|
||||
case item.Type == "reasoning":
|
||||
// Anthropic 无法摄入 OpenAI 的 reasoning:encrypted_content 是不透明的,
|
||||
// 而 thinking 块的重放需要 Anthropic 自己签发的 signature,无法伪造。
|
||||
// Codex 常见形态(只带 summary + encrypted_content)本来就会被丢弃,
|
||||
// 这里让带 content 数组的形态保持同样行为——否则 reasoning_text 块会被
|
||||
// 原样塞进 Anthropic 请求体,上游直接回 400。
|
||||
|
||||
case item.Role == "user":
|
||||
content, err := convertResponsesUserToAnthropicContent(item.Content)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
// 内容里只有网关不认识的分片时,sanitize 会得到空串。Anthropic 拒收
|
||||
// 空内容消息("all messages must have non-empty content"),整条丢掉
|
||||
// 比发一条必然 400 的消息更可用。
|
||||
if anthropicContentIsEmpty(content) {
|
||||
continue
|
||||
}
|
||||
messages = append(messages, AnthropicMessage{
|
||||
Role: "user",
|
||||
Content: content,
|
||||
@@ -179,19 +192,35 @@ func convertResponsesInputToAnthropic(instructions string, inputRaw json.RawMess
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
// 同上:分片全不认识时会退化成单个空 text 块,而 Anthropic 拒收
|
||||
// 空文本块("text content blocks must contain non-whitespace text")。
|
||||
if anthropicContentIsEmpty(content) || anthropicContentIsOnlyBlankText(content) {
|
||||
continue
|
||||
}
|
||||
messages = append(messages, AnthropicMessage{
|
||||
Role: "assistant",
|
||||
Content: content,
|
||||
})
|
||||
|
||||
default:
|
||||
// Unknown role/type — attempt as user message
|
||||
if item.Content != nil {
|
||||
messages = append(messages, AnthropicMessage{
|
||||
Role: "user",
|
||||
Content: item.Content,
|
||||
})
|
||||
// 未知 role/type —— 尽量当作 user 消息保留其中的文本/图片。
|
||||
// 必须走与真实 user 消息同一套白名单转换:直接透传 item.Content 会把
|
||||
// Responses 专有的分片类型(reasoning_text、web_search_call 的载荷等)
|
||||
// 原样发给 Anthropic,上游只会回 400 把整轮打挂。
|
||||
if item.Content == nil {
|
||||
continue
|
||||
}
|
||||
content, err := convertResponsesUserToAnthropicContent(item.Content)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if anthropicContentIsEmpty(content) {
|
||||
continue
|
||||
}
|
||||
messages = append(messages, AnthropicMessage{
|
||||
Role: "user",
|
||||
Content: content,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -393,6 +422,32 @@ func extractTextFromContent(raw json.RawMessage) string {
|
||||
|
||||
// convertResponsesUserToAnthropicContent converts a Responses user message
|
||||
// content field into Anthropic content blocks JSON.
|
||||
// anthropicContentIsEmpty 判断转换结果是否为"空内容"。
|
||||
// convertResponsesUserToAnthropicContent 在没有任何可识别分片时返回 JSON 空串,
|
||||
// 而 Anthropic 拒收空内容消息。
|
||||
func anthropicContentIsEmpty(content json.RawMessage) bool {
|
||||
trimmed := strings.TrimSpace(string(content))
|
||||
switch trimmed {
|
||||
case "", "null", `""`, "[]":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// anthropicContentIsOnlyBlankText 判断内容是否只由空白 text 块组成。
|
||||
func anthropicContentIsOnlyBlankText(content json.RawMessage) bool {
|
||||
blocks := parseContentBlocks(content)
|
||||
if len(blocks) == 0 {
|
||||
return false
|
||||
}
|
||||
for _, b := range blocks {
|
||||
if b.Type != "text" || strings.TrimSpace(b.Text) != "" {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func convertResponsesUserToAnthropicContent(raw json.RawMessage) (json.RawMessage, error) {
|
||||
if len(raw) == 0 {
|
||||
return json.Marshal("") // empty string content
|
||||
|
||||
@@ -50,6 +50,12 @@ const (
|
||||
// OpenAIImageGenerationIntent 标识 OpenAI 请求会触发生图能力(用于图片能力维度限流)
|
||||
OpenAIImageGenerationIntent Key = "ctx_openai_image_generation_intent"
|
||||
|
||||
// OpenAIImagesEndpoint 标识请求是从 /v1/images/* 入站的。
|
||||
// 与 OpenAIImageGenerationIntent 的区别:后者只表示"这次请求会生图",
|
||||
// /v1/responses 带图片模型时也会置位;本 key 只在专用生图端点置位,
|
||||
// 用于区分"用错端点"与"端点用对了但账号没能力"。
|
||||
OpenAIImagesEndpoint Key = "ctx_openai_images_endpoint"
|
||||
|
||||
// Group 认证后的分组信息,由 API Key 认证中间件设置
|
||||
Group Key = "ctx_group"
|
||||
|
||||
|
||||
@@ -16,10 +16,29 @@ import (
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"golang.org/x/net/proxy"
|
||||
)
|
||||
|
||||
const (
|
||||
// socks5DialTimeout 限制到 SOCKS5 代理自身的 TCP 建连耗时。
|
||||
socks5DialTimeout = 10 * time.Second
|
||||
// socks5DialKeepAlive 与 Go 默认 keepalive 探测间隔保持一致。
|
||||
socks5DialKeepAlive = 30 * time.Second
|
||||
)
|
||||
|
||||
// socks5ForwardDialer 是 SOCKS5 dialer 的底层拨号器。
|
||||
//
|
||||
// proxy.FromURL 的默认 forward dialer 是 proxy.Direct(零值 net.Dialer,无超时),
|
||||
// 代理地址不可达时会一直卡到内核 TCP 重传耗尽(Linux 约 130 秒)。SOCKS5 分支会
|
||||
// 覆盖 Transport.DialContext,因此调用方在 Transport 上设置的建连超时对这条路径
|
||||
// 无效,必须在这里补上。
|
||||
var socks5ForwardDialer = &net.Dialer{
|
||||
Timeout: socks5DialTimeout,
|
||||
KeepAlive: socks5DialKeepAlive,
|
||||
}
|
||||
|
||||
// ConfigureTransportProxy 根据代理 URL 配置 Transport
|
||||
//
|
||||
// 支持的协议:
|
||||
@@ -45,7 +64,7 @@ func ConfigureTransportProxy(transport *http.Transport, proxyURL *url.URL) error
|
||||
return nil
|
||||
|
||||
case "socks5", "socks5h":
|
||||
dialer, err := proxy.FromURL(proxyURL, proxy.Direct)
|
||||
dialer, err := proxy.FromURL(proxyURL, socks5ForwardDialer)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create socks5 dialer: %w", err)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
package proxyutil
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
var errStub = errors.New("stub dial")
|
||||
|
||||
// 回归:SOCKS5 分支覆盖了调用方在 Transport 上设置的 DialContext,
|
||||
// 底层 forward dialer 必须自带建连超时。proxy.Direct 是零值 net.Dialer,
|
||||
// 代理不可达时会一直卡到内核 TCP 重传耗尽(Linux 约 130 秒)。
|
||||
func TestSOCKS5ForwardDialerHasBoundedTimeout(t *testing.T) {
|
||||
require.Greater(t, socks5ForwardDialer.Timeout, time.Duration(0))
|
||||
require.Equal(t, socks5DialTimeout, socks5ForwardDialer.Timeout)
|
||||
require.Equal(t, socks5DialKeepAlive, socks5ForwardDialer.KeepAlive)
|
||||
}
|
||||
|
||||
func TestConfigureTransportProxySOCKS5SetsDialContext(t *testing.T) {
|
||||
for _, scheme := range []string{"socks5", "socks5h"} {
|
||||
t.Run(scheme, func(t *testing.T) {
|
||||
proxyURL, err := url.Parse(scheme + "://127.0.0.1:1080")
|
||||
require.NoError(t, err)
|
||||
|
||||
transport := &http.Transport{}
|
||||
require.NoError(t, ConfigureTransportProxy(transport, proxyURL))
|
||||
require.NotNil(t, transport.DialContext)
|
||||
require.Nil(t, transport.Proxy, "SOCKS5 不应设置 Transport.Proxy")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// HTTP 代理走 Transport.Proxy,不得覆盖调用方设置的 DialContext。
|
||||
func TestConfigureTransportProxyHTTPPreservesDialContext(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://127.0.0.1:8080")
|
||||
require.NoError(t, err)
|
||||
|
||||
called := false
|
||||
transport := &http.Transport{}
|
||||
transport.DialContext = func(_ context.Context, _, _ string) (net.Conn, error) {
|
||||
called = true
|
||||
return nil, errStub
|
||||
}
|
||||
|
||||
require.NoError(t, ConfigureTransportProxy(transport, proxyURL))
|
||||
require.NotNil(t, transport.Proxy)
|
||||
require.NotNil(t, transport.DialContext)
|
||||
|
||||
_, _ = transport.DialContext(context.Background(), "tcp", "127.0.0.1:1")
|
||||
require.True(t, called, "HTTP 代理分支不应替换调用方的 DialContext")
|
||||
}
|
||||
@@ -274,13 +274,14 @@ type UsageLogFilters struct {
|
||||
RequestID string
|
||||
Model string
|
||||
// ModelFilterSource controls how Model is matched. Empty preserves raw usage_logs.model semantics.
|
||||
ModelFilterSource string
|
||||
RequestType *int16
|
||||
Stream *bool
|
||||
BillingType *int8
|
||||
BillingMode string
|
||||
StartTime *time.Time
|
||||
EndTime *time.Time
|
||||
ModelFilterSource string
|
||||
RequestType *int16
|
||||
Stream *bool
|
||||
BillingType *int8
|
||||
BillingMode string
|
||||
UpstreamModelMismatch *bool
|
||||
StartTime *time.Time
|
||||
EndTime *time.Time
|
||||
// ExactTotal requests exact COUNT(*) for pagination. Default false for fast large-table paging.
|
||||
ExactTotal bool
|
||||
}
|
||||
|
||||
@@ -55,6 +55,18 @@ const (
|
||||
// defaultResponseHeaderTimeout: 默认等待响应头超时时间(5分钟)
|
||||
// LLM 请求可能排队较久,需要较长超时
|
||||
defaultResponseHeaderTimeout = 300 * time.Second
|
||||
// defaultUpstreamDialTimeout: 默认 TCP/DNS 建连超时(10秒)
|
||||
// Transport 不设置 DialContext 时会退化为零值 net.Dialer(无超时),建连阶段
|
||||
// 只能依赖内核默认 TCP 重传(Linux 约 130 秒)。ResponseHeaderTimeout 只约束
|
||||
// 连接建立之后等待响应头的阶段,覆盖不到 DNS 解析与 TCP 握手。
|
||||
// 上游域名被解析到 443 不可达的 IP 时(DNS 污染/路由异常),单个账号就要卡满
|
||||
// 内核超时;而多账号故障转移是串行的,一次请求会阻塞数分钟且不写中间错误。
|
||||
defaultUpstreamDialTimeout = 10 * time.Second
|
||||
// defaultUpstreamDialKeepAlive: TCP keepalive 探测间隔,与 Go 默认值保持一致
|
||||
defaultUpstreamDialKeepAlive = 30 * time.Second
|
||||
// defaultUpstreamTLSHandshakeTimeout: TLS 握手超时(10秒)
|
||||
// 与建连超时同量级,避免 TCP 已连通但对端不推进握手时无限等待
|
||||
defaultUpstreamTLSHandshakeTimeout = 10 * time.Second
|
||||
// defaultMaxUpstreamClients: 默认最大客户端缓存数量
|
||||
// 超出后会淘汰最久未使用的客户端
|
||||
defaultMaxUpstreamClients = 5000
|
||||
@@ -1252,6 +1264,17 @@ func defaultPoolSettings(cfg *config.Config) poolSettings {
|
||||
}
|
||||
}
|
||||
|
||||
// newUpstreamDialer 构建上游 Transport 的 TCP dialer。
|
||||
//
|
||||
// 必须显式提供:http.Transport 的 DialContext 为 nil 时使用零值 net.Dialer,
|
||||
// 建连没有任何超时上限,只能等内核 TCP 重传耗尽(Linux 约 130 秒)。
|
||||
func newUpstreamDialer() *net.Dialer {
|
||||
return &net.Dialer{
|
||||
Timeout: defaultUpstreamDialTimeout,
|
||||
KeepAlive: defaultUpstreamDialKeepAlive,
|
||||
}
|
||||
}
|
||||
|
||||
// buildUpstreamTransport 构建上游请求的 Transport
|
||||
// 使用配置文件中的连接池参数,支持生产环境调优
|
||||
//
|
||||
@@ -1264,6 +1287,8 @@ func defaultPoolSettings(cfg *config.Config) poolSettings {
|
||||
// - error: 代理配置错误
|
||||
//
|
||||
// Transport 参数说明:
|
||||
// - DialContext: DNS 解析 + TCP 建连超时(不设置则无上限,退化为内核默认重传)
|
||||
// - TLSHandshakeTimeout: TLS 握手超时
|
||||
// - MaxIdleConns: 所有主机的最大空闲连接总数
|
||||
// - MaxIdleConnsPerHost: 每主机最大空闲连接数(影响连接复用率)
|
||||
// - MaxConnsPerHost: 每主机最大连接数(达到后新请求等待)
|
||||
@@ -1271,6 +1296,8 @@ func defaultPoolSettings(cfg *config.Config) poolSettings {
|
||||
// - ResponseHeaderTimeout: 等待响应头超时(不影响流式传输)
|
||||
func buildUpstreamTransport(settings poolSettings, proxyURL *url.URL, protocolMode string) (*http.Transport, error) {
|
||||
transport := &http.Transport{
|
||||
DialContext: newUpstreamDialer().DialContext,
|
||||
TLSHandshakeTimeout: defaultUpstreamTLSHandshakeTimeout,
|
||||
MaxIdleConns: settings.maxIdleConns,
|
||||
MaxIdleConnsPerHost: settings.maxIdleConnsPerHost,
|
||||
MaxConnsPerHost: settings.maxConnsPerHost,
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"net/url"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// 回归:上游 Transport 必须显式配置建连超时。
|
||||
//
|
||||
// http.Transport.DialContext 为 nil 时 Go 使用零值 net.Dialer(Timeout=0),
|
||||
// DNS 解析与 TCP 握手没有任何上限,只能等内核重传耗尽(Linux 约 130 秒)。
|
||||
// ResponseHeaderTimeout 只覆盖连接建立之后的阶段,管不到建连。
|
||||
// 上游域名被解析到不可达 IP 时,串行的多账号故障转移会把一次请求拖到数分钟。
|
||||
func TestBuildUpstreamTransportSetsDialTimeout(t *testing.T) {
|
||||
settings := defaultPoolSettings(nil)
|
||||
|
||||
transport, err := buildUpstreamTransport(settings, nil, upstreamProtocolModeDefault)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, transport.DialContext, "DialContext 缺失会退化为无超时的零值 dialer")
|
||||
require.Equal(t, defaultUpstreamTLSHandshakeTimeout, transport.TLSHandshakeTimeout)
|
||||
}
|
||||
|
||||
func TestNewUpstreamDialerHasBoundedTimeout(t *testing.T) {
|
||||
dialer := newUpstreamDialer()
|
||||
|
||||
require.Greater(t, dialer.Timeout, time.Duration(0), "建连超时必须有上限")
|
||||
require.Equal(t, defaultUpstreamDialTimeout, dialer.Timeout)
|
||||
require.Equal(t, defaultUpstreamDialKeepAlive, dialer.KeepAlive)
|
||||
}
|
||||
|
||||
// 建连超时对 HTTP 代理同样生效:Transport.Proxy 走的仍是 DialContext,
|
||||
// 代理地址不可达时必须快速失败而不是挂满内核超时。
|
||||
func TestBuildUpstreamTransportKeepsDialTimeoutWithHTTPProxy(t *testing.T) {
|
||||
proxyURL, err := url.Parse("http://127.0.0.1:1080")
|
||||
require.NoError(t, err)
|
||||
|
||||
transport, err := buildUpstreamTransport(defaultPoolSettings(nil), proxyURL, upstreamProtocolModeDefault)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, transport.Proxy)
|
||||
require.NotNil(t, transport.DialContext)
|
||||
}
|
||||
|
||||
// SOCKS5 分支会覆盖 Transport.DialContext,覆盖后仍必须是有超时的拨号器。
|
||||
func TestBuildUpstreamTransportKeepsDialContextWithSOCKS5Proxy(t *testing.T) {
|
||||
proxyURL, err := url.Parse("socks5h://127.0.0.1:1080")
|
||||
require.NoError(t, err)
|
||||
|
||||
transport, err := buildUpstreamTransport(defaultPoolSettings(nil), proxyURL, upstreamProtocolModeDefault)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, transport.DialContext)
|
||||
}
|
||||
|
||||
// Timeout 字段确实被 net.Dialer 用于建连:拨一个已被 close 的本地监听端口,
|
||||
// 断言 Dialer 走的是自己的超时路径而不是无限等待。
|
||||
// (不依赖外网可达性,CI 中确定性执行。)
|
||||
func TestUpstreamDialerRespectsContextCancellation(t *testing.T) {
|
||||
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
require.NoError(t, err)
|
||||
addr := listener.Addr().String()
|
||||
require.NoError(t, listener.Close())
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
conn, err := newUpstreamDialer().DialContext(ctx, "tcp", addr)
|
||||
if conn != nil {
|
||||
_ = conn.Close()
|
||||
}
|
||||
require.Error(t, err, "已取消的 context 必须立即中止拨号")
|
||||
}
|
||||
@@ -57,6 +57,8 @@ const schedulerOutboxPendingDedupKeyMigration = "153_scheduler_outbox_pending_de
|
||||
const schedulerOutboxPendingDedupKeyIndex = "idx_scheduler_outbox_pending_dedup_key"
|
||||
const latestAPIKeyIPIndexMigration = "174_add_usage_logs_api_key_latest_ip_index_notx.sql"
|
||||
const latestAPIKeyIPIndex = "idx_usage_logs_api_key_latest_ip"
|
||||
const usageLogsUpstreamModelMismatchIndexMigration = "195_add_usage_log_upstream_model_mismatch_index_notx.sql"
|
||||
const usageLogsUpstreamModelMismatchIndex = "idx_usage_logs_upstream_model_mismatch_created_at"
|
||||
|
||||
type migrationChecksumCompatibilityRule struct {
|
||||
fileChecksum string
|
||||
@@ -283,6 +285,8 @@ func prepareNonTransactionalMigration(ctx context.Context, db migrationConnectio
|
||||
return dropInvalidIndexIfPresent(ctx, db, schedulerOutboxPendingDedupKeyIndex)
|
||||
case latestAPIKeyIPIndexMigration:
|
||||
return dropInvalidIndexIfPresent(ctx, db, latestAPIKeyIPIndex)
|
||||
case usageLogsUpstreamModelMismatchIndexMigration:
|
||||
return dropInvalidIndexIfPresent(ctx, db, usageLogsUpstreamModelMismatchIndex)
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -155,6 +155,42 @@ CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_logs_api_key_latest_ip
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestApplyMigrationsFS_NonTransactionalMigration_UsageModelMismatchIndexDropsInvalidIndexBeforeRetry(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = db.Close() }()
|
||||
|
||||
prepareMigrationsBootstrapExpectations(mock)
|
||||
mock.ExpectQuery("SELECT checksum FROM schema_migrations WHERE filename = \\$1").
|
||||
WithArgs(usageLogsUpstreamModelMismatchIndexMigration).
|
||||
WillReturnError(sql.ErrNoRows)
|
||||
mock.ExpectQuery("SELECT EXISTS \\(").
|
||||
WithArgs(usageLogsUpstreamModelMismatchIndex).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
|
||||
mock.ExpectExec("DROP INDEX CONCURRENTLY IF EXISTS idx_usage_logs_upstream_model_mismatch_created_at").
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
mock.ExpectExec("CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_logs_upstream_model_mismatch_created_at").
|
||||
WillReturnResult(sqlmock.NewResult(0, 0))
|
||||
mock.ExpectExec("INSERT INTO schema_migrations \\(filename, checksum\\) VALUES \\(\\$1, \\$2\\)").
|
||||
WithArgs(usageLogsUpstreamModelMismatchIndexMigration, sqlmock.AnyArg()).
|
||||
WillReturnResult(sqlmock.NewResult(1, 1))
|
||||
mock.ExpectExec("SELECT pg_advisory_unlock\\(\\$1\\)").
|
||||
WithArgs(migrationsAdvisoryLockID).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
|
||||
fsys := fstest.MapFS{
|
||||
usageLogsUpstreamModelMismatchIndexMigration: &fstest.MapFile{Data: []byte(`
|
||||
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_usage_logs_upstream_model_mismatch_created_at
|
||||
ON usage_logs (created_at DESC, id DESC)
|
||||
WHERE upstream_model_mismatch IS TRUE;
|
||||
`)},
|
||||
}
|
||||
|
||||
err = applyMigrationsFS(context.Background(), db, fsys)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestApplyMigrationsFS_PaymentOrdersOutTradeNoUniqueMigration_FailsFastOnDuplicatePrecheck(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -76,6 +76,24 @@ func TestMigrationsRunner_IsIdempotent_AndSchemaIsUpToDate(t *testing.T) {
|
||||
requireColumn(t, tx, "usage_logs", "video_count", "integer", 0, false)
|
||||
requireColumn(t, tx, "usage_logs", "video_resolution", "character varying", 10, true)
|
||||
requireColumn(t, tx, "usage_logs", "video_duration_seconds", "integer", 0, true)
|
||||
requireColumn(t, tx, "usage_logs", "upstream_response_model", "character varying", 200, true)
|
||||
requireColumn(t, tx, "usage_logs", "upstream_model_mismatch", "boolean", 0, true)
|
||||
requireIndex(t, tx, "usage_logs", usageLogsUpstreamModelMismatchIndex)
|
||||
|
||||
var mismatchIndexDef string
|
||||
require.NoError(t, tx.QueryRowContext(context.Background(), `
|
||||
SELECT pg_get_indexdef(i.indexrelid)
|
||||
FROM pg_class idx
|
||||
JOIN pg_index i ON i.indexrelid = idx.oid
|
||||
JOIN pg_class tbl ON tbl.oid = i.indrelid
|
||||
JOIN pg_namespace ns ON ns.oid = tbl.relnamespace
|
||||
WHERE ns.nspname = 'public'
|
||||
AND tbl.relname = 'usage_logs'
|
||||
AND idx.relname = $1
|
||||
`, usageLogsUpstreamModelMismatchIndex).Scan(&mismatchIndexDef))
|
||||
require.Contains(t, mismatchIndexDef, "created_at DESC")
|
||||
require.Contains(t, mismatchIndexDef, "id DESC")
|
||||
require.Contains(t, mismatchIndexDef, "WHERE (upstream_model_mismatch IS TRUE)")
|
||||
requireConstraintDefinitionContains(
|
||||
t,
|
||||
tx,
|
||||
|
||||
@@ -31,6 +31,8 @@ var usageLogInsertArgTypes = [...]string{
|
||||
"text", // model
|
||||
"text", // requested_model
|
||||
"text", // upstream_model
|
||||
"text", // upstream_response_model
|
||||
"boolean", // upstream_model_mismatch
|
||||
"bigint", // group_id
|
||||
"bigint", // subscription_id
|
||||
"integer", // input_tokens
|
||||
@@ -227,6 +229,8 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor,
|
||||
model,
|
||||
requested_model,
|
||||
upstream_model,
|
||||
upstream_response_model,
|
||||
upstream_model_mismatch,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -278,12 +282,12 @@ func (r *usageLogRepository) createSingle(ctx context.Context, sqlq sqlExecutor,
|
||||
session_id,
|
||||
created_at
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7,
|
||||
$8, $9,
|
||||
$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, $57
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9,
|
||||
$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, $57, $58, $59
|
||||
)
|
||||
ON CONFLICT (request_id, api_key_id) DO NOTHING
|
||||
RETURNING id, created_at
|
||||
@@ -682,6 +686,8 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage
|
||||
model,
|
||||
requested_model,
|
||||
upstream_model,
|
||||
upstream_response_model,
|
||||
upstream_model_mismatch,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -734,9 +740,9 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage
|
||||
created_at
|
||||
) AS (VALUES `)
|
||||
|
||||
// Each batch row prepends the synthetic input_index before the 57
|
||||
// Each batch row prepends the synthetic input_index before the 59
|
||||
// usage-log column values.
|
||||
args := make([]any, 0, len(keys)*58)
|
||||
args := make([]any, 0, len(keys)*60)
|
||||
argPos := 1
|
||||
for idx, key := range keys {
|
||||
if idx > 0 {
|
||||
@@ -772,6 +778,8 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage
|
||||
model,
|
||||
requested_model,
|
||||
upstream_model,
|
||||
upstream_response_model,
|
||||
upstream_model_mismatch,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -831,6 +839,8 @@ func buildUsageLogBatchInsertQuery(keys []string, preparedByKey map[string]usage
|
||||
model,
|
||||
requested_model,
|
||||
upstream_model,
|
||||
upstream_response_model,
|
||||
upstream_model_mismatch,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -930,6 +940,8 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (
|
||||
model,
|
||||
requested_model,
|
||||
upstream_model,
|
||||
upstream_response_model,
|
||||
upstream_model_mismatch,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -982,7 +994,7 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (
|
||||
created_at
|
||||
) AS (VALUES `)
|
||||
|
||||
args := make([]any, 0, len(preparedList)*57)
|
||||
args := make([]any, 0, len(preparedList)*59)
|
||||
argPos := 1
|
||||
for idx, prepared := range preparedList {
|
||||
if idx > 0 {
|
||||
@@ -1015,6 +1027,8 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (
|
||||
model,
|
||||
requested_model,
|
||||
upstream_model,
|
||||
upstream_response_model,
|
||||
upstream_model_mismatch,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -1074,6 +1088,8 @@ func buildUsageLogBestEffortInsertQuery(preparedList []usageLogInsertPrepared) (
|
||||
model,
|
||||
requested_model,
|
||||
upstream_model,
|
||||
upstream_response_model,
|
||||
upstream_model_mismatch,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -1141,6 +1157,8 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared
|
||||
model,
|
||||
requested_model,
|
||||
upstream_model,
|
||||
upstream_response_model,
|
||||
upstream_model_mismatch,
|
||||
group_id,
|
||||
subscription_id,
|
||||
input_tokens,
|
||||
@@ -1192,12 +1210,12 @@ func execUsageLogInsertNoResult(ctx context.Context, sqlq sqlExecutor, prepared
|
||||
session_id,
|
||||
created_at
|
||||
) VALUES (
|
||||
$1, $2, $3, $4, $5, $6, $7,
|
||||
$8, $9,
|
||||
$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, $57
|
||||
$1, $2, $3, $4, $5, $6, $7, $8, $9,
|
||||
$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, $57, $58, $59
|
||||
)
|
||||
ON CONFLICT (request_id, api_key_id) DO NOTHING
|
||||
`, prepared.args...)
|
||||
@@ -1244,6 +1262,8 @@ func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared {
|
||||
requestedModel = strings.TrimSpace(log.Model)
|
||||
}
|
||||
upstreamModel := nullString(log.UpstreamModel)
|
||||
upstreamResponseModel := nullString(log.UpstreamResponseModel)
|
||||
upstreamModelMismatch := nullBool(log.UpstreamModelMismatch)
|
||||
|
||||
var requestIDArg any
|
||||
if requestID != "" {
|
||||
@@ -1263,6 +1283,8 @@ func prepareUsageLogInsert(log *service.UsageLog) usageLogInsertPrepared {
|
||||
log.Model,
|
||||
nullString(&requestedModel),
|
||||
upstreamModel,
|
||||
upstreamResponseModel,
|
||||
upstreamModelMismatch,
|
||||
groupID,
|
||||
subscriptionID,
|
||||
log.InputTokens,
|
||||
|
||||
@@ -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, session_id, created_at"
|
||||
const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, requested_model, upstream_model, upstream_response_model, upstream_model_mismatch, 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"
|
||||
@@ -127,6 +127,9 @@ func (r *usageLogRepository) ListWithFilters(ctx context.Context, params paginat
|
||||
args = append(args, int16(*filters.BillingType))
|
||||
}
|
||||
conditions, args = appendUsageLogBillingModeWhereCondition(conditions, args, filters.BillingMode)
|
||||
if filters.UpstreamModelMismatch != nil {
|
||||
conditions = append(conditions, upstreamModelMismatchCondition("upstream_model_mismatch", *filters.UpstreamModelMismatch))
|
||||
}
|
||||
if filters.StartTime != nil {
|
||||
conditions = append(conditions, fmt.Sprintf("created_at >= $%d", len(args)+1))
|
||||
args = append(args, *filters.StartTime)
|
||||
@@ -157,6 +160,13 @@ func (r *usageLogRepository) ListWithFilters(ctx context.Context, params paginat
|
||||
return logs, page, nil
|
||||
}
|
||||
|
||||
func upstreamModelMismatchCondition(column string, mismatch bool) string {
|
||||
if mismatch {
|
||||
return column + " IS TRUE"
|
||||
}
|
||||
return column + " IS FALSE"
|
||||
}
|
||||
|
||||
func shouldUseFastUsageLogTotal(filters UsageLogFilters) bool {
|
||||
if filters.ExactTotal {
|
||||
return false
|
||||
@@ -437,6 +447,8 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
|
||||
model string
|
||||
requestedModel sql.NullString
|
||||
upstreamModel sql.NullString
|
||||
upstreamResponseModel sql.NullString
|
||||
upstreamModelMismatch sql.NullBool
|
||||
groupID sql.NullInt64
|
||||
subscriptionID sql.NullInt64
|
||||
inputTokens int
|
||||
@@ -498,6 +510,8 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
|
||||
&model,
|
||||
&requestedModel,
|
||||
&upstreamModel,
|
||||
&upstreamResponseModel,
|
||||
&upstreamModelMismatch,
|
||||
&groupID,
|
||||
&subscriptionID,
|
||||
&inputTokens,
|
||||
@@ -651,6 +665,13 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
|
||||
if upstreamModel.Valid {
|
||||
log.UpstreamModel = &upstreamModel.String
|
||||
}
|
||||
if upstreamResponseModel.Valid {
|
||||
log.UpstreamResponseModel = &upstreamResponseModel.String
|
||||
}
|
||||
if upstreamModelMismatch.Valid {
|
||||
value := upstreamModelMismatch.Bool
|
||||
log.UpstreamModelMismatch = &value
|
||||
}
|
||||
if channelID.Valid {
|
||||
value := channelID.Int64
|
||||
log.ChannelID = &value
|
||||
@@ -703,6 +724,13 @@ func nullString(v *string) sql.NullString {
|
||||
return sql.NullString{String: *v, Valid: true}
|
||||
}
|
||||
|
||||
func nullBool(v *bool) sql.NullBool {
|
||||
if v == nil {
|
||||
return sql.NullBool{}
|
||||
}
|
||||
return sql.NullBool{Bool: *v, Valid: true}
|
||||
}
|
||||
|
||||
func nullStringIntMapJSON(v map[string]int) any {
|
||||
if len(v) == 0 {
|
||||
return nil
|
||||
|
||||
@@ -48,6 +48,8 @@ func TestUsageLogRepositoryCreateSyncRequestTypeAndLegacyFields(t *testing.T) {
|
||||
log.Model,
|
||||
log.RequestedModel,
|
||||
sqlmock.AnyArg(), // upstream_model
|
||||
sqlmock.AnyArg(), // upstream_response_model
|
||||
sqlmock.AnyArg(), // upstream_model_mismatch
|
||||
sqlmock.AnyArg(), // group_id
|
||||
sqlmock.AnyArg(), // subscription_id
|
||||
log.InputTokens,
|
||||
@@ -137,9 +139,11 @@ func TestUsageLogRepositoryCreate_PersistsServiceTier(t *testing.T) {
|
||||
log.RequestID,
|
||||
log.Model,
|
||||
log.RequestedModel,
|
||||
sqlmock.AnyArg(),
|
||||
sqlmock.AnyArg(),
|
||||
sqlmock.AnyArg(),
|
||||
sqlmock.AnyArg(), // upstream_model
|
||||
sqlmock.AnyArg(), // upstream_response_model
|
||||
sqlmock.AnyArg(), // upstream_model_mismatch
|
||||
sqlmock.AnyArg(), // group_id
|
||||
sqlmock.AnyArg(), // subscription_id
|
||||
log.InputTokens,
|
||||
log.OutputTokens,
|
||||
log.CacheCreationTokens,
|
||||
@@ -211,8 +215,8 @@ func TestBuildUsageLogBestEffortInsertQuery_IncludesRequestedModelColumn(t *test
|
||||
query, args := buildUsageLogBestEffortInsertQuery([]usageLogInsertPrepared{prepared})
|
||||
|
||||
require.Contains(t, query, "INSERT INTO usage_logs (")
|
||||
require.Contains(t, query, "\n\t\t\tmodel,\n\t\t\trequested_model,\n\t\t\tupstream_model,")
|
||||
require.Contains(t, query, "\n\t\t\trequest_id,\n\t\t\tmodel,\n\t\t\trequested_model,\n\t\t\tupstream_model,")
|
||||
require.Contains(t, query, "\n\t\t\tmodel,\n\t\t\trequested_model,\n\t\t\tupstream_model,\n\t\t\tupstream_response_model,\n\t\t\tupstream_model_mismatch,")
|
||||
require.Contains(t, query, "\n\t\t\trequest_id,\n\t\t\tmodel,\n\t\t\trequested_model,\n\t\t\tupstream_model,\n\t\t\tupstream_response_model,\n\t\t\tupstream_model_mismatch,")
|
||||
require.Len(t, args, len(prepared.args))
|
||||
require.Equal(t, prepared.args[5], args[5])
|
||||
}
|
||||
@@ -273,11 +277,11 @@ func TestPrepareUsageLogInsert_PersistsImageSizeMetadata(t *testing.T) {
|
||||
CreatedAt: time.Date(2025, 1, 6, 12, 0, 0, 0, time.UTC),
|
||||
})
|
||||
|
||||
require.Equal(t, sql.NullString{String: imageSize, Valid: true}, prepared.args[36])
|
||||
require.Equal(t, sql.NullString{String: inputSize, Valid: true}, prepared.args[37])
|
||||
require.Equal(t, sql.NullString{String: outputSize, Valid: true}, prepared.args[38])
|
||||
require.Equal(t, sql.NullString{String: source, Valid: true}, prepared.args[39])
|
||||
breakdownJSON, ok := prepared.args[40].(string)
|
||||
require.Equal(t, sql.NullString{String: imageSize, Valid: true}, prepared.args[38])
|
||||
require.Equal(t, sql.NullString{String: inputSize, Valid: true}, prepared.args[39])
|
||||
require.Equal(t, sql.NullString{String: outputSize, Valid: true}, prepared.args[40])
|
||||
require.Equal(t, sql.NullString{String: source, Valid: true}, prepared.args[41])
|
||||
breakdownJSON, ok := prepared.args[42].(string)
|
||||
require.True(t, ok)
|
||||
require.JSONEq(t, `{"1K":1,"4K":1}`, breakdownJSON)
|
||||
}
|
||||
@@ -809,6 +813,8 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
|
||||
"gpt-image-2",
|
||||
sql.NullString{Valid: true, String: "gpt-image-2"},
|
||||
sql.NullString{},
|
||||
sql.NullString{},
|
||||
sql.NullBool{},
|
||||
sql.NullInt64{},
|
||||
sql.NullInt64{},
|
||||
0, 0, 0, 0, 0, 0,
|
||||
@@ -872,6 +878,8 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
|
||||
"gpt-5", // model
|
||||
sql.NullString{Valid: true, String: "gpt-5"}, // requested_model
|
||||
sql.NullString{}, // upstream_model
|
||||
sql.NullString{}, // upstream_response_model
|
||||
sql.NullBool{}, // upstream_model_mismatch
|
||||
sql.NullInt64{}, // group_id
|
||||
sql.NullInt64{}, // subscription_id
|
||||
1, // input_tokens
|
||||
@@ -942,6 +950,8 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
|
||||
"gpt-5",
|
||||
sql.NullString{Valid: true, String: "gpt-5"},
|
||||
sql.NullString{},
|
||||
sql.NullString{},
|
||||
sql.NullBool{},
|
||||
sql.NullInt64{},
|
||||
sql.NullInt64{},
|
||||
1, 2, 3, 4, 5, 6,
|
||||
@@ -1000,6 +1010,8 @@ func TestScanUsageLogRequestTypeAndLegacyFallback(t *testing.T) {
|
||||
"gpt-5.4",
|
||||
sql.NullString{Valid: true, String: "gpt-5.4"},
|
||||
sql.NullString{},
|
||||
sql.NullString{},
|
||||
sql.NullBool{},
|
||||
sql.NullInt64{},
|
||||
sql.NullInt64{},
|
||||
1, 2, 3, 4, 5, 6,
|
||||
|
||||
@@ -683,6 +683,9 @@ func (r *usageLogRepository) GetStatsWithFilters(ctx context.Context, filters Us
|
||||
args = append(args, int16(*filters.BillingType))
|
||||
}
|
||||
conditions, args = appendUsageLogBillingModeWhereCondition(conditions, args, filters.BillingMode)
|
||||
if filters.UpstreamModelMismatch != nil {
|
||||
conditions = append(conditions, upstreamModelMismatchCondition("upstream_model_mismatch", *filters.UpstreamModelMismatch))
|
||||
}
|
||||
if filters.StartTime != nil {
|
||||
conditions = append(conditions, fmt.Sprintf("created_at >= $%d", len(args)+1))
|
||||
args = append(args, *filters.StartTime)
|
||||
|
||||
@@ -4,6 +4,7 @@ package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -12,6 +13,66 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestUsageLog_UpstreamModelMismatchFilterAndPartialIndex(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
client := tx.Client()
|
||||
repo := newUsageLogRepositoryWithSQL(client, tx)
|
||||
|
||||
user := mustCreateUser(t, client, &service.User{Email: "model-audit@test.com"})
|
||||
apiKey := mustCreateApiKey(t, client, &service.APIKey{UserID: user.ID, Key: "sk-model-audit", Name: "model-audit"})
|
||||
account := mustCreateAccount(t, client, &service.Account{Name: "model-audit-account"})
|
||||
now := time.Now().UTC()
|
||||
responseModel := "gpt-5.4"
|
||||
for _, mismatch := range []bool{true, false} {
|
||||
mismatchValue := mismatch
|
||||
_, err := repo.Create(ctx, &service.UsageLog{
|
||||
UserID: user.ID, APIKeyID: apiKey.ID, AccountID: account.ID,
|
||||
Model: "gpt-5.5", InputTokens: 1, OutputTokens: 1,
|
||||
UpstreamResponseModel: &responseModel, UpstreamModelMismatch: &mismatchValue,
|
||||
CreatedAt: now,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
start := now.Add(-time.Hour)
|
||||
end := now.Add(time.Hour)
|
||||
trueValue := true
|
||||
stats, err := repo.GetStatsWithFilters(ctx, usagestats.UsageLogFilters{
|
||||
UserID: user.ID, StartTime: &start, EndTime: &end, UpstreamModelMismatch: &trueValue,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), stats.TotalRequests)
|
||||
|
||||
trend, err := repo.GetUsageTrendWithUsageFilters(ctx, start, end, "hour", usagestats.UsageLogFilters{
|
||||
UserID: user.ID, UpstreamModelMismatch: &trueValue,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, trend, 1)
|
||||
require.Equal(t, int64(1), trend[0].Requests)
|
||||
|
||||
_, err = tx.ExecContext(ctx, "SET LOCAL enable_seqscan = off")
|
||||
require.NoError(t, err)
|
||||
rows, err := tx.QueryContext(ctx, `
|
||||
EXPLAIN (COSTS OFF)
|
||||
SELECT id
|
||||
FROM usage_logs
|
||||
WHERE upstream_model_mismatch IS TRUE
|
||||
ORDER BY created_at DESC, id DESC
|
||||
LIMIT 100
|
||||
`)
|
||||
require.NoError(t, err)
|
||||
defer func() { require.NoError(t, rows.Close()) }()
|
||||
var planLines []string
|
||||
for rows.Next() {
|
||||
var line string
|
||||
require.NoError(t, rows.Scan(&line))
|
||||
planLines = append(planLines, line)
|
||||
}
|
||||
require.NoError(t, rows.Err())
|
||||
require.Contains(t, strings.Join(planLines, "\n"), usageLogsUpstreamModelMismatchIndex)
|
||||
}
|
||||
|
||||
func TestUsageLog_GetStatsWithFilters_AggregatesAndEndpoints(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
tx := testEntTx(t)
|
||||
|
||||
@@ -265,20 +265,20 @@ func (r *usageLogRepository) GetUserUsageTrendByUserID(ctx context.Context, user
|
||||
|
||||
// GetUserModelStats 获取指定用户的模型统计
|
||||
func (r *usageLogRepository) GetUserModelStats(ctx context.Context, userID int64, startTime, endTime time.Time) (results []ModelStat, err error) {
|
||||
return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, 0, 0, 0, "", nil, nil, nil, usagestats.ModelSourceRequested, "")
|
||||
return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, 0, 0, 0, "", nil, nil, nil, usagestats.ModelSourceRequested, "", nil)
|
||||
}
|
||||
|
||||
// GetUsageTrendWithFilters returns usage trend data with optional filters
|
||||
func (r *usageLogRepository) GetUsageTrendWithFilters(ctx context.Context, startTime, endTime time.Time, granularity string, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8) (results []TrendDataPoint, err error) {
|
||||
return r.getUsageTrendWithFilters(ctx, startTime, endTime, granularity, userID, apiKeyID, accountID, groupID, model, "", requestType, stream, billingType, "")
|
||||
return r.getUsageTrendWithFilters(ctx, startTime, endTime, granularity, userID, apiKeyID, accountID, groupID, model, "", requestType, stream, billingType, "", nil)
|
||||
}
|
||||
|
||||
func (r *usageLogRepository) GetUsageTrendWithUsageFilters(ctx context.Context, startTime, endTime time.Time, granularity string, filters UsageLogFilters) (results []TrendDataPoint, err error) {
|
||||
return r.getUsageTrendWithFilters(ctx, startTime, endTime, granularity, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.ModelFilterSource, filters.RequestType, filters.Stream, filters.BillingType, filters.BillingMode)
|
||||
return r.getUsageTrendWithFilters(ctx, startTime, endTime, granularity, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.ModelFilterSource, filters.RequestType, filters.Stream, filters.BillingType, filters.BillingMode, filters.UpstreamModelMismatch)
|
||||
}
|
||||
|
||||
func (r *usageLogRepository) getUsageTrendWithFilters(ctx context.Context, startTime, endTime time.Time, granularity string, userID, apiKeyID, accountID, groupID int64, model string, modelSource string, requestType *int16, stream *bool, billingType *int8, billingMode string) (results []TrendDataPoint, err error) {
|
||||
if shouldUsePreaggregatedTrend(granularity, userID, apiKeyID, accountID, groupID, model, requestType, stream, billingType, billingMode) {
|
||||
func (r *usageLogRepository) getUsageTrendWithFilters(ctx context.Context, startTime, endTime time.Time, granularity string, userID, apiKeyID, accountID, groupID int64, model string, modelSource string, requestType *int16, stream *bool, billingType *int8, billingMode string, upstreamModelMismatch *bool) (results []TrendDataPoint, err error) {
|
||||
if shouldUsePreaggregatedTrend(granularity, userID, apiKeyID, accountID, groupID, model, requestType, stream, billingType, billingMode, upstreamModelMismatch) {
|
||||
aggregated, aggregatedErr := r.getUsageTrendFromAggregates(ctx, startTime, endTime, granularity)
|
||||
if aggregatedErr == nil && len(aggregated) > 0 {
|
||||
return aggregated, nil
|
||||
@@ -326,6 +326,9 @@ func (r *usageLogRepository) getUsageTrendWithFilters(ctx context.Context, start
|
||||
args = append(args, int16(*billingType))
|
||||
}
|
||||
query, args = appendUsageLogBillingModeQueryFilter(query, args, billingMode, "")
|
||||
if upstreamModelMismatch != nil {
|
||||
query += " AND " + upstreamModelMismatchCondition("upstream_model_mismatch", *upstreamModelMismatch)
|
||||
}
|
||||
query += " GROUP BY date ORDER BY date ASC"
|
||||
|
||||
rows, err := r.sql.QueryContext(ctx, query, args...)
|
||||
@@ -348,7 +351,7 @@ func (r *usageLogRepository) getUsageTrendWithFilters(ctx context.Context, start
|
||||
return results, nil
|
||||
}
|
||||
|
||||
func shouldUsePreaggregatedTrend(granularity string, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8, billingMode string) bool {
|
||||
func shouldUsePreaggregatedTrend(granularity string, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8, billingMode string, upstreamModelMismatch *bool) bool {
|
||||
if granularity != "day" && granularity != "hour" {
|
||||
return false
|
||||
}
|
||||
@@ -360,7 +363,8 @@ func shouldUsePreaggregatedTrend(granularity string, userID, apiKeyID, accountID
|
||||
requestType == nil &&
|
||||
stream == nil &&
|
||||
billingType == nil &&
|
||||
billingMode == ""
|
||||
billingMode == "" &&
|
||||
upstreamModelMismatch == nil
|
||||
}
|
||||
|
||||
func (r *usageLogRepository) getUsageTrendFromAggregates(ctx context.Context, startTime, endTime time.Time, granularity string) (results []TrendDataPoint, err error) {
|
||||
@@ -425,20 +429,20 @@ func (r *usageLogRepository) getUsageTrendFromAggregates(ctx context.Context, st
|
||||
|
||||
// GetModelStatsWithFilters returns model statistics with optional filters
|
||||
func (r *usageLogRepository) GetModelStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8) (results []ModelStat, err error) {
|
||||
return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, "", requestType, stream, billingType, usagestats.ModelSourceRequested, "")
|
||||
return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, "", requestType, stream, billingType, usagestats.ModelSourceRequested, "", nil)
|
||||
}
|
||||
|
||||
// GetModelStatsWithFiltersBySource returns model statistics with optional filters and model source dimension.
|
||||
// source: requested | upstream | mapping.
|
||||
func (r *usageLogRepository) GetModelStatsWithFiltersBySource(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8, source string) (results []ModelStat, err error) {
|
||||
return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, "", requestType, stream, billingType, source, "")
|
||||
return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, "", requestType, stream, billingType, source, "", nil)
|
||||
}
|
||||
|
||||
func (r *usageLogRepository) GetModelStatsWithUsageFiltersBySource(ctx context.Context, startTime, endTime time.Time, filters UsageLogFilters, source string) (results []ModelStat, err error) {
|
||||
return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType, source, filters.BillingMode)
|
||||
return r.getModelStatsWithFiltersBySource(ctx, startTime, endTime, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType, source, filters.BillingMode, filters.UpstreamModelMismatch)
|
||||
}
|
||||
|
||||
func (r *usageLogRepository) getModelStatsWithFiltersBySource(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8, source string, billingMode string) (results []ModelStat, err error) {
|
||||
func (r *usageLogRepository) getModelStatsWithFiltersBySource(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8, source string, billingMode string, upstreamModelMismatch *bool) (results []ModelStat, err error) {
|
||||
actualCostExpr := "COALESCE(SUM(actual_cost), 0) as actual_cost"
|
||||
// 当仅按 account_id 聚合时,实际费用使用账号倍率(total_cost * account_rate_multiplier)。
|
||||
if accountID > 0 && userID == 0 && apiKeyID == 0 {
|
||||
@@ -490,6 +494,9 @@ func (r *usageLogRepository) getModelStatsWithFiltersBySource(ctx context.Contex
|
||||
args = append(args, int16(*billingType))
|
||||
}
|
||||
query, args = appendUsageLogBillingModeQueryFilter(query, args, billingMode, "")
|
||||
if upstreamModelMismatch != nil {
|
||||
query += " AND " + upstreamModelMismatchCondition("upstream_model_mismatch", *upstreamModelMismatch)
|
||||
}
|
||||
query += fmt.Sprintf(" GROUP BY %s ORDER BY total_tokens DESC", modelExpr)
|
||||
|
||||
rows, err := r.sql.QueryContext(ctx, query, args...)
|
||||
@@ -514,14 +521,14 @@ func (r *usageLogRepository) getModelStatsWithFiltersBySource(ctx context.Contex
|
||||
|
||||
// GetGroupStatsWithFilters returns group usage statistics with optional filters
|
||||
func (r *usageLogRepository) GetGroupStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8) (results []usagestats.GroupStat, err error) {
|
||||
return r.getGroupStatsWithFilters(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, "", requestType, stream, billingType, "")
|
||||
return r.getGroupStatsWithFilters(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, "", requestType, stream, billingType, "", nil)
|
||||
}
|
||||
|
||||
func (r *usageLogRepository) GetGroupStatsWithUsageFilters(ctx context.Context, startTime, endTime time.Time, filters UsageLogFilters) (results []usagestats.GroupStat, err error) {
|
||||
return r.getGroupStatsWithFilters(ctx, startTime, endTime, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType, filters.BillingMode)
|
||||
return r.getGroupStatsWithFilters(ctx, startTime, endTime, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType, filters.BillingMode, filters.UpstreamModelMismatch)
|
||||
}
|
||||
|
||||
func (r *usageLogRepository) getGroupStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8, billingMode string) (results []usagestats.GroupStat, err error) {
|
||||
func (r *usageLogRepository) getGroupStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, model string, requestType *int16, stream *bool, billingType *int8, billingMode string, upstreamModelMismatch *bool) (results []usagestats.GroupStat, err error) {
|
||||
query := `
|
||||
SELECT
|
||||
COALESCE(ul.group_id, 0) as group_id,
|
||||
@@ -564,6 +571,9 @@ func (r *usageLogRepository) getGroupStatsWithFilters(ctx context.Context, start
|
||||
args = append(args, int16(*billingType))
|
||||
}
|
||||
query, args = appendUsageLogBillingModeQueryFilter(query, args, billingMode, "ul")
|
||||
if upstreamModelMismatch != nil {
|
||||
query += " AND " + upstreamModelMismatchCondition("ul.upstream_model_mismatch", *upstreamModelMismatch)
|
||||
}
|
||||
query += " GROUP BY ul.group_id, g.name ORDER BY total_tokens DESC"
|
||||
|
||||
rows, err := r.sql.QueryContext(ctx, query, args...)
|
||||
|
||||
@@ -32,7 +32,7 @@ func newSessionIDUsageLog(sessionID *string) *service.UsageLog {
|
||||
// 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")
|
||||
require.Len(t, usageLogInsertArgTypes, 59, "arg-type table must include session_id")
|
||||
|
||||
sessionID := "sess-persisted-123"
|
||||
prepared := prepareUsageLogInsert(newSessionIDUsageLog(&sessionID))
|
||||
|
||||
@@ -368,7 +368,7 @@ func (r *userSubscriptionRepository) UpdateNotes(ctx context.Context, subscripti
|
||||
return translatePersistenceError(err, service.ErrSubscriptionNotFound, nil)
|
||||
}
|
||||
|
||||
func (r *userSubscriptionRepository) ActivateWindows(ctx context.Context, id int64, start time.Time) error {
|
||||
func (r *userSubscriptionRepository) ActivateWindows(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error {
|
||||
client := clientFromContext(ctx, r.client)
|
||||
n, err := client.UserSubscription.Update().
|
||||
Where(
|
||||
@@ -377,24 +377,24 @@ func (r *userSubscriptionRepository) ActivateWindows(ctx context.Context, id int
|
||||
usersubscription.WeeklyWindowStartIsNil(),
|
||||
usersubscription.MonthlyWindowStartIsNil(),
|
||||
).
|
||||
SetDailyWindowStart(start).
|
||||
SetWeeklyWindowStart(start).
|
||||
SetMonthlyWindowStart(start).
|
||||
SetDailyWindowStart(dailyStart).
|
||||
SetWeeklyWindowStart(periodicStart).
|
||||
SetMonthlyWindowStart(periodicStart).
|
||||
Save(ctx)
|
||||
return r.translateConditionalWindowReset(ctx, client, id, n, err)
|
||||
}
|
||||
|
||||
func (r *userSubscriptionRepository) ResetUsageWindows(ctx context.Context, id int64, resetDaily, resetWeekly, resetMonthly bool, newWindowStart time.Time) error {
|
||||
func (r *userSubscriptionRepository) ResetUsageWindows(ctx context.Context, id int64, resetDaily, resetWeekly, resetMonthly bool, dailyStart, periodicStart time.Time) error {
|
||||
client := clientFromContext(ctx, r.client)
|
||||
update := client.UserSubscription.UpdateOneID(id)
|
||||
if resetDaily {
|
||||
update.SetDailyUsageUsd(0).SetDailyWindowStart(newWindowStart)
|
||||
update.SetDailyUsageUsd(0).SetDailyWindowStart(dailyStart)
|
||||
}
|
||||
if resetWeekly {
|
||||
update.SetWeeklyUsageUsd(0).SetWeeklyWindowStart(newWindowStart)
|
||||
update.SetWeeklyUsageUsd(0).SetWeeklyWindowStart(periodicStart)
|
||||
}
|
||||
if resetMonthly {
|
||||
update.SetMonthlyUsageUsd(0).SetMonthlyWindowStart(newWindowStart)
|
||||
update.SetMonthlyUsageUsd(0).SetMonthlyWindowStart(periodicStart)
|
||||
}
|
||||
_, err := update.Save(ctx)
|
||||
return translatePersistenceError(err, service.ErrSubscriptionNotFound, nil)
|
||||
|
||||
@@ -451,8 +451,9 @@ func (s *UserSubscriptionRepoSuite) TestActivateWindows() {
|
||||
group := s.mustCreateGroup("g-activate")
|
||||
sub := s.mustCreateSubscription(user.ID, group.ID, nil)
|
||||
|
||||
dailyStart := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)
|
||||
activateAt := time.Date(2025, 1, 1, 12, 0, 0, 0, time.UTC)
|
||||
err := s.repo.ActivateWindows(s.ctx, sub.ID, activateAt)
|
||||
err := s.repo.ActivateWindows(s.ctx, sub.ID, dailyStart, activateAt)
|
||||
s.Require().NoError(err, "ActivateWindows")
|
||||
|
||||
got, err := s.repo.GetByID(s.ctx, sub.ID)
|
||||
@@ -460,7 +461,9 @@ func (s *UserSubscriptionRepoSuite) TestActivateWindows() {
|
||||
s.Require().NotNil(got.DailyWindowStart)
|
||||
s.Require().NotNil(got.WeeklyWindowStart)
|
||||
s.Require().NotNil(got.MonthlyWindowStart)
|
||||
s.Require().WithinDuration(activateAt, *got.DailyWindowStart, time.Microsecond)
|
||||
s.Require().WithinDuration(dailyStart, *got.DailyWindowStart, time.Microsecond)
|
||||
s.Require().WithinDuration(activateAt, *got.WeeklyWindowStart, time.Microsecond)
|
||||
s.Require().WithinDuration(activateAt, *got.MonthlyWindowStart, time.Microsecond)
|
||||
}
|
||||
|
||||
func (s *UserSubscriptionRepoSuite) TestActivateWindows_StaleActivationPreservesExistingWindows() {
|
||||
@@ -469,15 +472,16 @@ func (s *UserSubscriptionRepoSuite) TestActivateWindows_StaleActivationPreserves
|
||||
sub := s.mustCreateSubscription(user.ID, group.ID, nil)
|
||||
activatedAt := time.Date(2025, 1, 1, 12, 0, 0, 0, time.UTC)
|
||||
manualResetAt := activatedAt.Add(2 * time.Hour)
|
||||
manualDailyStart := time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC)
|
||||
|
||||
s.Require().NoError(s.repo.ActivateWindows(s.ctx, sub.ID, activatedAt))
|
||||
s.Require().NoError(s.repo.ResetUsageWindows(s.ctx, sub.ID, true, true, true, manualResetAt))
|
||||
s.Require().NoError(s.repo.ActivateWindows(s.ctx, sub.ID, activatedAt, activatedAt))
|
||||
s.Require().NoError(s.repo.ResetUsageWindows(s.ctx, sub.ID, true, true, true, manualDailyStart, manualResetAt))
|
||||
// Simulate a concurrent request carrying the original unactivated snapshot.
|
||||
s.Require().NoError(s.repo.ActivateWindows(s.ctx, sub.ID, activatedAt.Add(time.Hour)))
|
||||
s.Require().NoError(s.repo.ActivateWindows(s.ctx, sub.ID, activatedAt.Add(time.Hour), activatedAt.Add(time.Hour)))
|
||||
|
||||
got, err := s.repo.GetByID(s.ctx, sub.ID)
|
||||
s.Require().NoError(err)
|
||||
s.Require().WithinDuration(manualResetAt, *got.DailyWindowStart, time.Microsecond)
|
||||
s.Require().WithinDuration(manualDailyStart, *got.DailyWindowStart, time.Microsecond)
|
||||
s.Require().WithinDuration(manualResetAt, *got.WeeklyWindowStart, time.Microsecond)
|
||||
s.Require().WithinDuration(manualResetAt, *got.MonthlyWindowStart, time.Microsecond)
|
||||
}
|
||||
@@ -535,7 +539,7 @@ func (s *UserSubscriptionRepoSuite) TestResetUsageWindows_ClearsUsageAfterAutoma
|
||||
newWindowStart := oldWindowStart.Add(24 * time.Hour)
|
||||
s.Require().NoError(s.repo.ResetDailyUsage(s.ctx, sub.ID, &oldWindowStart, newWindowStart))
|
||||
s.Require().NoError(s.repo.IncrementUsage(s.ctx, sub.ID, 3))
|
||||
s.Require().NoError(s.repo.ResetUsageWindows(s.ctx, sub.ID, true, false, false, newWindowStart))
|
||||
s.Require().NoError(s.repo.ResetUsageWindows(s.ctx, sub.ID, true, false, false, newWindowStart, newWindowStart))
|
||||
|
||||
got, err := s.repo.GetByID(s.ctx, sub.ID)
|
||||
s.Require().NoError(err)
|
||||
@@ -770,7 +774,7 @@ func (s *UserSubscriptionRepoSuite) TestActiveExpiredBoundaries_UsageAndReset_Ba
|
||||
s.Require().Equal(active.ID, got.ID, "expected active subscription")
|
||||
|
||||
activateAt := time.Now().Add(-25 * time.Hour)
|
||||
s.Require().NoError(s.repo.ActivateWindows(s.ctx, active.ID, activateAt), "ActivateWindows")
|
||||
s.Require().NoError(s.repo.ActivateWindows(s.ctx, active.ID, activateAt, activateAt), "ActivateWindows")
|
||||
s.Require().NoError(s.repo.IncrementUsage(s.ctx, active.ID, 1.25), "IncrementUsage")
|
||||
|
||||
after, err := s.repo.GetByID(s.ctx, active.ID)
|
||||
|
||||
@@ -2259,10 +2259,10 @@ func (stubUserSubscriptionRepo) UpdateStatus(ctx context.Context, subscriptionID
|
||||
func (stubUserSubscriptionRepo) UpdateNotes(ctx context.Context, subscriptionID int64, notes string) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (stubUserSubscriptionRepo) ActivateWindows(ctx context.Context, id int64, start time.Time) error {
|
||||
func (stubUserSubscriptionRepo) ActivateWindows(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (stubUserSubscriptionRepo) ResetUsageWindows(ctx context.Context, id int64, resetDaily, resetWeekly, resetMonthly bool, newWindowStart time.Time) error {
|
||||
func (stubUserSubscriptionRepo) ResetUsageWindows(ctx context.Context, id int64, resetDaily, resetWeekly, resetMonthly bool, dailyStart, periodicStart time.Time) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (stubUserSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, expectedWindowStart *time.Time, newWindowStart time.Time) error {
|
||||
|
||||
@@ -82,7 +82,7 @@ type fakeGoogleSubscriptionRepo struct {
|
||||
getByID func(ctx context.Context, id int64) (*service.UserSubscription, error)
|
||||
getActive func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error)
|
||||
updateStatus func(ctx context.Context, subscriptionID int64, status string) error
|
||||
activateWindow func(ctx context.Context, id int64, start time.Time) error
|
||||
activateWindow func(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error
|
||||
resetDaily func(ctx context.Context, id int64, start time.Time) error
|
||||
resetWeekly func(ctx context.Context, id int64, start time.Time) error
|
||||
resetMonthly func(ctx context.Context, id int64, start time.Time) error
|
||||
@@ -231,13 +231,13 @@ func (f fakeGoogleSubscriptionRepo) UpdateStatus(ctx context.Context, subscripti
|
||||
func (f fakeGoogleSubscriptionRepo) UpdateNotes(ctx context.Context, subscriptionID int64, notes string) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ActivateWindows(ctx context.Context, id int64, start time.Time) error {
|
||||
func (f fakeGoogleSubscriptionRepo) ActivateWindows(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error {
|
||||
if f.activateWindow != nil {
|
||||
return f.activateWindow(ctx, id, start)
|
||||
return f.activateWindow(ctx, id, dailyStart, periodicStart)
|
||||
}
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time) error {
|
||||
func (f fakeGoogleSubscriptionRepo) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time, time.Time) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
func (f fakeGoogleSubscriptionRepo) ResetDailyUsage(ctx context.Context, id int64, _ *time.Time, start time.Time) error {
|
||||
@@ -886,7 +886,7 @@ func TestApiKeyAuthWithSubscriptionGoogle_SubscriptionLimitExceededReturns429(t
|
||||
return &clone, nil
|
||||
},
|
||||
updateStatus: func(ctx context.Context, subscriptionID int64, status string) error { return nil },
|
||||
activateWindow: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
activateWindow: func(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error { return nil },
|
||||
resetDaily: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
resetWeekly: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
resetMonthly: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
|
||||
@@ -119,7 +119,7 @@ func TestSimpleModeBypassesQuotaCheck(t *testing.T) {
|
||||
return &clone, nil
|
||||
},
|
||||
updateStatus: func(ctx context.Context, subscriptionID int64, status string) error { return nil },
|
||||
activateWindow: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
activateWindow: func(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error { return nil },
|
||||
resetDaily: func(ctx context.Context, id int64, start time.Time) error {
|
||||
sub.DailyWindowStart = &start
|
||||
sub.DailyUsageUSD = 0
|
||||
@@ -252,7 +252,7 @@ func TestSimpleModeBypassesQuotaCheck(t *testing.T) {
|
||||
return &clone, nil
|
||||
},
|
||||
updateStatus: func(ctx context.Context, subscriptionID int64, status string) error { return nil },
|
||||
activateWindow: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
activateWindow: func(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error { return nil },
|
||||
resetDaily: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
resetWeekly: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
resetMonthly: func(ctx context.Context, id int64, start time.Time) error { return nil },
|
||||
@@ -1631,7 +1631,7 @@ type stubUserSubscriptionRepo struct {
|
||||
getByID func(ctx context.Context, id int64) (*service.UserSubscription, error)
|
||||
getActive func(ctx context.Context, userID, groupID int64) (*service.UserSubscription, error)
|
||||
updateStatus func(ctx context.Context, subscriptionID int64, status string) error
|
||||
activateWindow func(ctx context.Context, id int64, start time.Time) error
|
||||
activateWindow func(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error
|
||||
resetDaily func(ctx context.Context, id int64, start time.Time) error
|
||||
resetWeekly func(ctx context.Context, id int64, start time.Time) error
|
||||
resetMonthly func(ctx context.Context, id int64, start time.Time) error
|
||||
@@ -1753,14 +1753,14 @@ func (r *stubUserSubscriptionRepo) UpdateNotes(ctx context.Context, subscription
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) ActivateWindows(ctx context.Context, id int64, start time.Time) error {
|
||||
func (r *stubUserSubscriptionRepo) ActivateWindows(ctx context.Context, id int64, dailyStart, periodicStart time.Time) error {
|
||||
if r.activateWindow != nil {
|
||||
return r.activateWindow(ctx, id, start)
|
||||
return r.activateWindow(ctx, id, dailyStart, periodicStart)
|
||||
}
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (r *stubUserSubscriptionRepo) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time) error {
|
||||
func (r *stubUserSubscriptionRepo) ResetUsageWindows(context.Context, int64, bool, bool, bool, time.Time, time.Time) error {
|
||||
return errors.New("not implemented")
|
||||
}
|
||||
|
||||
|
||||
@@ -625,6 +625,11 @@ func (a *Account) resolveModelMapping(rawMapping map[string]any) map[string]stri
|
||||
"gemini-3-flash",
|
||||
"gemini-3.1-pro-high",
|
||||
"gemini-3.1-pro-low",
|
||||
"gemini-3.6-flash",
|
||||
"gemini-3.6-flash-high",
|
||||
"gemini-3.6-flash-low",
|
||||
"gemini-3.6-flash-medium",
|
||||
"gemini-3.6-flash-tiered",
|
||||
})
|
||||
applyAntigravityGemini31ProAliases(result)
|
||||
}
|
||||
|
||||
@@ -28,6 +28,7 @@ import (
|
||||
// ├─ 成功 → 正常返回
|
||||
// └─ 失败 → 设置模型限流 + 清除粘性绑定 → 切换账号
|
||||
func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte, isStickySession bool) (*ForwardResult, error) {
|
||||
beginUpstreamResponseModelObservation(c)
|
||||
// 上游透传账号直接转发,不走 OAuth token 刷新
|
||||
if account.Type == AccountTypeUpstream {
|
||||
return s.ForwardUpstream(ctx, c, account, body)
|
||||
@@ -448,14 +449,16 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context,
|
||||
}
|
||||
|
||||
return &ForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: *usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: billingModel,
|
||||
Stream: claudeReq.Stream,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ClientDisconnect: clientDisconnect,
|
||||
RequestID: requestID,
|
||||
Usage: *usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: billingModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
Stream: claudeReq.Stream,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ClientDisconnect: clientDisconnect,
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -168,6 +168,7 @@ func (s *AntigravityGatewayService) forwardAntigravityCompat(
|
||||
account *Account,
|
||||
request antigravityCompatRequest,
|
||||
) (*ForwardResult, error) {
|
||||
beginUpstreamResponseModelObservation(c)
|
||||
call, err := s.prepareAntigravityCompatCall(ctx, c, account, request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -324,15 +325,17 @@ func (s *AntigravityGatewayService) consumeAntigravityCompatResponse(
|
||||
}
|
||||
|
||||
return &ForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: *streamResult.usage,
|
||||
Model: call.request.originalModel,
|
||||
UpstreamModel: call.billingModel,
|
||||
Stream: call.request.clientStream,
|
||||
Duration: time.Since(call.request.startTime),
|
||||
FirstTokenMs: streamResult.firstTokenMs,
|
||||
ReasoningEffort: call.request.reasoningEffort,
|
||||
ClientDisconnect: streamResult.clientDisconnect,
|
||||
RequestID: requestID,
|
||||
Usage: *streamResult.usage,
|
||||
Model: call.request.originalModel,
|
||||
UpstreamModel: call.billingModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
Stream: call.request.clientStream,
|
||||
Duration: time.Since(call.request.startTime),
|
||||
FirstTokenMs: streamResult.firstTokenMs,
|
||||
ReasoningEffort: call.request.reasoningEffort,
|
||||
ClientDisconnect: streamResult.clientDisconnect,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -477,7 +480,7 @@ func (s *AntigravityGatewayService) handleChatCompletionsNonStreamingFromAntigra
|
||||
startTime time.Time,
|
||||
originalModel string,
|
||||
) (*antigravityStreamResult, error) {
|
||||
claudeResponse, result, err := s.collectClaudeStreamResponse(resp, startTime, originalModel)
|
||||
claudeResponse, result, err := s.collectClaudeStreamResponse(c, resp, startTime, originalModel)
|
||||
if err != nil {
|
||||
return nil, s.mapAntigravityCompatCollectionError(c, err)
|
||||
}
|
||||
@@ -496,7 +499,7 @@ func (s *AntigravityGatewayService) handleResponsesNonStreamingFromAntigravity(
|
||||
startTime time.Time,
|
||||
originalModel string,
|
||||
) (*antigravityStreamResult, error) {
|
||||
claudeResponse, result, err := s.collectClaudeStreamResponse(resp, startTime, originalModel)
|
||||
claudeResponse, result, err := s.collectClaudeStreamResponse(c, resp, startTime, originalModel)
|
||||
if err != nil {
|
||||
return nil, s.mapAntigravityCompatCollectionError(c, err)
|
||||
}
|
||||
|
||||
@@ -301,6 +301,7 @@ func (s *AntigravityGatewayService) handleAntigravityCompatStream(
|
||||
return s.handleAntigravityCompatReadError(c, session, event.err, maxLineSize, prefix)
|
||||
}
|
||||
resetAntigravityCompatTimer(timeoutTimer, timeout)
|
||||
s.observeAntigravityGeminiSSELine(c, event.line)
|
||||
session.consume(event.line)
|
||||
|
||||
case <-timeoutCh:
|
||||
|
||||
@@ -42,6 +42,7 @@ func WithForwardGeminiSession(groupID int64, sessionHash string) ForwardGeminiOp
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte, isStickySession bool, options ...ForwardGeminiOption) (*ForwardResult, error) {
|
||||
beginUpstreamResponseModelObservation(c)
|
||||
startTime := time.Now()
|
||||
forwardOpts := forwardGeminiOptions{}
|
||||
for _, apply := range options {
|
||||
@@ -203,6 +204,7 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
|
||||
if err == nil && fallbackResp.StatusCode < 400 {
|
||||
_ = resp.Body.Close()
|
||||
resp = fallbackResp
|
||||
billingModel = fallbackModel
|
||||
} else if fallbackResp != nil {
|
||||
_ = fallbackResp.Body.Close()
|
||||
}
|
||||
@@ -431,17 +433,19 @@ handleSuccess:
|
||||
}
|
||||
|
||||
return &ForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: *usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: billingModel,
|
||||
Stream: stream,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ClientDisconnect: clientDisconnect,
|
||||
ImageCount: imageCount,
|
||||
ImageSize: imageSize,
|
||||
ImageInputSize: imageInputSize,
|
||||
RequestID: requestID,
|
||||
Usage: *usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: billingModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
Stream: stream,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ClientDisconnect: clientDisconnect,
|
||||
ImageCount: imageCount,
|
||||
ImageSize: imageSize,
|
||||
ImageInputSize: imageInputSize,
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -22,6 +22,26 @@ type antigravityStreamResult struct {
|
||||
clientDisconnect bool // 客户端是否在流式传输过程中断开
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) observeAntigravityGeminiSSELine(c *gin.Context, line string) {
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if !strings.HasPrefix(trimmed, "data:") {
|
||||
return
|
||||
}
|
||||
payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
|
||||
if payload == "" || payload == "[DONE]" {
|
||||
return
|
||||
}
|
||||
raw := []byte(payload)
|
||||
if inner, err := s.unwrapV1InternalResponse(raw); err == nil && len(inner) > 0 {
|
||||
raw = inner
|
||||
}
|
||||
observer.ObserveGemini(raw)
|
||||
}
|
||||
|
||||
// antigravityClientWriter 封装流式响应的客户端写入,自动检测断开并标记。
|
||||
// 断开后所有写入操作变为 no-op,调用方通过 Disconnected() 判断是否继续 drain 上游。
|
||||
type antigravityClientWriter struct {
|
||||
@@ -95,6 +115,9 @@ func handleStreamReadError(err error, clientDisconnected bool, prefix string) (d
|
||||
}
|
||||
|
||||
func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context, resp *http.Response, startTime time.Time) (*antigravityStreamResult, error) {
|
||||
if upstreamResponseModelObserverFromContext(c) == nil {
|
||||
beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
c.Status(resp.StatusCode)
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
@@ -220,6 +243,7 @@ func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context
|
||||
lastDataAt = time.Now()
|
||||
|
||||
line := ev.line
|
||||
s.observeAntigravityGeminiSSELine(c, line)
|
||||
trimmed := strings.TrimRight(line, "\r\n")
|
||||
if strings.HasPrefix(trimmed, "data:") {
|
||||
payload := strings.TrimSpace(strings.TrimPrefix(trimmed, "data:"))
|
||||
@@ -298,6 +322,9 @@ func (s *AntigravityGatewayService) handleGeminiStreamingResponse(c *gin.Context
|
||||
// handleGeminiStreamToNonStreaming 读取上游流式响应,合并为非流式响应返回给客户端
|
||||
// Gemini 流式响应是增量的,需要累积所有 chunk 的内容
|
||||
func (s *AntigravityGatewayService) handleGeminiStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time) (*antigravityStreamResult, error) {
|
||||
if upstreamResponseModelObserverFromContext(c) == nil {
|
||||
beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
maxLineSize := defaultMaxLineSize
|
||||
if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 {
|
||||
@@ -377,6 +404,7 @@ func (s *AntigravityGatewayService) handleGeminiStreamToNonStreaming(c *gin.Cont
|
||||
}
|
||||
|
||||
line := ev.line
|
||||
s.observeAntigravityGeminiSSELine(c, line)
|
||||
trimmed := strings.TrimRight(line, "\r\n")
|
||||
|
||||
if !strings.HasPrefix(trimmed, "data:") {
|
||||
@@ -765,7 +793,10 @@ func (s *AntigravityGatewayService) writeGoogleError(c *gin.Context, status int,
|
||||
|
||||
// collectClaudeStreamResponse 收集上游流式响应,转换为 Claude 非流式格式返回
|
||||
// 用于处理客户端非流式请求但上游只支持流式的情况
|
||||
func (s *AntigravityGatewayService) collectClaudeStreamResponse(resp *http.Response, startTime time.Time, originalModel string) ([]byte, *antigravityStreamResult, error) {
|
||||
func (s *AntigravityGatewayService) collectClaudeStreamResponse(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) ([]byte, *antigravityStreamResult, error) {
|
||||
if upstreamResponseModelObserverFromContext(c) == nil {
|
||||
beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
scanner := bufio.NewScanner(resp.Body)
|
||||
maxLineSize := defaultMaxLineSize
|
||||
if s.settingService.cfg != nil && s.settingService.cfg.Gateway.MaxLineSize > 0 {
|
||||
@@ -860,6 +891,7 @@ func (s *AntigravityGatewayService) collectClaudeStreamResponse(resp *http.Respo
|
||||
if parseErr != nil {
|
||||
continue
|
||||
}
|
||||
upstreamResponseModelObserverFromContext(c).ObserveGemini(inner)
|
||||
|
||||
var parsed map[string]any
|
||||
if err := json.Unmarshal(inner, &parsed); err != nil {
|
||||
@@ -941,7 +973,7 @@ returnResponse:
|
||||
// handleClaudeStreamToNonStreaming 收集上游流式响应,转换为 Claude 非流式格式返回
|
||||
// 用于处理客户端非流式请求但上游只支持流式的情况
|
||||
func (s *AntigravityGatewayService) handleClaudeStreamToNonStreaming(c *gin.Context, resp *http.Response, startTime time.Time, originalModel string) (*antigravityStreamResult, error) {
|
||||
claudeResp, streamRes, err := s.collectClaudeStreamResponse(resp, startTime, originalModel)
|
||||
claudeResp, streamRes, err := s.collectClaudeStreamResponse(c, resp, startTime, originalModel)
|
||||
if err != nil {
|
||||
var failoverErr *UpstreamFailoverError
|
||||
if errors.As(err, &failoverErr) {
|
||||
@@ -1120,6 +1152,7 @@ func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context
|
||||
}
|
||||
|
||||
lastDataAt = time.Now()
|
||||
s.observeAntigravityGeminiSSELine(c, ev.line)
|
||||
|
||||
// 处理 SSE 行,转换为 Claude 格式
|
||||
claudeEvents := processor.ProcessLine(strings.TrimRight(ev.line, "\r\n"))
|
||||
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
|
||||
// ForwardUpstream 使用 base_url + /v1/messages + 双 header 认证透传上游 Claude 请求
|
||||
func (s *AntigravityGatewayService) ForwardUpstream(ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
||||
beginUpstreamResponseModelObservation(c)
|
||||
startTime := time.Now()
|
||||
sessionID := getSessionID(c)
|
||||
prefix := logPrefix(sessionID, account.Name)
|
||||
@@ -129,6 +130,7 @@ func (s *AntigravityGatewayService) ForwardUpstream(ctx context.Context, c *gin.
|
||||
}
|
||||
|
||||
// 提取 usage
|
||||
upstreamResponseModelObserverFromContext(c).ObserveAnthropic(respBody)
|
||||
usage = s.extractClaudeUsage(respBody)
|
||||
|
||||
c.Header("Content-Type", resp.Header.Get("Content-Type"))
|
||||
@@ -141,11 +143,13 @@ func (s *AntigravityGatewayService) ForwardUpstream(ctx context.Context, c *gin.
|
||||
logger.LegacyPrintf("service.antigravity_gateway", "%s status=success duration_ms=%d", prefix, duration.Milliseconds())
|
||||
|
||||
return &ForwardResult{
|
||||
Model: originalModel,
|
||||
Stream: claudeReq.Stream,
|
||||
Duration: duration,
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ClientDisconnect: clientDisconnect,
|
||||
Model: originalModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
Stream: claudeReq.Stream,
|
||||
Duration: duration,
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ClientDisconnect: clientDisconnect,
|
||||
Usage: ClaudeUsage{
|
||||
InputTokens: usage.InputTokens,
|
||||
OutputTokens: usage.OutputTokens,
|
||||
@@ -247,6 +251,9 @@ func (s *AntigravityGatewayService) streamUpstreamResponse(c *gin.Context, resp
|
||||
lastDataAt = time.Now()
|
||||
|
||||
line := ev.line
|
||||
if data, ok := extractAnthropicSSEDataLine(line); ok {
|
||||
upstreamResponseModelObserverFromContext(c).ObserveAnthropic([]byte(strings.TrimSpace(data)))
|
||||
}
|
||||
|
||||
// 记录首 token 时间
|
||||
if firstTokenMs == nil && len(line) > 0 {
|
||||
|
||||
@@ -28,7 +28,8 @@ type PlazaModel struct {
|
||||
// PlazaGroup 模型广场中以分组为顶层的条目。
|
||||
//
|
||||
// 与 AvailableGroupRef 相比多了 Description 与 Models;Models 来自该分组关联渠道的
|
||||
// 支持模型(按分组平台隔离,防跨平台泄漏),与「可用渠道」页口径一致。
|
||||
// 支持模型(普通分组按分组平台隔离,Composite 分组展开关联渠道已配置的
|
||||
// 具体平台),与「可用渠道」页口径一致。
|
||||
type PlazaGroup struct {
|
||||
ID int64
|
||||
Name string
|
||||
@@ -99,8 +100,12 @@ func (s *ChannelService) ListPlazaGroups(ctx context.Context) ([]PlazaGroup, err
|
||||
order = append(order, g.ID)
|
||||
}
|
||||
|
||||
// modelIdx[groupID][modelName] = index into byGroup[groupID].Models
|
||||
modelIdx := make(map[int64]map[string]int, len(groups))
|
||||
type modelKey struct {
|
||||
platform string
|
||||
name string
|
||||
}
|
||||
// modelIdx[groupID][platform+modelName] = index into byGroup[groupID].Models
|
||||
modelIdx := make(map[int64]map[modelKey]int, len(groups))
|
||||
for i := range channels {
|
||||
ch := &channels[i]
|
||||
if ch.Status != StatusActive {
|
||||
@@ -117,23 +122,28 @@ func (s *ChannelService) ListPlazaGroups(ctx context.Context) ([]PlazaGroup, err
|
||||
}
|
||||
idx := modelIdx[gid]
|
||||
if idx == nil {
|
||||
idx = make(map[string]int, len(supported))
|
||||
idx = make(map[modelKey]int, len(supported))
|
||||
modelIdx[gid] = idx
|
||||
}
|
||||
for j := range supported {
|
||||
m := supported[j]
|
||||
if m.Platform != pg.Platform {
|
||||
if pg.Platform == PlatformComposite {
|
||||
if !isConcreteRequestPlatform(m.Platform) {
|
||||
continue
|
||||
}
|
||||
} else if m.Platform != pg.Platform {
|
||||
continue
|
||||
}
|
||||
pricing := plazaImageDisplayPricing(m.Pricing, groupEnt[gid])
|
||||
if at, seen := idx[m.Name]; seen {
|
||||
key := modelKey{platform: m.Platform, name: m.Name}
|
||||
if at, seen := idx[key]; seen {
|
||||
// 先见者胜;仅当已存条目无定价而新条目有定价时升级。
|
||||
if pg.Models[at].Pricing == nil && pricing != nil {
|
||||
pg.Models[at].Pricing = pricing
|
||||
}
|
||||
continue
|
||||
}
|
||||
idx[m.Name] = len(pg.Models)
|
||||
idx[key] = len(pg.Models)
|
||||
pg.Models = append(pg.Models, PlazaModel{
|
||||
Name: m.Name,
|
||||
Platform: m.Platform,
|
||||
@@ -150,7 +160,12 @@ func (s *ChannelService) ListPlazaGroups(ctx context.Context) ([]PlazaGroup, err
|
||||
if len(pg.Models) == 0 {
|
||||
continue
|
||||
}
|
||||
sort.SliceStable(pg.Models, func(i, j int) bool { return pg.Models[i].Name < pg.Models[j].Name })
|
||||
sort.SliceStable(pg.Models, func(i, j int) bool {
|
||||
if pg.Models[i].Name != pg.Models[j].Name {
|
||||
return pg.Models[i].Name < pg.Models[j].Name
|
||||
}
|
||||
return pg.Models[i].Platform < pg.Models[j].Platform
|
||||
})
|
||||
for j := range pg.Models {
|
||||
pg.Models[j].OfficialPricing = s.lookupOfficialPricing(pg.Models[j].Name, officialMemo)
|
||||
}
|
||||
|
||||
@@ -107,6 +107,68 @@ func TestListPlazaGroups_PlatformIsolation(t *testing.T) {
|
||||
require.Equal(t, "gpt-5", byName["g-gpt"][0].Name)
|
||||
}
|
||||
|
||||
func TestListPlazaGroups_CompositeIncludesConfiguredConcretePlatforms(t *testing.T) {
|
||||
anthropicPrice := 3e-6
|
||||
openAIPrice := 2e-6
|
||||
ch := Channel{
|
||||
ID: 1, Name: "multi", Status: StatusActive, GroupIDs: []int64{10},
|
||||
ModelPricing: []ChannelModelPricing{
|
||||
{Platform: PlatformAnthropic, Models: []string{"shared-model"}, InputPrice: &anthropicPrice},
|
||||
{Platform: PlatformOpenAI, Models: []string{"shared-model"}, InputPrice: &openAIPrice},
|
||||
{Platform: "", Models: []string{"empty-platform"}},
|
||||
{Platform: PlatformComposite, Models: []string{"nested-composite"}},
|
||||
{Platform: "unknown-platform", Models: []string{"unknown-platform"}},
|
||||
},
|
||||
}
|
||||
groups := []Group{{ID: 10, Name: "composite", Platform: PlatformComposite, RateMultiplier: 1}}
|
||||
|
||||
out, err := newPlazaChannelService([]Channel{ch}, groups, nil).ListPlazaGroups(context.Background())
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Len(t, out, 1)
|
||||
require.Len(t, out[0].Models, 2, "only concrete platforms are included and same-named models remain distinct")
|
||||
require.Equal(t, PlatformAnthropic, out[0].Models[0].Platform)
|
||||
require.Equal(t, PlatformOpenAI, out[0].Models[1].Platform)
|
||||
require.InDelta(t, anthropicPrice, *out[0].Models[0].Pricing.InputPrice, 1e-12)
|
||||
require.InDelta(t, openAIPrice, *out[0].Models[1].Pricing.InputPrice, 1e-12)
|
||||
}
|
||||
|
||||
func TestListPlazaGroups_CompositeAndOrdinaryGroupsDoNotLeakPlatforms(t *testing.T) {
|
||||
ch := Channel{
|
||||
ID: 1, Name: "multi", Status: StatusActive, GroupIDs: []int64{10, 20},
|
||||
ModelPricing: []ChannelModelPricing{
|
||||
{Platform: PlatformAnthropic, Models: []string{"claude-sonnet"}, InputPrice: testPtrFloat64(3e-6)},
|
||||
{Platform: PlatformOpenAI, Models: []string{"gpt-5"}, InputPrice: testPtrFloat64(2e-6)},
|
||||
},
|
||||
}
|
||||
groups := []Group{
|
||||
{ID: 10, Name: "anthropic-only", Platform: PlatformAnthropic, RateMultiplier: 1},
|
||||
{ID: 20, Name: "composite", Platform: PlatformComposite, RateMultiplier: 1},
|
||||
}
|
||||
|
||||
out, err := newPlazaChannelService([]Channel{ch}, groups, nil).ListPlazaGroups(context.Background())
|
||||
|
||||
require.NoError(t, err)
|
||||
require.Len(t, out, 2)
|
||||
byName := map[string]PlazaGroup{}
|
||||
for _, group := range out {
|
||||
byName[group.Name] = group
|
||||
}
|
||||
require.Len(t, byName["anthropic-only"].Models, 1)
|
||||
require.Equal(t, []PlazaModel{{
|
||||
Name: "claude-sonnet", Platform: PlatformAnthropic, Pricing: byName["anthropic-only"].Models[0].Pricing,
|
||||
}}, byName["anthropic-only"].Models)
|
||||
require.Len(t, byName["composite"].Models, 2)
|
||||
require.Equal(t, []string{"claude-sonnet", "gpt-5"}, []string{
|
||||
byName["composite"].Models[0].Name,
|
||||
byName["composite"].Models[1].Name,
|
||||
})
|
||||
require.Equal(t, []string{PlatformAnthropic, PlatformOpenAI}, []string{
|
||||
byName["composite"].Models[0].Platform,
|
||||
byName["composite"].Models[1].Platform,
|
||||
})
|
||||
}
|
||||
|
||||
func TestListPlazaGroups_InactiveChannelSkipped(t *testing.T) {
|
||||
inactive := plazaPricedChannel(1, "off", []int64{10}, "anthropic", "claude-sonnet")
|
||||
inactive.Status = "inactive"
|
||||
|
||||
@@ -132,6 +132,20 @@ func (s *DashboardService) GetUsageTrendWithFilters(ctx context.Context, startTi
|
||||
return trend, nil
|
||||
}
|
||||
|
||||
func (s *DashboardService) GetUsageTrendWithUsageFilters(ctx context.Context, startTime, endTime time.Time, granularity string, filters usagestats.UsageLogFilters) ([]usagestats.TrendDataPoint, error) {
|
||||
type usageTrendWithFiltersRepo interface {
|
||||
GetUsageTrendWithUsageFilters(context.Context, time.Time, time.Time, string, usagestats.UsageLogFilters) ([]usagestats.TrendDataPoint, error)
|
||||
}
|
||||
if repo, ok := s.usageRepo.(usageTrendWithFiltersRepo); ok {
|
||||
trend, err := repo.GetUsageTrendWithUsageFilters(ctx, startTime, endTime, granularity, filters)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get usage trend with usage filters: %w", err)
|
||||
}
|
||||
return trend, nil
|
||||
}
|
||||
return s.GetUsageTrendWithFilters(ctx, startTime, endTime, granularity, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.Model, filters.RequestType, filters.Stream, filters.BillingType)
|
||||
}
|
||||
|
||||
func (s *DashboardService) GetModelStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8) ([]usagestats.ModelStat, error) {
|
||||
stats, err := s.usageRepo.GetModelStatsWithFilters(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, requestType, stream, billingType)
|
||||
if err != nil {
|
||||
@@ -161,6 +175,21 @@ func (s *DashboardService) GetModelStatsWithFiltersBySource(ctx context.Context,
|
||||
return s.GetModelStatsWithFilters(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, requestType, stream, billingType)
|
||||
}
|
||||
|
||||
func (s *DashboardService) GetModelStatsWithUsageFiltersBySource(ctx context.Context, startTime, endTime time.Time, filters usagestats.UsageLogFilters, modelSource string) ([]usagestats.ModelStat, error) {
|
||||
normalizedSource := usagestats.NormalizeModelSource(modelSource)
|
||||
type modelStatsWithFiltersRepo interface {
|
||||
GetModelStatsWithUsageFiltersBySource(context.Context, time.Time, time.Time, usagestats.UsageLogFilters, string) ([]usagestats.ModelStat, error)
|
||||
}
|
||||
if repo, ok := s.usageRepo.(modelStatsWithFiltersRepo); ok {
|
||||
stats, err := repo.GetModelStatsWithUsageFiltersBySource(ctx, startTime, endTime, filters, normalizedSource)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get model stats with usage filters by source: %w", err)
|
||||
}
|
||||
return stats, nil
|
||||
}
|
||||
return s.GetModelStatsWithFiltersBySource(ctx, startTime, endTime, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.RequestType, filters.Stream, filters.BillingType, normalizedSource)
|
||||
}
|
||||
|
||||
func (s *DashboardService) GetGroupStatsWithFilters(ctx context.Context, startTime, endTime time.Time, userID, apiKeyID, accountID, groupID int64, requestType *int16, stream *bool, billingType *int8) ([]usagestats.GroupStat, error) {
|
||||
stats, err := s.usageRepo.GetGroupStatsWithFilters(ctx, startTime, endTime, userID, apiKeyID, accountID, groupID, requestType, stream, billingType)
|
||||
if err != nil {
|
||||
@@ -169,6 +198,20 @@ func (s *DashboardService) GetGroupStatsWithFilters(ctx context.Context, startTi
|
||||
return stats, nil
|
||||
}
|
||||
|
||||
func (s *DashboardService) GetGroupStatsWithUsageFilters(ctx context.Context, startTime, endTime time.Time, filters usagestats.UsageLogFilters) ([]usagestats.GroupStat, error) {
|
||||
type groupStatsWithFiltersRepo interface {
|
||||
GetGroupStatsWithUsageFilters(context.Context, time.Time, time.Time, usagestats.UsageLogFilters) ([]usagestats.GroupStat, error)
|
||||
}
|
||||
if repo, ok := s.usageRepo.(groupStatsWithFiltersRepo); ok {
|
||||
stats, err := repo.GetGroupStatsWithUsageFilters(ctx, startTime, endTime, filters)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("get group stats with usage filters: %w", err)
|
||||
}
|
||||
return stats, nil
|
||||
}
|
||||
return s.GetGroupStatsWithFilters(ctx, startTime, endTime, filters.UserID, filters.APIKeyID, filters.AccountID, filters.GroupID, filters.RequestType, filters.Stream, filters.BillingType)
|
||||
}
|
||||
|
||||
// GetGroupUsageSummary returns today's and cumulative cost for all groups.
|
||||
func (s *DashboardService) GetGroupUsageSummary(ctx context.Context, todayStart time.Time) ([]usagestats.GroupUsageSummary, error) {
|
||||
results, err := s.usageRepo.GetAllGroupUsageSummary(ctx, todayStart)
|
||||
|
||||
@@ -271,7 +271,7 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
|
||||
if err != nil {
|
||||
// 流中断时保留已观测到的 usage 与错误一起返回,避免上游已计量的请求
|
||||
// 完全漏记漏计费(issue #5148)。
|
||||
if partial := partialStreamUsageResult(resp, streamResult, input.OriginalModel, input.RequestModel, input.StartTime, err); partial != nil {
|
||||
if partial := partialStreamUsageResult(c, resp, streamResult, input.OriginalModel, input.RequestModel, input.StartTime, err); partial != nil {
|
||||
return partial, err
|
||||
}
|
||||
return nil, err
|
||||
@@ -290,14 +290,16 @@ func (s *GatewayService) forwardAnthropicAPIKeyPassthroughWithInput(
|
||||
}
|
||||
|
||||
return &ForwardResult{
|
||||
RequestID: resp.Header.Get("x-request-id"),
|
||||
Usage: *usage,
|
||||
Model: input.OriginalModel,
|
||||
UpstreamModel: input.RequestModel,
|
||||
Stream: input.RequestStream,
|
||||
Duration: time.Since(input.StartTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ClientDisconnect: clientDisconnect,
|
||||
RequestID: resp.Header.Get("x-request-id"),
|
||||
Usage: *usage,
|
||||
Model: input.OriginalModel,
|
||||
UpstreamModel: input.RequestModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
Stream: input.RequestStream,
|
||||
Duration: time.Since(input.StartTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ClientDisconnect: clientDisconnect,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -379,6 +381,10 @@ func (s *GatewayService) handleStreamingResponseAnthropicAPIKeyPassthrough(
|
||||
startTime time.Time,
|
||||
model string,
|
||||
) (*streamingResult, error) {
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
if s.rateLimitService != nil {
|
||||
s.rateLimitService.UpdateSessionWindow(ctx, account, resp.Header)
|
||||
}
|
||||
@@ -532,6 +538,7 @@ func (s *GatewayService) handleStreamingResponseAnthropicAPIKeyPassthrough(
|
||||
line := ev.line
|
||||
if data, ok := extractAnthropicSSEDataLine(line); ok {
|
||||
trimmed := strings.TrimSpace(data)
|
||||
observer.ObserveAnthropic([]byte(trimmed))
|
||||
if anthropicStreamEventIsTerminal("", trimmed) {
|
||||
sawTerminalEvent = true
|
||||
}
|
||||
@@ -782,6 +789,11 @@ func (s *GatewayService) handleNonStreamingResponseAnthropicAPIKeyPassthrough(
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
observer.ObserveAnthropic(body)
|
||||
|
||||
if resp.StatusCode >= http.StatusOK && resp.StatusCode < http.StatusMultipleChoices {
|
||||
var raw json.RawMessage
|
||||
|
||||
@@ -92,6 +92,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
if parsed == nil {
|
||||
return nil, fmt.Errorf("parse request: empty request")
|
||||
}
|
||||
beginUpstreamResponseModelObservation(c)
|
||||
|
||||
// Web Search 模拟:纯 web_search 请求时,直接调用搜索 API 构造响应
|
||||
if account != nil && s.shouldEmulateWebSearch(ctx, account, parsed.GroupID, parsed.Body.Bytes()) {
|
||||
@@ -845,7 +846,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
}
|
||||
// 流中断(缺失 terminal 事件、读错误、数据间隔超时等)时保留已观测到的
|
||||
// usage 与错误一起返回,handler 在错误处理完成后照常提交 usage 记录。
|
||||
if partial := partialStreamUsageResult(resp, streamResult, originalModel, mappedModel, startTime, err); partial != nil {
|
||||
if partial := partialStreamUsageResult(c, resp, streamResult, originalModel, mappedModel, startTime, err); partial != nil {
|
||||
return partial, err
|
||||
}
|
||||
return nil, err
|
||||
@@ -861,14 +862,16 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
|
||||
}
|
||||
|
||||
return &ForwardResult{
|
||||
RequestID: resp.Header.Get("x-request-id"),
|
||||
Usage: *usage,
|
||||
Model: originalModel, // 使用原始模型用于计费和日志
|
||||
UpstreamModel: mappedModel,
|
||||
Stream: reqStream,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ClientDisconnect: clientDisconnect,
|
||||
RequestID: resp.Header.Get("x-request-id"),
|
||||
Usage: *usage,
|
||||
Model: originalModel, // 使用原始模型用于计费和日志
|
||||
UpstreamModel: mappedModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
Stream: reqStream,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ClientDisconnect: clientDisconnect,
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -595,12 +595,16 @@ type ForwardResult struct {
|
||||
Model string
|
||||
// UpstreamModel is the actual upstream model after mapping.
|
||||
// Prefer empty when it is identical to Model; persistence normalizes equal values away as no-op mappings.
|
||||
UpstreamModel string
|
||||
Stream bool
|
||||
Duration time.Duration
|
||||
FirstTokenMs *int // 首字时间(流式请求)
|
||||
ClientDisconnect bool // 客户端是否在流式传输过程中断开
|
||||
ReasoningEffort *string
|
||||
UpstreamModel string
|
||||
// UpstreamResponseModel is captured from the raw successful upstream
|
||||
// response before any client-facing rewrite or protocol conversion.
|
||||
UpstreamResponseModel string
|
||||
UpstreamResponseModelConflict bool
|
||||
Stream bool
|
||||
Duration time.Duration
|
||||
FirstTokenMs *int // 首字时间(流式请求)
|
||||
ClientDisconnect bool // 客户端是否在流式传输过程中断开
|
||||
ReasoningEffort *string
|
||||
|
||||
// 图片生成计费字段(图片生成模型使用)
|
||||
ImageCount int // 生成的图片数量
|
||||
|
||||
@@ -665,7 +665,7 @@ func (u *ClaudeUsage) hasObservedTokens() bool {
|
||||
//
|
||||
// 不变式:UpstreamFailoverError 必须保持 result=nil——failover 重试成功后按成功请求
|
||||
// 计费,若同时返回部分 usage 会造成双重计费,此处显式拦截兜底。
|
||||
func partialStreamUsageResult(resp *http.Response, streamResult *streamingResult, model, upstreamModel string, startTime time.Time, err error) *ForwardResult {
|
||||
func partialStreamUsageResult(c *gin.Context, resp *http.Response, streamResult *streamingResult, model, upstreamModel string, startTime time.Time, err error) *ForwardResult {
|
||||
if streamResult == nil || !streamResult.usage.hasObservedTokens() {
|
||||
return nil
|
||||
}
|
||||
@@ -674,18 +674,24 @@ func partialStreamUsageResult(resp *http.Response, streamResult *streamingResult
|
||||
return nil
|
||||
}
|
||||
return &ForwardResult{
|
||||
RequestID: resp.Header.Get("x-request-id"),
|
||||
Usage: *streamResult.usage,
|
||||
Model: model,
|
||||
UpstreamModel: upstreamModel,
|
||||
Stream: true,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: streamResult.firstTokenMs,
|
||||
ClientDisconnect: streamResult.clientDisconnect,
|
||||
RequestID: resp.Header.Get("x-request-id"),
|
||||
Usage: *streamResult.usage,
|
||||
Model: model,
|
||||
UpstreamModel: upstreamModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
Stream: true,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: streamResult.firstTokenMs,
|
||||
ClientDisconnect: streamResult.clientDisconnect,
|
||||
}
|
||||
}
|
||||
|
||||
func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, startTime time.Time, originalModel, mappedModel string, mimicClaudeCode bool) (*streamingResult, error) {
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
// 更新5h窗口状态
|
||||
s.rateLimitService.UpdateSessionWindow(ctx, account, resp.Header)
|
||||
|
||||
@@ -881,6 +887,7 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http
|
||||
}
|
||||
|
||||
eventType, _ := event["type"].(string)
|
||||
observer.ObserveAnthropic([]byte(dataLine))
|
||||
if eventName == "" {
|
||||
eventName = eventType
|
||||
}
|
||||
@@ -1376,6 +1383,11 @@ func (s *GatewayService) handleNonStreamingResponse(ctx context.Context, resp *h
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
observer.ObserveAnthropic(body)
|
||||
|
||||
// 解析usage
|
||||
var response struct {
|
||||
|
||||
@@ -1052,6 +1052,16 @@ func (s *GatewayService) buildRecordUsageLog(
|
||||
) *UsageLog {
|
||||
durationMs := int(result.Duration.Milliseconds())
|
||||
requestID := resolveUsageBillingRequestID(ctx, result.RequestID)
|
||||
sentModel := upstreamSentModel(result.Model, result.UpstreamModel)
|
||||
if result.UpstreamResponseModelConflict {
|
||||
slog.Warn("upstream_response_model_conflict",
|
||||
"platform", account.Platform,
|
||||
"account_id", account.ID,
|
||||
"request_id", requestID,
|
||||
"sent_model", sentModel,
|
||||
"selected_response_model", strings.TrimSpace(result.UpstreamResponseModel),
|
||||
)
|
||||
}
|
||||
usageLog := &UsageLog{
|
||||
UserID: user.ID,
|
||||
APIKeyID: apiKey.ID,
|
||||
@@ -1060,6 +1070,8 @@ func (s *GatewayService) buildRecordUsageLog(
|
||||
Model: result.Model,
|
||||
RequestedModel: requestedModel,
|
||||
UpstreamModel: optionalTrimmedStringPtr(result.UpstreamModel),
|
||||
UpstreamResponseModel: optionalTrimmedStringPtr(result.UpstreamResponseModel),
|
||||
UpstreamModelMismatch: upstreamModelMismatch(sentModel, result.UpstreamResponseModel),
|
||||
ReasoningEffort: result.ReasoningEffort,
|
||||
InboundEndpoint: optionalTrimmedStringPtr(input.InboundEndpoint),
|
||||
UpstreamEndpoint: optionalTrimmedStringPtr(input.UpstreamEndpoint),
|
||||
|
||||
@@ -581,6 +581,7 @@ func (s *GeminiMessagesCompatService) SelectAccountForAIStudioEndpoints(ctx cont
|
||||
}
|
||||
|
||||
func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte) (*ForwardResult, error) {
|
||||
beginUpstreamResponseModelObservation(c)
|
||||
startTime := time.Now()
|
||||
|
||||
var req struct {
|
||||
@@ -1072,6 +1073,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)
|
||||
upstreamResponseModelObserverFromContext(c).ObserveGemini(collectedBytes)
|
||||
claudeResp, usageObj2 := convertGeminiToClaudeMessage(collected, originalModel, collectedBytes, false)
|
||||
c.JSON(http.StatusOK, claudeResp)
|
||||
usage = usageObj2
|
||||
@@ -1095,16 +1097,18 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
|
||||
}
|
||||
|
||||
return &ForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: *usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: mappedModel,
|
||||
Stream: req.Stream,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ImageCount: imageCount,
|
||||
ImageSize: imageSize,
|
||||
ImageInputSize: imageInputSize,
|
||||
RequestID: requestID,
|
||||
Usage: *usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: mappedModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
Stream: req.Stream,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ImageCount: imageCount,
|
||||
ImageSize: imageSize,
|
||||
ImageInputSize: imageInputSize,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -1117,6 +1121,7 @@ func isGeminiSignatureRelatedError(respBody []byte) bool {
|
||||
}
|
||||
|
||||
func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.Context, account *Account, originalModel string, action string, stream bool, body []byte) (*ForwardResult, error) {
|
||||
beginUpstreamResponseModelObservation(c)
|
||||
startTime := time.Now()
|
||||
|
||||
if strings.TrimSpace(originalModel) == "" {
|
||||
@@ -1603,6 +1608,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
||||
return nil, s.writeGoogleError(c, http.StatusBadGateway, "Failed to read upstream stream")
|
||||
}
|
||||
b, _ := json.Marshal(collected)
|
||||
upstreamResponseModelObserverFromContext(c).ObserveGemini(b)
|
||||
c.Data(http.StatusOK, "application/json", b)
|
||||
usage = usageObj
|
||||
} else {
|
||||
@@ -1627,16 +1633,18 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
|
||||
}
|
||||
|
||||
return &ForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: *usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: mappedModel,
|
||||
Stream: stream,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ImageCount: imageCount,
|
||||
ImageSize: imageSize,
|
||||
ImageInputSize: imageInputSize,
|
||||
RequestID: requestID,
|
||||
Usage: *usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: mappedModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
Stream: stream,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ImageCount: imageCount,
|
||||
ImageSize: imageSize,
|
||||
ImageInputSize: imageInputSize,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -1999,6 +2007,11 @@ func (s *GeminiMessagesCompatService) handleNonStreamingResponse(c *gin.Context,
|
||||
if err != nil {
|
||||
return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response")
|
||||
}
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
observer.ObserveGemini(unwrappedBody)
|
||||
|
||||
var geminiResp map[string]any
|
||||
if err := json.Unmarshal(unwrappedBody, &geminiResp); err != nil {
|
||||
@@ -2082,6 +2095,11 @@ func (s *GeminiMessagesCompatService) handleStreamingResponse(c *gin.Context, re
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
observer.ObserveGemini(unwrappedBytes)
|
||||
|
||||
var geminiResp map[string]any
|
||||
if err := json.Unmarshal(unwrappedBytes, &geminiResp); err != nil {
|
||||
@@ -2583,6 +2601,11 @@ func (s *GeminiMessagesCompatService) handleNativeNonStreamingResponse(c *gin.Co
|
||||
respBody = unwrappedBody
|
||||
}
|
||||
}
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
observer.ObserveGemini(respBody)
|
||||
|
||||
responseheaders.WriteFilteredHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter)
|
||||
|
||||
@@ -2631,6 +2654,10 @@ func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Conte
|
||||
|
||||
reader := bufio.NewReader(resp.Body)
|
||||
usage := &ClaudeUsage{}
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
var firstTokenMs *int
|
||||
|
||||
for {
|
||||
@@ -2661,6 +2688,7 @@ func (s *GeminiMessagesCompatService) handleNativeStreamingResponse(c *gin.Conte
|
||||
if u := extractGeminiUsage(rawBytes); u != nil {
|
||||
usage = u
|
||||
}
|
||||
observer.ObserveGemini(rawBytes)
|
||||
|
||||
if firstTokenMs == nil {
|
||||
ms := int(time.Since(startTime).Milliseconds())
|
||||
|
||||
@@ -1204,7 +1204,7 @@ func extractGrokMediaVideoRequestID(body []byte) string {
|
||||
if len(body) == 0 || !gjson.ValidBytes(body) {
|
||||
return ""
|
||||
}
|
||||
for _, path := range []string{"request_id", "id", "data.request_id", "data.id", "video.request_id", "video.id"} {
|
||||
for _, path := range []string{"request_id", "id", "data.request_id", "data.id", "video.request_id", "video.id", "task_id", "data.task_id", "video.task_id"} {
|
||||
if id := strings.TrimSpace(gjson.GetBytes(body, path).String()); id != "" {
|
||||
return id
|
||||
}
|
||||
|
||||
@@ -120,6 +120,23 @@ func OpenAIImageGenerationIntentFromContext(ctx context.Context) bool {
|
||||
return ok && enabled
|
||||
}
|
||||
|
||||
// WithOpenAIImagesEndpoint 标记请求从 /v1/images/* 专用生图端点入站。
|
||||
func WithOpenAIImagesEndpoint(ctx context.Context) context.Context {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
return context.WithValue(ctx, ctxkey.OpenAIImagesEndpoint, true)
|
||||
}
|
||||
|
||||
// OpenAIImagesEndpointFromContext 报告请求是否来自 /v1/images/*。
|
||||
func OpenAIImagesEndpointFromContext(ctx context.Context) bool {
|
||||
if ctx == nil {
|
||||
return false
|
||||
}
|
||||
enabled, ok := ctx.Value(ctxkey.OpenAIImagesEndpoint).(bool)
|
||||
return ok && enabled
|
||||
}
|
||||
|
||||
func resolveFinalAntigravityModelKey(ctx context.Context, account *Account, requestedModel string) string {
|
||||
modelKey := mapAntigravityModel(account, requestedModel)
|
||||
if modelKey == "" {
|
||||
|
||||
@@ -2,12 +2,17 @@ package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
@@ -72,6 +77,172 @@ func TestStreamFailedEventCapacityShedRetriesOnSameAccount(t *testing.T) {
|
||||
require.False(t, openAIStreamFailedEventRetryableOnSameAccount(nonPool, other, "boom"))
|
||||
}
|
||||
|
||||
// 上游降载的真实序列是「event: error → event: response.failed」。error 帧不算
|
||||
// 客户端输出:若把它当首输出 flush,clientOutputStarted 被固化,随后的 failed
|
||||
// 事件就进不了 pre-output failover 分支,只能把致命错误原样转发给客户端。
|
||||
func TestOpenAIStreamErrorFrameDoesNotStartClientOutput(t *testing.T) {
|
||||
cases := []struct {
|
||||
data string
|
||||
eventType string
|
||||
want bool
|
||||
}{
|
||||
{`{"type":"error","error":{"code":"server_is_overloaded","message":"overloaded"}}`, "error", false},
|
||||
{`{"type":"error","error":{"code":"slow_down","message":"slow down"}}`, "error", false},
|
||||
{`{"type":"error","error":{"type":"rate_limit_error","code":"rate_limit_exceeded","message":"limited"}}`, "error", false},
|
||||
// 不可重试类错误帧维持原样转发(不进 failover),保留上游错误细节。
|
||||
{`{"type":"error","error":{"type":"invalid_request_error","code":"content_policy_violation","message":"blocked"}}`, "error", true},
|
||||
{`{"type":"response.failed","response":{"error":{"code":"server_is_overloaded"}}}`, "response.failed", false},
|
||||
{`{"type":"response.created","response":{"id":"resp_1"}}`, "response.created", false},
|
||||
{`{"type":"response.in_progress","response":{"id":"resp_1"}}`, "response.in_progress", false},
|
||||
{`{"type":"response.output_text.delta","delta":"hi"}`, "response.output_text.delta", true},
|
||||
{`[DONE]`, "", true},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
require.Equal(t, tc.want, openAIStreamDataStartsClientOutput(tc.data, tc.eventType), "data=%s type=%s", tc.data, tc.eventType)
|
||||
}
|
||||
}
|
||||
|
||||
// 回归用例(真实上游降载序列):created → in_progress → error 帧 → response.failed。
|
||||
// 期望仍然走 pre-output failover(同账号重试 + 请求级瞬时标记),且不向客户端写出任何字节。
|
||||
func TestOpenAIStreamCapacityShedErrorFramePrecedingFailedStillFailsOver(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
cfg := &config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
}
|
||||
svc := &OpenAIGatewayService{cfg: cfg}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||||
"event: response.created",
|
||||
`data: {"type":"response.created","response":{"id":"resp_1"},"sequence_number":0}`,
|
||||
"",
|
||||
"event: response.in_progress",
|
||||
`data: {"type":"response.in_progress","response":{"id":"resp_1"},"sequence_number":1}`,
|
||||
"",
|
||||
"event: error",
|
||||
`data: {"type":"error","error":{"type":"service_unavailable_error","code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later."},"sequence_number":2}`,
|
||||
"",
|
||||
"event: response.failed",
|
||||
`data: {"type":"response.failed","response":{"id":"resp_1","status":"failed","error":{"code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later."}},"sequence_number":3}`,
|
||||
"",
|
||||
}, "\n"))),
|
||||
Header: http.Header{"X-Request-Id": []string{"rid-shed-error-then-failed"}},
|
||||
}
|
||||
|
||||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Name: "acc"}, time.Now(), "model", "model")
|
||||
require.Error(t, err)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.ErrorAs(t, err, &failoverErr)
|
||||
require.True(t, failoverErr.RetryableOnSameAccount)
|
||||
require.True(t, failoverErr.RequestScopedTransient)
|
||||
require.False(t, c.Writer.Written())
|
||||
require.Empty(t, rec.Body.String())
|
||||
}
|
||||
|
||||
// 流中途(已有真实输出)降载时无法再 failover,此时必须把降载码改写为客户端
|
||||
// 可重试的 server_error 再转发——Codex 对 server_is_overloaded/slow_down 判致命
|
||||
// 并终止会话,对其余错误码执行内置退避重试。消息原样保留。
|
||||
func TestOpenAIStreamCapacityShedAfterOutputRewritesCodeForClient(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
cfg := &config.Config{
|
||||
Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize},
|
||||
}
|
||||
svc := &OpenAIGatewayService{cfg: cfg}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Body: io.NopCloser(strings.NewReader(strings.Join([]string{
|
||||
"event: response.created",
|
||||
`data: {"type":"response.created","response":{"id":"resp_1"}}`,
|
||||
"",
|
||||
"event: response.output_text.delta",
|
||||
`data: {"type":"response.output_text.delta","delta":"partial"}`,
|
||||
"",
|
||||
"event: error",
|
||||
`data: {"type":"error","error":{"type":"service_unavailable_error","code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later."},"sequence_number":2}`,
|
||||
"",
|
||||
"event: response.failed",
|
||||
`data: {"type":"response.failed","response":{"id":"resp_1","status":"failed","error":{"code":"server_is_overloaded","message":"Our servers are currently overloaded. Please try again later."}},"sequence_number":3}`,
|
||||
"",
|
||||
}, "\n"))),
|
||||
Header: http.Header{"X-Request-Id": []string{"rid-shed-after-output"}},
|
||||
}
|
||||
|
||||
_, err := svc.handleStreamingResponse(c.Request.Context(), resp, c, &Account{ID: 1, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Name: "acc"}, time.Now(), "model", "model")
|
||||
require.Error(t, err)
|
||||
var failoverErr *UpstreamFailoverError
|
||||
require.False(t, errors.As(err, &failoverErr))
|
||||
|
||||
body := rec.Body.String()
|
||||
require.Contains(t, body, "partial")
|
||||
require.Contains(t, body, "event: response.failed")
|
||||
require.Contains(t, body, `"code":"server_error"`)
|
||||
require.NotContains(t, body, "server_is_overloaded")
|
||||
require.Contains(t, body, "Our servers are currently overloaded")
|
||||
}
|
||||
|
||||
// helper 单测:只有降载码被改写,其余错误码(尤其 rate_limit_exceeded,客户端
|
||||
// 依赖其原码解析重试延时)必须原样保留。
|
||||
func TestSanitizeOpenAICapacityShedErrorCodeForClient(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
payload string
|
||||
wantChanged bool
|
||||
wantContain string
|
||||
}{
|
||||
{
|
||||
name: "failed事件嵌套code改写",
|
||||
payload: `{"type":"response.failed","response":{"error":{"code":"server_is_overloaded","message":"overloaded"}}}`,
|
||||
wantChanged: true,
|
||||
wantContain: `"code":"server_error"`,
|
||||
},
|
||||
{
|
||||
name: "error帧裸code改写",
|
||||
payload: `{"type":"error","error":{"code":"slow_down","message":"slow down"}}`,
|
||||
wantChanged: true,
|
||||
wantContain: `"code":"server_error"`,
|
||||
},
|
||||
{
|
||||
name: "rate_limit不改写",
|
||||
payload: `{"type":"response.failed","response":{"error":{"code":"rate_limit_exceeded","message":"try again in 3s"}}}`,
|
||||
wantChanged: false,
|
||||
wantContain: `"code":"rate_limit_exceeded"`,
|
||||
},
|
||||
{
|
||||
name: "普通server_error不改写",
|
||||
payload: `{"type":"response.failed","response":{"error":{"code":"server_error","message":"boom"}}}`,
|
||||
wantChanged: false,
|
||||
wantContain: `"code":"server_error"`,
|
||||
},
|
||||
{
|
||||
name: "非JSON不改写",
|
||||
payload: `not-json`,
|
||||
wantChanged: false,
|
||||
wantContain: `not-json`,
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
out, changed := sanitizeOpenAICapacityShedErrorCodeForClient([]byte(tc.payload))
|
||||
require.Equal(t, tc.wantChanged, changed)
|
||||
require.Contains(t, string(out), tc.wantContain)
|
||||
if changed {
|
||||
require.NotContains(t, string(out), "server_is_overloaded")
|
||||
require.NotContains(t, string(out), "slow_down")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 出站身份的版本声明只能有一个来源:UA 的版本段、version 头、探针版本三处必须同源,
|
||||
// 各自硬编码会漂移成互相矛盾的身份,而自相矛盾或陈旧的身份会被上游优先降载。
|
||||
func TestCodexOutboundVersionHasSingleSource(t *testing.T) {
|
||||
|
||||
@@ -0,0 +1,100 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestNormalizeOpenAICompactRequestBodyPreservesServiceTier(t *testing.T) {
|
||||
body := []byte(`{
|
||||
"model":"gpt-5.6-sol",
|
||||
"input":[{"type":"message","role":"user","content":"hello"}],
|
||||
"service_tier":"priority",
|
||||
"prompt_cache_key":"compact-cache-key",
|
||||
"store":false,
|
||||
"stream":true
|
||||
}`)
|
||||
|
||||
normalized, changed, err := normalizeOpenAICompactRequestBody(body)
|
||||
require.NoError(t, err)
|
||||
require.True(t, changed)
|
||||
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(normalized, "model").String())
|
||||
require.Equal(t, "priority", gjson.GetBytes(normalized, "service_tier").String())
|
||||
require.False(t, gjson.GetBytes(normalized, "prompt_cache_key").Exists())
|
||||
require.False(t, gjson.GetBytes(normalized, "store").Exists())
|
||||
require.False(t, gjson.GetBytes(normalized, "stream").Exists())
|
||||
}
|
||||
|
||||
func TestOpenAIOAuthCompactHTTPBuildersUsePreservedServiceTierInRoutingHint(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
body := []byte(`{
|
||||
"model":"gpt-5.6-sol",
|
||||
"input":[{"type":"message","role":"user","content":"hello"}],
|
||||
"service_tier":"priority",
|
||||
"stream":true
|
||||
}`)
|
||||
normalized, changed, err := normalizeOpenAICompactRequestBody(body)
|
||||
require.NoError(t, err)
|
||||
require.True(t, changed)
|
||||
require.Equal(t, "priority", gjson.GetBytes(normalized, "service_tier").String())
|
||||
|
||||
account := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"chatgpt_account_id": "test-account",
|
||||
},
|
||||
}
|
||||
svc := &OpenAIGatewayService{}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
build func(*gin.Context) (*http.Request, error)
|
||||
}{
|
||||
{
|
||||
name: "ordinary",
|
||||
build: func(c *gin.Context) (*http.Request, error) {
|
||||
return svc.buildUpstreamRequest(
|
||||
context.Background(), c, account, normalized, "test-token",
|
||||
false, "", true,
|
||||
)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "passthrough",
|
||||
build: func(c *gin.Context) (*http.Request, error) {
|
||||
return svc.buildUpstreamRequestOpenAIPassthrough(
|
||||
context.Background(), c, account, normalized, "test-token",
|
||||
)
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(
|
||||
http.MethodPost,
|
||||
"/v1/responses/compact",
|
||||
bytes.NewReader(normalized),
|
||||
)
|
||||
|
||||
req, buildErr := tt.build(c)
|
||||
require.NoError(t, buildErr)
|
||||
require.Equal(
|
||||
t,
|
||||
"model=gpt-5.6-sol;tier=priority",
|
||||
req.Header.Get(openAICodexRoutingHintHeader),
|
||||
)
|
||||
require.Equal(t, "priority", gjson.GetBytes(normalized, "service_tier").String())
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -58,6 +58,8 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions(
|
||||
promptCacheKey string,
|
||||
defaultMappedModel string,
|
||||
) (*OpenAIForwardResult, error) {
|
||||
beginUpstreamResponseModelObservation(c)
|
||||
|
||||
restrictionResult := s.detectCodexClientRestriction(c, account, body)
|
||||
logCodexCLIOnlyDetection(ctx, c, account, getAPIKeyIDFromContext(c), restrictionResult, body)
|
||||
if restrictionResult.Enabled && !restrictionResult.Matched {
|
||||
@@ -419,6 +421,11 @@ func (s *OpenAIGatewayService) handleChatBufferedStreamingResponse(
|
||||
writeChatCompletionsError(c, http.StatusBadGateway, "api_error", "Upstream stream ended without a terminal response event")
|
||||
return nil, fmt.Errorf("upstream stream ended without terminal event")
|
||||
}
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
observer.Observe(finalResponse.Model, true)
|
||||
if strings.TrimSpace(finalResponse.Status) == "failed" {
|
||||
payload, _ := json.Marshal(gin.H{"type": "response.failed", "response": finalResponse})
|
||||
// cyber_policy 致命不可重试:不 failover,以 Chat Completions 错误格式回写(F4),
|
||||
@@ -477,13 +484,15 @@ func (s *OpenAIGatewayService) handleChatBufferedStreamingResponse(
|
||||
c.JSON(http.StatusOK, chatResp)
|
||||
|
||||
result := &OpenAIForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
BillingModel: billingModel,
|
||||
UpstreamModel: upstreamModel,
|
||||
Stream: false,
|
||||
Duration: time.Since(startTime),
|
||||
RequestID: requestID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
BillingModel: billingModel,
|
||||
UpstreamModel: upstreamModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
Stream: false,
|
||||
Duration: time.Since(startTime),
|
||||
}
|
||||
// Grok chat bridge: bill native search tools found in the terminal Responses body.
|
||||
if account != nil && account.IsGrok() && finalResponse != nil {
|
||||
@@ -530,6 +539,10 @@ func (s *OpenAIGatewayService) handleChatStreamingResponse(
|
||||
searchCount := 0
|
||||
streamSearchSeen := make(map[string]struct{})
|
||||
countSearch := account != nil && account.IsGrok()
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
|
||||
scanner := s.newUpstreamSSEScanner(resp.Body)
|
||||
|
||||
@@ -549,14 +562,16 @@ func (s *OpenAIGatewayService) handleChatStreamingResponse(
|
||||
|
||||
resultWithUsage := func() *OpenAIForwardResult {
|
||||
out := &OpenAIForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
BillingModel: billingModel,
|
||||
UpstreamModel: upstreamModel,
|
||||
Stream: true,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
RequestID: requestID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
BillingModel: billingModel,
|
||||
UpstreamModel: upstreamModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
Stream: true,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
}
|
||||
if searchCount > 0 {
|
||||
out.SearchCount = searchCount
|
||||
@@ -582,6 +597,7 @@ func (s *OpenAIGatewayService) handleChatStreamingResponse(
|
||||
)
|
||||
return false
|
||||
}
|
||||
observer.ObserveOpenAI([]byte(payload), event.Type)
|
||||
refusalDetector.ObservePayload([]byte(payload))
|
||||
|
||||
isTerminalEvent := isOpenAICompatResponsesTerminalEvent(event.Type)
|
||||
|
||||
@@ -257,6 +257,10 @@ func (s *OpenAIGatewayService) streamRawChatCompletions(
|
||||
startTime time.Time,
|
||||
requestBodyLen int,
|
||||
) (*OpenAIForwardResult, error) {
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
requestID := resp.Header.Get("x-request-id")
|
||||
writeStreamHeaders := s.newStreamHeaderWriter(c, resp.Header)
|
||||
scanner := s.newUpstreamSSEScanner(resp.Body)
|
||||
@@ -306,6 +310,7 @@ func (s *OpenAIGatewayService) streamRawChatCompletions(
|
||||
if payload, ok := extractOpenAISSEDataLine(line); ok {
|
||||
trimmedPayload := strings.TrimSpace(payload)
|
||||
if trimmedPayload != "[DONE]" {
|
||||
observer.ObserveOpenAI([]byte(payload), strings.TrimSpace(gjson.Get(payload, "type").String()))
|
||||
usageOnlyChunk := isOpenAIChatUsageOnlyStreamChunk(payload)
|
||||
if u := extractCCStreamUsage(payload); u != nil {
|
||||
usage = *u
|
||||
@@ -360,16 +365,18 @@ func (s *OpenAIGatewayService) streamRawChatCompletions(
|
||||
}
|
||||
|
||||
return &OpenAIForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
BillingModel: billingModel,
|
||||
UpstreamModel: upstreamModel,
|
||||
ReasoningEffort: reasoningEffort,
|
||||
ServiceTier: serviceTier,
|
||||
Stream: true,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
RequestID: requestID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
BillingModel: billingModel,
|
||||
UpstreamModel: upstreamModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
ReasoningEffort: reasoningEffort,
|
||||
ServiceTier: serviceTier,
|
||||
Stream: true,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -429,6 +436,11 @@ func (s *OpenAIGatewayService) bufferRawChatCompletions(
|
||||
}
|
||||
return nil, fmt.Errorf("read upstream body: %w", err)
|
||||
}
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
observer.ObserveOpenAI(respBody, strings.TrimSpace(gjson.GetBytes(respBody, "type").String()))
|
||||
|
||||
var usage OpenAIUsage
|
||||
if parsedUsage, ok := extractOpenAIUsageFromJSONBytes(respBody); ok {
|
||||
@@ -447,15 +459,17 @@ func (s *OpenAIGatewayService) bufferRawChatCompletions(
|
||||
_, _ = c.Writer.Write(respBody)
|
||||
|
||||
return &OpenAIForwardResult{
|
||||
RequestID: requestID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
BillingModel: billingModel,
|
||||
UpstreamModel: upstreamModel,
|
||||
ReasoningEffort: reasoningEffort,
|
||||
ServiceTier: serviceTier,
|
||||
Stream: false,
|
||||
Duration: time.Since(startTime),
|
||||
RequestID: requestID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
BillingModel: billingModel,
|
||||
UpstreamModel: upstreamModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
ReasoningEffort: reasoningEffort,
|
||||
ServiceTier: serviceTier,
|
||||
Stream: false,
|
||||
Duration: time.Since(startTime),
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -345,12 +345,26 @@ func isOpenAIOAuthInputTokensUnsupported(statusCode int, body []byte) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// OAuth's platform endpoint can be blocked by an upstream proxy before it
|
||||
// reaches the API and return an HTML 403 page without a structured error.
|
||||
// Treat that endpoint-level response like the other unsupported cases so
|
||||
// count_tokens remains a local, non-health-affecting convenience request.
|
||||
if statusCode == http.StatusForbidden && isHTMLResponse(body) {
|
||||
return true
|
||||
}
|
||||
|
||||
return strings.Contains(msg, "input_tokens") &&
|
||||
(strings.Contains(msg, "not found") ||
|
||||
strings.Contains(msg, "not supported") ||
|
||||
strings.Contains(msg, "unsupported"))
|
||||
}
|
||||
|
||||
func isHTMLResponse(body []byte) bool {
|
||||
trimmed := strings.TrimSpace(strings.ToLower(string(body)))
|
||||
return strings.HasPrefix(trimmed, "<!doctype html") ||
|
||||
strings.HasPrefix(trimmed, "<html")
|
||||
}
|
||||
|
||||
func estimateOpenAIInputTokens(req openAIInputTokensCountRequest) (int, error) {
|
||||
codec, err := openAIInputTokensCodecForModel(req.Model)
|
||||
if err != nil {
|
||||
|
||||
@@ -123,6 +123,11 @@ func TestOpenAIGatewayService_ForwardCountTokensAsAnthropic_OAuthFallsBackWhenPl
|
||||
statusCode: http.StatusForbidden,
|
||||
body: `{"error":{"type":"invalid_request_error","code":"missing_scope","message":"Missing scopes: api.responses.write"}}`,
|
||||
},
|
||||
{
|
||||
name: "403_html_proxy_page",
|
||||
statusCode: http.StatusForbidden,
|
||||
body: "<!doctype html><html><body>Forbidden</body></html>",
|
||||
},
|
||||
{
|
||||
name: "404_input_tokens_unsupported",
|
||||
statusCode: http.StatusNotFound,
|
||||
|
||||
@@ -19,6 +19,7 @@ import (
|
||||
|
||||
// Forward forwards request to OpenAI API
|
||||
func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, account *Account, body []byte) (*OpenAIForwardResult, error) {
|
||||
beginUpstreamResponseModelObservation(c)
|
||||
clearGrokResponsesClientToolMapping(c)
|
||||
clearOpenAIResponsesNamespaceNames(c)
|
||||
startTime := time.Now()
|
||||
@@ -46,6 +47,16 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
if normalized {
|
||||
body = normalizedBody
|
||||
}
|
||||
// 在分流到 passthrough / Codex transform / 原生 ChatCompletions 之前统一修正
|
||||
// 显式为 null 的工具 Schema type,否则 upstream 的 400 会被归一成可重试的 502,
|
||||
// 同一份坏定义在账号池里反复重放。
|
||||
sanitizedToolBody, toolSchemaSanitized, toolSchemaErr := sanitizeOpenAIResponsesToolParameterTypes(body)
|
||||
if toolSchemaErr != nil {
|
||||
return nil, fmt.Errorf("sanitize OpenAI Responses tool parameters: %w", toolSchemaErr)
|
||||
}
|
||||
if toolSchemaSanitized {
|
||||
body = sanitizedToolBody
|
||||
}
|
||||
if account.IsOpenAIOAuth() && isOpenAIResponsesLiteHeader(c.GetHeader(responsesLiteHeader)) {
|
||||
liteBody, changed, liteErr := normalizeOpenAIResponsesLiteToolsPayload(body)
|
||||
if liteErr != nil {
|
||||
@@ -967,18 +978,20 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco
|
||||
}
|
||||
|
||||
forwardResult := &OpenAIForwardResult{
|
||||
RequestID: resp.Header.Get("x-request-id"),
|
||||
ResponseID: responseID,
|
||||
Usage: *usage,
|
||||
Model: originalModel,
|
||||
BillingModel: billingModel,
|
||||
UpstreamModel: upstreamModel,
|
||||
ServiceTier: serviceTier,
|
||||
ReasoningEffort: reasoningEffort,
|
||||
Stream: reqStream,
|
||||
OpenAIWSMode: false,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
RequestID: resp.Header.Get("x-request-id"),
|
||||
ResponseID: responseID,
|
||||
Usage: *usage,
|
||||
Model: originalModel,
|
||||
BillingModel: billingModel,
|
||||
UpstreamModel: upstreamModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
ServiceTier: serviceTier,
|
||||
ReasoningEffort: reasoningEffort,
|
||||
Stream: reqStream,
|
||||
OpenAIWSMode: false,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
}
|
||||
if imageCount > 0 {
|
||||
forwardResult.ImageCount = imageCount
|
||||
@@ -1068,7 +1081,6 @@ func (s *OpenAIGatewayService) buildUpstreamRequest(ctx context.Context, c *gin.
|
||||
req.Header.Del("OpenAI-Beta")
|
||||
req.Header.Del("originator")
|
||||
} else {
|
||||
req.Header.Set("OpenAI-Beta", "responses=experimental")
|
||||
req.Header.Set("originator", resolveOpenAIUpstreamOriginator(c, isCodexCLI))
|
||||
}
|
||||
apiKeyID := getAPIKeyIDFromContext(c)
|
||||
@@ -1120,6 +1132,8 @@ func (s *OpenAIGatewayService) buildUpstreamRequest(ctx context.Context, c *gin.
|
||||
|
||||
// 账号级请求头覆写(仅 openai api_key 账号启用时生效;OAuth 路径 no-op)
|
||||
account.ApplyHeaderOverrides(req.Header)
|
||||
setOpenAICodexRoutingHintFromBody(req.Header, account, body)
|
||||
logOpenAIRoutingDiagnosticsFromBody(ctx, account, "http", req.Header, body, "not_applicable")
|
||||
|
||||
return req, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestShouldFailoverUpstreamError_405IsFailoverEligible(t *testing.T) {
|
||||
svc := &OpenAIGatewayService{}
|
||||
|
||||
assert.True(t, svc.shouldFailoverUpstreamError(http.StatusMethodNotAllowed),
|
||||
"405 should trigger failover so sticky sessions can escape to healthy accounts")
|
||||
}
|
||||
|
||||
func TestShouldFailoverUpstreamError_ExistingCodesStillWork(t *testing.T) {
|
||||
svc := &OpenAIGatewayService{}
|
||||
|
||||
failoverCodes := []int{401, 402, 403, 405, 429, 529, 500, 502, 503, 504}
|
||||
for _, code := range failoverCodes {
|
||||
assert.True(t, svc.shouldFailoverUpstreamError(code), "status %d should trigger failover", code)
|
||||
}
|
||||
|
||||
nonFailoverCodes := []int{200, 201, 400, 404, 408, 422}
|
||||
for _, code := range nonFailoverCodes {
|
||||
assert.False(t, svc.shouldFailoverUpstreamError(code), "status %d should NOT trigger failover", code)
|
||||
}
|
||||
}
|
||||
@@ -1299,6 +1299,51 @@ func TestForwardGrokMediaVideoGenerationReturnsUsageAndResponseID(t *testing.T)
|
||||
require.Equal(t, 10, result.VideoDurationSeconds)
|
||||
}
|
||||
|
||||
func TestForwardGrokMediaVideoGenerationReturnsTaskIDAsResponseID(t *testing.T) {
|
||||
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
body := []byte(`{"model":"grok-imagine-video","prompt":"waves"}`)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/generations", bytes.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
account := &Account{
|
||||
ID: 63,
|
||||
Name: "grok",
|
||||
Platform: PlatformGrok,
|
||||
Type: AccountTypeAPIKey,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "api-key",
|
||||
"base_url": "https://xai.test/v1",
|
||||
},
|
||||
}
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(`{"task_id":"video-task-123"}`)),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{httpUpstream: upstream}
|
||||
|
||||
result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideosGenerations, "", body, "application/json")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "video-task-123", result.ResponseID)
|
||||
}
|
||||
|
||||
func TestExtractGrokMediaVideoRequestIDPreservesExistingPrecedence(t *testing.T) {
|
||||
body := []byte(`{
|
||||
"request_id":"request-id",
|
||||
"id":"id",
|
||||
"task_id":"task-id",
|
||||
"data":{"request_id":"data-request-id","id":"data-id","task_id":"data-task-id"},
|
||||
"video":{"request_id":"video-request-id","id":"video-id","task_id":"video-task-id"}
|
||||
}`)
|
||||
|
||||
require.Equal(t, "request-id", extractGrokMediaVideoRequestID(body))
|
||||
}
|
||||
|
||||
func TestForwardGrokMediaVideoGenerationPreservesImageToVideoModel(t *testing.T) {
|
||||
t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -34,6 +34,8 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic(
|
||||
promptCacheKey string,
|
||||
defaultMappedModel string,
|
||||
) (*OpenAIForwardResult, error) {
|
||||
beginUpstreamResponseModelObservation(c)
|
||||
|
||||
// 入口分流:APIKey 账号 + 上游不支持 Responses API → 走 CC 直转(与
|
||||
// ForwardAsChatCompletions 对称)。缺少此分流时,/v1/messages 入站请求
|
||||
// 会被无条件转为 Responses 格式发往上游 /v1/responses,导致只支持
|
||||
@@ -549,6 +551,11 @@ func (s *OpenAIGatewayService) handleAnthropicBufferedStreamingResponse(
|
||||
writeAnthropicError(c, http.StatusBadGateway, "api_error", "Upstream stream ended without a terminal response event")
|
||||
return nil, fmt.Errorf("upstream stream ended without terminal event")
|
||||
}
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
observer.Observe(finalResponse.Model, true)
|
||||
|
||||
if strings.TrimSpace(finalResponse.Status) == "failed" {
|
||||
payload, _ := json.Marshal(gin.H{"type": "response.failed", "response": finalResponse})
|
||||
@@ -602,14 +609,16 @@ func (s *OpenAIGatewayService) handleAnthropicBufferedStreamingResponse(
|
||||
c.JSON(http.StatusOK, anthropicResp)
|
||||
|
||||
result := &OpenAIForwardResult{
|
||||
RequestID: requestID,
|
||||
ResponseID: finalResponse.ID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
BillingModel: billingModel,
|
||||
UpstreamModel: upstreamModel,
|
||||
Stream: false,
|
||||
Duration: time.Since(startTime),
|
||||
RequestID: requestID,
|
||||
ResponseID: finalResponse.ID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
BillingModel: billingModel,
|
||||
UpstreamModel: upstreamModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
Stream: false,
|
||||
Duration: time.Since(startTime),
|
||||
}
|
||||
// Grok /v1/messages uses Responses upstream; count native search for surcharge.
|
||||
if account != nil && account.IsGrok() && finalResponse != nil {
|
||||
@@ -858,20 +867,26 @@ func (s *OpenAIGatewayService) handleAnthropicStreamingResponse(
|
||||
if intervalTicker != nil {
|
||||
intervalCh = intervalTicker.C
|
||||
}
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
|
||||
// resultWithUsage builds the final result snapshot.
|
||||
resultWithUsage := func() *OpenAIForwardResult {
|
||||
out := &OpenAIForwardResult{
|
||||
RequestID: requestID,
|
||||
ResponseID: responseID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
BillingModel: billingModel,
|
||||
UpstreamModel: upstreamModel,
|
||||
Stream: true,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ClientDisconnect: clientDisconnected,
|
||||
RequestID: requestID,
|
||||
ResponseID: responseID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
BillingModel: billingModel,
|
||||
UpstreamModel: upstreamModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
Stream: true,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
ClientDisconnect: clientDisconnected,
|
||||
}
|
||||
if searchCount > 0 {
|
||||
out.SearchCount = searchCount
|
||||
@@ -898,6 +913,7 @@ func (s *OpenAIGatewayService) handleAnthropicStreamingResponse(
|
||||
)
|
||||
return false
|
||||
}
|
||||
observer.ObserveOpenAI([]byte(payload), event.Type)
|
||||
|
||||
eventType := strings.TrimSpace(event.Type)
|
||||
isBareErrorEvent := eventType == "error"
|
||||
|
||||
@@ -278,17 +278,19 @@ func (s *OpenAIGatewayService) forwardOpenAIPassthrough(
|
||||
}
|
||||
|
||||
forwardResult := &OpenAIForwardResult{
|
||||
RequestID: resp.Header.Get("x-request-id"),
|
||||
ResponseID: responseID,
|
||||
Usage: *usage,
|
||||
Model: reqModel,
|
||||
UpstreamModel: upstreamPassthroughModel,
|
||||
ServiceTier: serviceTier,
|
||||
ReasoningEffort: reasoningEffort,
|
||||
Stream: reqStream,
|
||||
OpenAIWSMode: false,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
RequestID: resp.Header.Get("x-request-id"),
|
||||
ResponseID: responseID,
|
||||
Usage: *usage,
|
||||
Model: reqModel,
|
||||
UpstreamModel: upstreamPassthroughModel,
|
||||
UpstreamResponseModel: observedUpstreamResponseModel(c),
|
||||
UpstreamResponseModelConflict: observedUpstreamResponseModelConflict(c),
|
||||
ServiceTier: serviceTier,
|
||||
ReasoningEffort: reasoningEffort,
|
||||
Stream: reqStream,
|
||||
OpenAIWSMode: false,
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
}
|
||||
if imageCount > 0 {
|
||||
forwardResult.ImageCount = imageCount
|
||||
@@ -390,6 +392,10 @@ func (s *OpenAIGatewayService) buildUpstreamRequestOpenAIPassthrough(
|
||||
|
||||
// OAuth 透传到 ChatGPT internal API 时补齐必要头。
|
||||
if account.Type == AccountTypeOAuth {
|
||||
// Current Codex OAuth HTTP no longer negotiates the legacy Responses
|
||||
// experiment. Passthrough may receive it from an older client, so remove
|
||||
// only that token while preserving any independent beta negotiation.
|
||||
stripOpenAILegacyResponsesBeta(req.Header)
|
||||
promptCacheKey := strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String())
|
||||
req.Host = "chatgpt.com"
|
||||
if err := resolveAndSetOpenAIChatGPTAccountHeaders(ctx, s.accountRepo, req.Header, account); err != nil {
|
||||
@@ -410,9 +416,6 @@ func (s *OpenAIGatewayService) buildUpstreamRequestOpenAIPassthrough(
|
||||
} else if req.Header.Get("accept") == "" {
|
||||
req.Header.Set("accept", "text/event-stream")
|
||||
}
|
||||
if req.Header.Get("OpenAI-Beta") == "" {
|
||||
req.Header.Set("OpenAI-Beta", "responses=experimental")
|
||||
}
|
||||
if req.Header.Get("originator") == "" {
|
||||
req.Header.Set("originator", openai.CodexDefaultOriginator)
|
||||
}
|
||||
@@ -456,10 +459,43 @@ func (s *OpenAIGatewayService) buildUpstreamRequestOpenAIPassthrough(
|
||||
|
||||
// 账号级请求头覆写(仅 openai api_key 账号启用时生效;OAuth 路径 no-op)
|
||||
account.ApplyHeaderOverrides(req.Header)
|
||||
setOpenAICodexRoutingHintFromBody(req.Header, account, body)
|
||||
logOpenAIRoutingDiagnosticsFromBody(ctx, account, "http_passthrough", req.Header, body, "not_applicable")
|
||||
|
||||
return req, nil
|
||||
}
|
||||
|
||||
func stripOpenAILegacyResponsesBeta(headers http.Header) {
|
||||
if headers == nil {
|
||||
return
|
||||
}
|
||||
|
||||
preserved := make([]string, 0)
|
||||
for key, values := range headers {
|
||||
if !strings.EqualFold(strings.TrimSpace(key), "OpenAI-Beta") {
|
||||
continue
|
||||
}
|
||||
delete(headers, key)
|
||||
for _, value := range values {
|
||||
parts := strings.Split(value, ",")
|
||||
kept := parts[:0]
|
||||
for _, part := range parts {
|
||||
part = strings.TrimSpace(part)
|
||||
if part == "" || strings.EqualFold(part, "responses=experimental") {
|
||||
continue
|
||||
}
|
||||
kept = append(kept, part)
|
||||
}
|
||||
if len(kept) > 0 {
|
||||
preserved = append(preserved, strings.Join(kept, ", "))
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, value := range preserved {
|
||||
headers.Add("OpenAI-Beta", value)
|
||||
}
|
||||
}
|
||||
|
||||
func shouldFailoverOpenAIPassthroughResponse(account *Account, statusCode int, responseBody []byte) bool {
|
||||
if isOpenAIContextWindowError("", responseBody) {
|
||||
return false
|
||||
@@ -757,8 +793,17 @@ func openAIStreamDataStartsClientOutput(data, eventType string) bool {
|
||||
if trimmed == "" {
|
||||
return false
|
||||
}
|
||||
if strings.TrimSpace(eventType) == "response.failed" {
|
||||
switch strings.TrimSpace(eventType) {
|
||||
case "response.failed":
|
||||
return false
|
||||
case "error":
|
||||
// 上游降载/瞬时故障会先推 {"type":"error"} 帧、再以 response.failed 收尾。
|
||||
// 可重试类错误帧不能算客户端输出:一旦把它当首输出 flush,
|
||||
// clientOutputStarted 即被固化,随后的 failed 事件永远进不了 pre-output
|
||||
// failover 分支,只能把致命错误原样转发给客户端。不可重试类
|
||||
// (content_policy / invalid_request 等)维持原样转发,保留上游错误细节。
|
||||
payload := []byte(trimmed)
|
||||
return !openAIStreamFailedEventShouldFailover(payload, extractOpenAISSEErrorMessage(payload))
|
||||
}
|
||||
return !openAIStreamEventIsPreamble(eventType)
|
||||
}
|
||||
@@ -785,6 +830,41 @@ func isOpenAIUpstreamCapacityShedEvent(payload []byte) bool {
|
||||
}
|
||||
}
|
||||
|
||||
// openAICapacityShedRetryableClientCode 是把上游容量降载错误转发给客户端时改写
|
||||
// 使用的错误码。Codex CLI 按闭集对错误码分类:server_is_overloaded / slow_down
|
||||
// 被判为致命错误(客户端提示 "Selected model is at capacity. Please try a
|
||||
// different model." 并直接终止会话),而 server_error 等致命集之外的错误码会进入
|
||||
// 客户端内置的退避重试。
|
||||
const openAICapacityShedRetryableClientCode = "server_error"
|
||||
|
||||
// sanitizeOpenAICapacityShedErrorCodeForClient 把即将写给下游客户端的
|
||||
// error / response.failed 事件中的容量降载错误码改写为客户端可重试的错误码。
|
||||
// 走到转发这一步说明网关侧 failover 已不可用(流中途)或已用尽;保留原始降载码
|
||||
// 只会让客户端就地终止会话。错误消息原样保留;监控与账号状态判定都基于改写前
|
||||
// 的原始 payload,不受影响。rate_limit 等其他错误码一律不动(客户端依赖
|
||||
// rate_limit_exceeded 原码解析重试延时)。
|
||||
func sanitizeOpenAICapacityShedErrorCodeForClient(payload []byte) ([]byte, bool) {
|
||||
if len(payload) == 0 || !gjson.ValidBytes(payload) || !isOpenAIUpstreamCapacityShedEvent(payload) {
|
||||
return payload, false
|
||||
}
|
||||
updated := payload
|
||||
changed := false
|
||||
for _, path := range []string{"response.error.code", "error.code"} {
|
||||
switch strings.ToLower(strings.TrimSpace(gjson.GetBytes(updated, path).String())) {
|
||||
case "server_is_overloaded", "slow_down":
|
||||
default:
|
||||
continue
|
||||
}
|
||||
next, err := sjson.SetBytes(updated, path, openAICapacityShedRetryableClientCode)
|
||||
if err != nil {
|
||||
return payload, false
|
||||
}
|
||||
updated = next
|
||||
changed = true
|
||||
}
|
||||
return updated, changed
|
||||
}
|
||||
|
||||
func openAIStreamFailedEventSemanticStatus(payload []byte, message string) int {
|
||||
if isOpenAIContextWindowError(message, payload) {
|
||||
return http.StatusBadRequest
|
||||
@@ -1051,6 +1131,10 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough(
|
||||
originalModel string,
|
||||
mappedModel string,
|
||||
) (*openaiStreamingResultPassthrough, error) {
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
writeOpenAIPassthroughResponseHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter)
|
||||
|
||||
// SSE headers
|
||||
@@ -1131,6 +1215,8 @@ func (s *OpenAIGatewayService) handleStreamingResponsePassthrough(
|
||||
if data, ok := extractOpenAISSEDataLine(line); ok {
|
||||
dataBytes := []byte(data)
|
||||
trimmedData := strings.TrimSpace(data)
|
||||
rawEventType := strings.TrimSpace(gjson.GetBytes(dataBytes, "type").String())
|
||||
observer.ObserveOpenAI(dataBytes, rawEventType)
|
||||
if needModelReplace && strings.Contains(data, mappedModel) {
|
||||
line = s.replaceModelInSSELine(line, mappedModel, originalModel)
|
||||
if replacedData, replaced := extractOpenAISSEDataLine(line); replaced {
|
||||
@@ -1317,6 +1403,15 @@ func (s *OpenAIGatewayService) handleNonStreamingResponsePassthrough(
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
if bodyHasSSEFraming(body) {
|
||||
observeOpenAISSEBody(observer, string(body))
|
||||
} else {
|
||||
observer.ObserveOpenAI(body, strings.TrimSpace(gjson.GetBytes(body, "type").String()))
|
||||
}
|
||||
|
||||
// Detect SSE responses from upstream and convert to JSON.
|
||||
// Some upstreams (e.g. other sub2api instances) may return SSE even when
|
||||
|
||||
@@ -308,6 +308,7 @@ func normalizeOpenAICompactRequestBody(body []byte) ([]byte, bool, error) {
|
||||
"tools",
|
||||
"parallel_tool_calls",
|
||||
"reasoning",
|
||||
"service_tier",
|
||||
"text",
|
||||
"previous_response_id",
|
||||
} {
|
||||
|
||||
@@ -387,13 +387,30 @@ func TestOpenAIResponseFlush_FailedAndErrorEventsFlushAtBoundaries(t *testing.T)
|
||||
require.Contains(t, flushes[1], "response.failed")
|
||||
})
|
||||
|
||||
t.Run("error event", func(t *testing.T) {
|
||||
t.Run("retryable error event buffered until terminal", func(t *testing.T) {
|
||||
// 可重试类 error 帧不算客户端输出:保持在 attempt 缓冲中不单独 flush,
|
||||
// 为随后可能到达的 response.failed 保留 pre-output failover 能力,
|
||||
// 与终止帧一起出站。
|
||||
body := "data: {\"type\":\"error\",\"error\":{\"message\":\"failed\"}}\n\n" +
|
||||
"data: [DONE]\n\n"
|
||||
recorder := newOpenAIResponseFlushRecorder()
|
||||
|
||||
result, err := runOpenAIResponseFlushTest(recorder, io.NopCloser(strings.NewReader(body)), config.GatewayConfig{})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
gotBody, flushes := recorder.snapshot()
|
||||
require.Equal(t, body, gotBody)
|
||||
require.Len(t, flushes, 1)
|
||||
})
|
||||
|
||||
t.Run("non-retryable error event flushes at boundary", func(t *testing.T) {
|
||||
body := "data: {\"type\":\"error\",\"error\":{\"code\":\"invalid_request\",\"message\":\"bad request\"}}\n\n" +
|
||||
"data: [DONE]\n\n"
|
||||
recorder := newOpenAIResponseFlushRecorder()
|
||||
|
||||
result, err := runOpenAIResponseFlushTest(recorder, io.NopCloser(strings.NewReader(body)), config.GatewayConfig{})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
gotBody, flushes := recorder.snapshot()
|
||||
|
||||
@@ -45,6 +45,10 @@ func (s *OpenAIGatewayService) handleStreamingResponse(ctx context.Context, resp
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context.Context, resp *http.Response, c *gin.Context, account *Account, startTime time.Time, originalModel, mappedModel, reasoningEffort string) (*openaiStreamingResult, error) {
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
firstOutputTimeout := time.Duration(0)
|
||||
if account != nil && account.Platform == PlatformOpenAI {
|
||||
firstOutputTimeout = s.openAIFirstOutputTimeout(reasoningEffort)
|
||||
@@ -423,6 +427,7 @@ func (s *OpenAIGatewayService) handleStreamingResponseWithReasoning(ctx context.
|
||||
dataBytes := []byte(data)
|
||||
eventTypeRaw := gjson.GetBytes(dataBytes, "type").String()
|
||||
eventType := strings.TrimSpace(eventTypeRaw)
|
||||
observer.ObserveOpenAI(dataBytes, eventTypeRaw)
|
||||
// 初始上游 data 的 type 只解析一次:原始值保持终止事件的精确匹配,规范化值供后续分支复用。
|
||||
if openAIStreamEventIsTerminalWithType(data, eventTypeRaw) {
|
||||
sawTerminalEvent = true
|
||||
@@ -1138,6 +1143,15 @@ func (s *OpenAIGatewayService) handleNonStreamingResponse(ctx context.Context, r
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
observer := upstreamResponseModelObserverFromContext(c)
|
||||
if observer == nil {
|
||||
observer = beginUpstreamResponseModelObservation(c)
|
||||
}
|
||||
if bodyHasSSEFraming(body) {
|
||||
observeOpenAISSEBody(observer, string(body))
|
||||
} else {
|
||||
observer.ObserveOpenAI(body, strings.TrimSpace(gjson.GetBytes(body, "type").String()))
|
||||
}
|
||||
|
||||
// Detect SSE responses for ALL account types via Content-Type header.
|
||||
// Some OpenAI-compatible upstreams (including other sub2api instances)
|
||||
@@ -1340,10 +1354,21 @@ func extractOpenAISSEErrorMessage(payload []byte) string {
|
||||
}
|
||||
|
||||
func sanitizeOpenAIResponseFailedEventForClient(payload []byte, eventType string, clientOutputStarted bool) ([]byte, bool) {
|
||||
if eventType != "response.failed" || len(payload) == 0 || !gjson.ValidBytes(payload) {
|
||||
eventType = strings.TrimSpace(eventType)
|
||||
isFailedEvent := eventType == "response.failed"
|
||||
if (!isFailedEvent && eventType != "error") || len(payload) == 0 || !gjson.ValidBytes(payload) {
|
||||
return payload, false
|
||||
}
|
||||
updated := payload
|
||||
// 容量降载码对 Codex CLI 是致命错误;事件既然要写给客户端(failover 已不可用),
|
||||
// 就改写为客户端可重试的错误码。error 帧与 response.failed 都要改:上游降载
|
||||
// 总是先推 error 帧再收 failed,两帧携带同一个错误。
|
||||
if rewritten, changed := sanitizeOpenAICapacityShedErrorCodeForClient(updated); changed {
|
||||
updated = rewritten
|
||||
}
|
||||
if !isFailedEvent {
|
||||
return updated, !bytes.Equal(updated, payload)
|
||||
}
|
||||
if clientOutputStarted && isOpenAIContextWindowError(extractOpenAISSEErrorMessage(payload), payload) {
|
||||
errorPath := ""
|
||||
switch {
|
||||
|
||||
@@ -243,6 +243,10 @@ type OpenAIForwardResult struct {
|
||||
// UpstreamModel is the actual model sent to the upstream provider after mapping.
|
||||
// Empty when no mapping was applied (requested model was used as-is).
|
||||
UpstreamModel string
|
||||
// UpstreamResponseModel is captured from the raw successful upstream
|
||||
// response before any client-facing rewrite or protocol conversion.
|
||||
UpstreamResponseModel string
|
||||
UpstreamResponseModelConflict bool
|
||||
// UpstreamEndpoint is the actual upstream API path used for this request.
|
||||
// It avoids guessing when one downstream protocol can use multiple upstream endpoints.
|
||||
UpstreamEndpoint string
|
||||
|
||||
@@ -2925,10 +2925,30 @@ func TestOpenAIBuildUpstreamRequestOpenAIPassthroughPreservesCompactPath(t *test
|
||||
require.Equal(t, chatgptCodexURL+"/compact", req.URL.String())
|
||||
require.Equal(t, "application/json", req.Header.Get("Accept"))
|
||||
require.Equal(t, codexCLIVersion, req.Header.Get("Version"))
|
||||
require.Empty(t, req.Header.Get("OpenAI-Beta"), "Codex OAuth HTTP must not synthesize the legacy responses beta header")
|
||||
require.NotEmpty(t, req.Header.Get("Session_Id"))
|
||||
require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(req.Context()))
|
||||
}
|
||||
|
||||
func TestOpenAIBuildUpstreamRequestOpenAIPassthroughPreservesExplicitAPIKeyBetaHeader(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader([]byte(`{"model":"gpt-5"}`)))
|
||||
c.Request.Header.Set("OpenAI-Beta", "api-key-specific-beta")
|
||||
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{
|
||||
Security: config.SecurityConfig{
|
||||
URLAllowlist: config.URLAllowlistConfig{Enabled: false},
|
||||
},
|
||||
}}
|
||||
account := &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
|
||||
|
||||
req, err := svc.buildUpstreamRequestOpenAIPassthrough(c.Request.Context(), c, account, []byte(`{"model":"gpt-5"}`), "token")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "api-key-specific-beta", req.Header.Get("OpenAI-Beta"), "OAuth-only backport must not alter API-key passthrough headers")
|
||||
}
|
||||
|
||||
func TestOpenAIBuildUpstreamRequestCompactForcesJSONAcceptForOAuth(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
@@ -2946,6 +2966,7 @@ func TestOpenAIBuildUpstreamRequestCompactForcesJSONAcceptForOAuth(t *testing.T)
|
||||
require.Equal(t, chatgptCodexURL+"/compact", req.URL.String())
|
||||
require.Equal(t, "application/json", req.Header.Get("Accept"))
|
||||
require.Equal(t, codexCLIVersion, req.Header.Get("Version"))
|
||||
require.Empty(t, req.Header.Get("OpenAI-Beta"), "Codex OAuth HTTP must not synthesize the legacy responses beta header")
|
||||
require.NotEmpty(t, req.Header.Get("Session_Id"))
|
||||
require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(req.Context()))
|
||||
}
|
||||
@@ -3493,6 +3514,32 @@ func TestHandleNonStreamingResponse_OAuthJSONBodyWithDataEventTextKeepsJSONUsage
|
||||
require.Contains(t, rec.Body.String(), "processing data: 1,2,3 then event: click finished")
|
||||
}
|
||||
|
||||
func TestHandleNonStreamingResponse_ObservesUpstreamModelBeforeClientRewrite(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{}}
|
||||
resp := &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(
|
||||
`{"id":"resp_model_audit","object":"response","model":"gpt-5.5","status":"completed","output":[],"usage":{"input_tokens":3,"output_tokens":2,"total_tokens":5}}`,
|
||||
)),
|
||||
}
|
||||
account := &Account{ID: 1, Type: AccountTypeAPIKey}
|
||||
|
||||
result, err := svc.handleNonStreamingResponse(context.Background(), resp, c, account, "gpt-5.6-sol", "gpt-5.5")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
|
||||
// 客户端仍看到自己请求的模型名,审计观察器则保留改写前的上游声明。
|
||||
require.Equal(t, "gpt-5.6-sol", gjson.Get(rec.Body.String(), "model").String())
|
||||
require.Equal(t, "gpt-5.5", observedUpstreamResponseModel(c))
|
||||
require.False(t, observedUpstreamResponseModelConflict(c))
|
||||
}
|
||||
|
||||
func TestHandleSSEToJSON_ReconstructsImageGenerationOutputItemDone(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
|
||||
@@ -211,7 +211,7 @@ func isOpenAIContextWindowError(upstreamMsg string, upstreamBody []byte) bool {
|
||||
|
||||
func (s *OpenAIGatewayService) shouldFailoverUpstreamError(statusCode int) bool {
|
||||
switch statusCode {
|
||||
case 401, 402, 403, 429, 529:
|
||||
case 401, 402, 403, 405, 429, 529:
|
||||
return true
|
||||
default:
|
||||
return statusCode >= 500
|
||||
|
||||
@@ -275,31 +275,43 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec
|
||||
if input.OriginalModel != "" {
|
||||
requestedModel = input.OriginalModel
|
||||
}
|
||||
sentModel := upstreamSentModel(result.Model, result.UpstreamModel)
|
||||
if result.UpstreamResponseModelConflict {
|
||||
logger.L().Warn("upstream_response_model_conflict",
|
||||
zap.String("platform", account.Platform),
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.String("request_id", requestID),
|
||||
zap.String("sent_model", sentModel),
|
||||
zap.String("selected_response_model", strings.TrimSpace(result.UpstreamResponseModel)),
|
||||
)
|
||||
}
|
||||
|
||||
usageLog := &UsageLog{
|
||||
UserID: user.ID,
|
||||
APIKeyID: apiKey.ID,
|
||||
AccountID: account.ID,
|
||||
RequestID: requestID,
|
||||
Model: result.Model,
|
||||
RequestedModel: requestedModel,
|
||||
UpstreamModel: optionalTrimmedStringPtr(result.UpstreamModel),
|
||||
ServiceTier: result.ServiceTier,
|
||||
ReasoningEffort: result.ReasoningEffort,
|
||||
InboundEndpoint: optionalTrimmedStringPtr(input.InboundEndpoint),
|
||||
UpstreamEndpoint: optionalTrimmedStringPtr(input.UpstreamEndpoint),
|
||||
InputTokens: actualInputTokens,
|
||||
OutputTokens: result.Usage.OutputTokens,
|
||||
CacheCreationTokens: result.Usage.CacheCreationInputTokens,
|
||||
CacheReadTokens: result.Usage.CacheReadInputTokens,
|
||||
ImageInputTokens: result.Usage.ImageInputTokens,
|
||||
ImageOutputTokens: result.Usage.ImageOutputTokens,
|
||||
ImageCount: result.ImageCount,
|
||||
ImageSize: optionalTrimmedStringPtr(result.ImageSize),
|
||||
ImageInputSize: optionalTrimmedStringPtr(result.ImageInputSize),
|
||||
ImageOutputSize: optionalTrimmedStringPtr(result.ImageOutputSize),
|
||||
ImageSizeSource: optionalTrimmedStringPtr(result.ImageSizeSource),
|
||||
ImageSizeBreakdown: result.ImageSizeBreakdown,
|
||||
UserID: user.ID,
|
||||
APIKeyID: apiKey.ID,
|
||||
AccountID: account.ID,
|
||||
RequestID: requestID,
|
||||
Model: result.Model,
|
||||
RequestedModel: requestedModel,
|
||||
UpstreamModel: optionalTrimmedStringPtr(result.UpstreamModel),
|
||||
UpstreamResponseModel: optionalTrimmedStringPtr(result.UpstreamResponseModel),
|
||||
UpstreamModelMismatch: upstreamModelMismatch(sentModel, result.UpstreamResponseModel),
|
||||
ServiceTier: result.ServiceTier,
|
||||
ReasoningEffort: result.ReasoningEffort,
|
||||
InboundEndpoint: optionalTrimmedStringPtr(input.InboundEndpoint),
|
||||
UpstreamEndpoint: optionalTrimmedStringPtr(input.UpstreamEndpoint),
|
||||
InputTokens: actualInputTokens,
|
||||
OutputTokens: result.Usage.OutputTokens,
|
||||
CacheCreationTokens: result.Usage.CacheCreationInputTokens,
|
||||
CacheReadTokens: result.Usage.CacheReadInputTokens,
|
||||
ImageInputTokens: result.Usage.ImageInputTokens,
|
||||
ImageOutputTokens: result.Usage.ImageOutputTokens,
|
||||
ImageCount: result.ImageCount,
|
||||
ImageSize: optionalTrimmedStringPtr(result.ImageSize),
|
||||
ImageInputSize: optionalTrimmedStringPtr(result.ImageInputSize),
|
||||
ImageOutputSize: optionalTrimmedStringPtr(result.ImageOutputSize),
|
||||
ImageSizeSource: optionalTrimmedStringPtr(result.ImageSizeSource),
|
||||
ImageSizeBreakdown: result.ImageSizeBreakdown,
|
||||
}
|
||||
isVideoUsage := isGrokVideoUsageResult(result, billingModels)
|
||||
if isVideoUsage {
|
||||
|
||||
@@ -767,7 +767,7 @@ func TestOpenAIGatewayServiceForwardImages_OAuthPassesNAndReturnsAllImages(t *te
|
||||
require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type"))
|
||||
require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept"))
|
||||
require.Equal(t, "acct-123", upstream.lastReq.Header.Get("chatgpt-account-id"))
|
||||
require.Equal(t, "responses=experimental", upstream.lastReq.Header.Get("OpenAI-Beta"))
|
||||
require.Empty(t, upstream.lastReq.Header.Get("OpenAI-Beta"))
|
||||
|
||||
require.Equal(t, openAIImagesResponsesMainModel, gjson.GetBytes(upstream.lastBody, "model").String())
|
||||
require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
|
||||
|
||||
@@ -0,0 +1,127 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"sort"
|
||||
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
const (
|
||||
// 工具定义在多轮历史里最多再嵌套一层 tools,留出余量后截断,避免畸形请求体
|
||||
// 造成无界递归。
|
||||
openAIResponsesToolSchemaMaxDepth = 4
|
||||
// JSON Schema 里 type 只能是字符串或字符串数组;显式 null 无论哪个方言都非法,
|
||||
// 补成 object 与 upstream 对该工具的实际期望一致。
|
||||
openAIResponsesToolSchemaFallbackType = `"object"`
|
||||
// 显式 null 在 JSON 里只有这一种字面量形态。
|
||||
openAIResponsesToolSchemaNullLiteral = "null"
|
||||
)
|
||||
|
||||
// openAIResponsesToolSchemaNullType 记录一处待修正的 null,用原始 body 上的
|
||||
// 绝对字节偏移表示,便于最后一次性拼接。
|
||||
type openAIResponsesToolSchemaNullType struct {
|
||||
offset int
|
||||
length int
|
||||
}
|
||||
|
||||
// sanitizeOpenAIResponsesToolParameterTypes 修正请求体中显式为 null 的
|
||||
// tools[].parameters.type。
|
||||
//
|
||||
// Codex Desktop 内置的 automation_update 工具会带 parameters.type = null,
|
||||
// OpenAI 直接回 400 invalid_function_parameters,而网关把该状态归一成可重试的
|
||||
// 502 upstream_error;该工具定义又会沉进多轮历史,导致之后每一轮继续失败并在
|
||||
// 账号池里反复重放同一份坏 Schema。
|
||||
//
|
||||
// 只修正显式 null:缺失 type 的 Schema 本身合法(等价于不约束),补写会收窄
|
||||
// 客户端语义,因此保持原样。
|
||||
//
|
||||
// 实现上先收集全部命中的绝对偏移,再一次性拼出新 body:逐个 sjson.SetBytes 每次
|
||||
// 都会重扫并全量拷贝整个文档,命中 N 处就是 N 次全量拷贝,而 /v1/responses 的
|
||||
// body 上限是 gateway.max_body_size(默认 256MB),构造请求能塞进百万级命中。
|
||||
func sanitizeOpenAIResponsesToolParameterTypes(body []byte) ([]byte, bool, error) {
|
||||
if len(body) == 0 {
|
||||
return body, false, nil
|
||||
}
|
||||
|
||||
hits := make([]openAIResponsesToolSchemaNullType, 0, 2)
|
||||
collectOpenAIResponsesToolSchemaNullTypes(body, gjson.GetBytes(body, "tools"), 0, &hits)
|
||||
if input := gjson.GetBytes(body, "input"); input.IsArray() {
|
||||
input.ForEach(func(_, item gjson.Result) bool {
|
||||
if item.IsObject() {
|
||||
collectOpenAIResponsesToolSchemaNullTypes(body, item.Get("tools"), 0, &hits)
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
if len(hits) == 0 {
|
||||
return body, false, nil
|
||||
}
|
||||
|
||||
// tools 与 input 在 body 里的先后顺序由客户端决定,收集顺序不保证单调。
|
||||
sort.Slice(hits, func(i, j int) bool { return hits[i].offset < hits[j].offset })
|
||||
|
||||
sanitized := make([]byte, 0, len(body)+len(hits)*len(openAIResponsesToolSchemaFallbackType))
|
||||
cursor := 0
|
||||
for _, hit := range hits {
|
||||
// 收集阶段已逐个校验过区间,这里再挡一次重叠,保证拼接严格单调向前。
|
||||
if hit.offset < cursor {
|
||||
continue
|
||||
}
|
||||
sanitized = append(sanitized, body[cursor:hit.offset]...)
|
||||
sanitized = append(sanitized, openAIResponsesToolSchemaFallbackType...)
|
||||
cursor = hit.offset + hit.length
|
||||
}
|
||||
sanitized = append(sanitized, body[cursor:]...)
|
||||
return sanitized, true, nil
|
||||
}
|
||||
|
||||
// collectOpenAIResponsesToolSchemaNullTypes 收集一个 tools 数组里所有需要修正的
|
||||
// parameters.type 位置。不按 tool type 过滤:null 的 schema type 在 function、
|
||||
// custom 以及任何 hosted 工具上都同样非法。
|
||||
func collectOpenAIResponsesToolSchemaNullTypes(
|
||||
body []byte, tools gjson.Result, depth int, hits *[]openAIResponsesToolSchemaNullType,
|
||||
) {
|
||||
if depth > openAIResponsesToolSchemaMaxDepth || !tools.IsArray() {
|
||||
return
|
||||
}
|
||||
tools.ForEach(func(_, tool gjson.Result) bool {
|
||||
if !tool.IsObject() {
|
||||
return true
|
||||
}
|
||||
// Responses 形态用顶层 parameters,ChatCompletions 形态用 function.parameters,
|
||||
// 两种都可能出现在 Responses 请求里(见 normalizeCodexTools)。
|
||||
for _, suffix := range []string{"parameters", "function.parameters"} {
|
||||
params := tool.Get(suffix)
|
||||
if !params.IsObject() {
|
||||
continue
|
||||
}
|
||||
// gjson 用 Type==Null 同时表示「显式 null」和「路径不存在」,靠 Raw
|
||||
// 区分:不存在时 Raw 为空串。
|
||||
if typ := params.Get("type"); typ.Type == gjson.Null && typ.Raw == openAIResponsesToolSchemaNullLiteral {
|
||||
appendOpenAIResponsesToolSchemaNullType(body, typ, hits)
|
||||
}
|
||||
}
|
||||
// 历史输入里的工具定义会再嵌套一层 tools(upstream 报错路径形如
|
||||
// input[234].tools[0].tools[3].parameters)。
|
||||
collectOpenAIResponsesToolSchemaNullTypes(body, tool.Get("tools"), depth+1, hits)
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
// appendOpenAIResponsesToolSchemaNullType 先校验 gjson 给出的偏移确实指向原始
|
||||
// body 上那段 null,再记录。gjson 对嵌套取值同样返回相对原始文档的绝对偏移,但
|
||||
// Index 为 0 表示未知;偏移不可用时跳过该处而不是猜位置——少修一个工具只是维持
|
||||
// 现状,拼错位置会损坏整个请求体。
|
||||
func appendOpenAIResponsesToolSchemaNullType(
|
||||
body []byte, typ gjson.Result, hits *[]openAIResponsesToolSchemaNullType,
|
||||
) {
|
||||
end := typ.Index + len(typ.Raw)
|
||||
if typ.Index <= 0 || end > len(body) {
|
||||
return
|
||||
}
|
||||
if !bytes.Equal(body[typ.Index:end], []byte(typ.Raw)) {
|
||||
return
|
||||
}
|
||||
*hits = append(*hits, openAIResponsesToolSchemaNullType{offset: typ.Index, length: len(typ.Raw)})
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// issue #5364 的最小复现体:Codex Desktop 内置 automation_update 带
|
||||
// parameters.type = null,upstream 回 400 invalid_function_parameters。
|
||||
func TestSanitizeOpenAIResponsesToolParameterTypes_TopLevelFunctionTool(t *testing.T) {
|
||||
body := []byte(`{
|
||||
"model": "gpt-5.6-sol",
|
||||
"input": "Reply with OK.",
|
||||
"stream": false,
|
||||
"tools": [
|
||||
{
|
||||
"type": "function",
|
||||
"name": "automation_update",
|
||||
"description": "Update an automation.",
|
||||
"parameters": {"type": null, "properties": {}}
|
||||
}
|
||||
]
|
||||
}`)
|
||||
|
||||
sanitized, changed, err := sanitizeOpenAIResponsesToolParameterTypes(body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, changed)
|
||||
require.Equal(t, "object", gjson.GetBytes(sanitized, "tools.0.parameters.type").String())
|
||||
// 只改 type,工具其余定义原样保留。
|
||||
require.Equal(t, "automation_update", gjson.GetBytes(sanitized, "tools.0.name").String())
|
||||
require.Equal(t, "Update an automation.", gjson.GetBytes(sanitized, "tools.0.description").String())
|
||||
require.True(t, gjson.GetBytes(sanitized, "tools.0.parameters.properties").IsObject())
|
||||
// 请求体其余字段不受影响。
|
||||
require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(sanitized, "model").String())
|
||||
require.Equal(t, "Reply with OK.", gjson.GetBytes(sanitized, "input").String())
|
||||
}
|
||||
|
||||
// 合法 Schema 必须原样返回:changed=false 且字节不变,避免无谓重写打散
|
||||
// prompt cache 前缀。
|
||||
func TestSanitizeOpenAIResponsesToolParameterTypes_ValidSchemaUntouched(t *testing.T) {
|
||||
body := []byte(`{"tools":[{"type":"function","name":"ok","parameters":{"type":"object","properties":{}}}]}`)
|
||||
|
||||
sanitized, changed, err := sanitizeOpenAIResponsesToolParameterTypes(body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.False(t, changed)
|
||||
require.Equal(t, string(body), string(sanitized))
|
||||
}
|
||||
|
||||
// 缺失 type 的 Schema 本身合法(等价于不约束),不得补写——补写会收窄客户端语义。
|
||||
func TestSanitizeOpenAIResponsesToolParameterTypes_MissingTypeNotInvented(t *testing.T) {
|
||||
body := []byte(`{"tools":[{"type":"function","name":"ok","parameters":{"properties":{}}}]}`)
|
||||
|
||||
sanitized, changed, err := sanitizeOpenAIResponsesToolParameterTypes(body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.False(t, changed)
|
||||
require.False(t, gjson.GetBytes(sanitized, "tools.0.parameters.type").Exists())
|
||||
}
|
||||
|
||||
// 多轮历史:工具定义沉进 input 后,upstream 报错路径形如
|
||||
// input[N].tools[i].tools[j].parameters,两层都要修。
|
||||
func TestSanitizeOpenAIResponsesToolParameterTypes_NestedHistoryTools(t *testing.T) {
|
||||
body := []byte(`{
|
||||
"input": [
|
||||
{"type": "message", "role": "user", "content": "hi"},
|
||||
{
|
||||
"type": "additional_tools",
|
||||
"role": "developer",
|
||||
"tools": [
|
||||
{
|
||||
"type": "namespace",
|
||||
"name": "codex_app",
|
||||
"tools": [
|
||||
{"type": "function", "name": "noop", "parameters": {"type": "object"}},
|
||||
{"type": "function", "name": "automation_update", "parameters": {"type": null}}
|
||||
]
|
||||
},
|
||||
{"type": "function", "name": "outer", "parameters": {"type": null}}
|
||||
]
|
||||
}
|
||||
]
|
||||
}`)
|
||||
|
||||
sanitized, changed, err := sanitizeOpenAIResponsesToolParameterTypes(body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, changed)
|
||||
require.Equal(t, "object", gjson.GetBytes(sanitized, "input.1.tools.0.tools.1.parameters.type").String())
|
||||
require.Equal(t, "object", gjson.GetBytes(sanitized, "input.1.tools.1.parameters.type").String())
|
||||
// 原本合法的兄弟条目保持不变。
|
||||
require.Equal(t, "object", gjson.GetBytes(sanitized, "input.1.tools.0.tools.0.parameters.type").String())
|
||||
require.Equal(t, "hi", gjson.GetBytes(sanitized, "input.0.content").String())
|
||||
}
|
||||
|
||||
// ChatCompletions 形态的工具({type:"function", function:{...}})同样可能出现在
|
||||
// Responses 请求里,见 normalizeCodexTools。
|
||||
func TestSanitizeOpenAIResponsesToolParameterTypes_ChatCompletionsShape(t *testing.T) {
|
||||
body := []byte(`{"tools":[{"type":"function","function":{"name":"legacy","parameters":{"type":null}}}]}`)
|
||||
|
||||
sanitized, changed, err := sanitizeOpenAIResponsesToolParameterTypes(body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, changed)
|
||||
require.Equal(t, "object", gjson.GetBytes(sanitized, "tools.0.function.parameters.type").String())
|
||||
require.Equal(t, "legacy", gjson.GetBytes(sanitized, "tools.0.function.name").String())
|
||||
}
|
||||
|
||||
// 索引映射:只有坏条目被改,前后兄弟条目按原下标保持不变。
|
||||
func TestSanitizeOpenAIResponsesToolParameterTypes_OnlyOffendingIndexRewritten(t *testing.T) {
|
||||
body := []byte(`{"tools":[
|
||||
{"type":"function","name":"a","parameters":{"type":"object"}},
|
||||
{"type":"function","name":"b","parameters":{"type":null}},
|
||||
{"type":"function","name":"c","parameters":{"type":"object"}}
|
||||
]}`)
|
||||
|
||||
sanitized, changed, err := sanitizeOpenAIResponsesToolParameterTypes(body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, changed)
|
||||
require.Equal(t, "a", gjson.GetBytes(sanitized, "tools.0.name").String())
|
||||
require.Equal(t, "b", gjson.GetBytes(sanitized, "tools.1.name").String())
|
||||
require.Equal(t, "c", gjson.GetBytes(sanitized, "tools.2.name").String())
|
||||
require.Equal(t, "object", gjson.GetBytes(sanitized, "tools.1.parameters.type").String())
|
||||
require.Equal(t, 3, int(gjson.GetBytes(sanitized, "tools.#").Int()))
|
||||
}
|
||||
|
||||
// 畸形/非常规形态不得 panic,且一律按不变处理。
|
||||
func TestSanitizeOpenAIResponsesToolParameterTypes_MalformedShapesAreNoOps(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
body string
|
||||
}{
|
||||
{"empty_body", ``},
|
||||
{"no_tools", `{"model":"gpt-5.6-sol","input":"hi"}`},
|
||||
{"tools_null", `{"tools":null}`},
|
||||
{"tools_object", `{"tools":{"type":"function"}}`},
|
||||
{"tool_is_string", `{"tools":["freeform"]}`},
|
||||
{"parameters_is_string", `{"tools":[{"type":"function","parameters":"nope"}]}`},
|
||||
{"parameters_null", `{"tools":[{"type":"function","parameters":null}]}`},
|
||||
{"input_string", `{"input":"hi","tools":[]}`},
|
||||
{"input_item_not_object", `{"input":["hi"]}`},
|
||||
{"type_already_array", `{"tools":[{"type":"function","parameters":{"type":["object","null"]}}]}`},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
sanitized, changed, err := sanitizeOpenAIResponsesToolParameterTypes([]byte(tc.body))
|
||||
require.NoError(t, err)
|
||||
require.False(t, changed)
|
||||
require.Equal(t, tc.body, string(sanitized))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 递归深度守卫:超深嵌套只做截断,不递归到栈溢出,也不报错。
|
||||
func TestSanitizeOpenAIResponsesToolParameterTypes_DepthGuard(t *testing.T) {
|
||||
tool := map[string]any{"type": "function", "name": "deep", "parameters": map[string]any{"type": nil}}
|
||||
for i := 0; i < 12; i++ {
|
||||
tool = map[string]any{"type": "namespace", "tools": []any{tool}}
|
||||
}
|
||||
body, err := json.Marshal(map[string]any{"tools": []any{tool}})
|
||||
require.NoError(t, err)
|
||||
|
||||
require.NotPanics(t, func() {
|
||||
_, _, sanitizeErr := sanitizeOpenAIResponsesToolParameterTypes(body)
|
||||
require.NoError(t, sanitizeErr)
|
||||
})
|
||||
}
|
||||
|
||||
// 输出必须是合法 JSON,且除目标字段外与输入等价。
|
||||
func TestSanitizeOpenAIResponsesToolParameterTypes_OutputStaysValidJSON(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-5.5","tool_choice":"none","store":false,"tools":[{"type":"function","name":"automation_update","parameters":{"type":null,"properties":{}}}]}`)
|
||||
|
||||
sanitized, changed, err := sanitizeOpenAIResponsesToolParameterTypes(body)
|
||||
require.NoError(t, err)
|
||||
require.True(t, changed)
|
||||
|
||||
var decoded map[string]any
|
||||
require.NoError(t, json.Unmarshal(sanitized, &decoded))
|
||||
require.Equal(t, "gpt-5.5", decoded["model"])
|
||||
require.Equal(t, "none", decoded["tool_choice"])
|
||||
require.Equal(t, false, decoded["store"])
|
||||
}
|
||||
|
||||
// 输入 body 是调用方持有的缓冲区(Forward 里 canonicalImageIntentBody 与它同源),
|
||||
// 净化必须返回新切片,绝不能就地改写。
|
||||
func TestSanitizeOpenAIResponsesToolParameterTypes_DoesNotMutateInputBody(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-5.6-sol","tools":[{"type":"function","name":"a","parameters":{"type":null}}]}`)
|
||||
original := append([]byte(nil), body...)
|
||||
|
||||
sanitized, changed, err := sanitizeOpenAIResponsesToolParameterTypes(body)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, changed)
|
||||
require.Equal(t, string(original), string(body), "调用方的 body 不得被就地改写")
|
||||
require.NotEqual(t, string(original), string(sanitized))
|
||||
}
|
||||
|
||||
func buildToolSchemaNullTypeBody(t *testing.T, hits int) []byte {
|
||||
t.Helper()
|
||||
tools := make([]any, 0, hits)
|
||||
for i := 0; i < hits; i++ {
|
||||
tools = append(tools, map[string]any{
|
||||
"type": "function",
|
||||
"name": "automation_update",
|
||||
"parameters": map[string]any{"type": nil, "properties": map[string]any{}},
|
||||
})
|
||||
}
|
||||
body, err := json.Marshal(map[string]any{"model": "gpt-5.6-sol", "tools": tools})
|
||||
require.NoError(t, err)
|
||||
return body
|
||||
}
|
||||
|
||||
// 复杂度守卫:重写次数必须与命中数无关。
|
||||
//
|
||||
// 逐个 sjson.SetBytes 的写法每命中一处就重扫并全量拷贝一次文档,命中 N 处即 N 次
|
||||
// 全量拷贝;/v1/responses 的 body 上限是 gateway.max_body_size(默认 256MB),
|
||||
// 构造请求可以塞进百万级命中,会被放大成 TB 级 memcpy。这里用分配次数锁死该行为:
|
||||
// 命中数放大 500 倍,分配次数不得随之增长。
|
||||
func TestSanitizeOpenAIResponsesToolParameterTypes_RewriteCountIndependentOfHits(t *testing.T) {
|
||||
small := buildToolSchemaNullTypeBody(t, 4)
|
||||
large := buildToolSchemaNullTypeBody(t, 2000)
|
||||
|
||||
smallAllocs := testing.AllocsPerRun(2, func() {
|
||||
_, _, _ = sanitizeOpenAIResponsesToolParameterTypes(small)
|
||||
})
|
||||
largeAllocs := testing.AllocsPerRun(2, func() {
|
||||
_, _, _ = sanitizeOpenAIResponsesToolParameterTypes(large)
|
||||
})
|
||||
|
||||
// 命中切片扩容是对数级,留出充裕余量;线性写法在这里会是 2000 量级。
|
||||
require.Less(t, largeAllocs, smallAllocs+40,
|
||||
"分配次数随命中数线性增长,说明退回了逐路径全量重写 (small=%v large=%v)", smallAllocs, largeAllocs)
|
||||
|
||||
// 同时确认大 body 的结果确实全部修好了。
|
||||
sanitized, changed, err := sanitizeOpenAIResponsesToolParameterTypes(large)
|
||||
require.NoError(t, err)
|
||||
require.True(t, changed)
|
||||
require.Equal(t, 2000, int(gjson.GetBytes(sanitized, "tools.#").Int()))
|
||||
gjson.GetBytes(sanitized, "tools").ForEach(func(_, tool gjson.Result) bool {
|
||||
require.Equal(t, "object", tool.Get("parameters.type").String())
|
||||
return true
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,128 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/tidwall/gjson"
|
||||
"go.uber.org/zap"
|
||||
"golang.org/x/net/http/httpguts"
|
||||
)
|
||||
|
||||
const openAICodexRoutingHintHeader = "x-codex-routing-hint"
|
||||
|
||||
// setOpenAICodexRoutingHint mirrors the Codex backend routing-hint contract for
|
||||
// OpenAI OAuth requests. The request model must already be the final upstream
|
||||
// slug and serviceTier must already reflect any local policy rewrite/filter.
|
||||
func setOpenAICodexRoutingHint(headers http.Header, account *Account, model string, serviceTier string) {
|
||||
if headers == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// The routing hint is gateway-owned. Strip every spelling before deciding
|
||||
// whether to synthesize it so API-key/provider credential paths cannot pass
|
||||
// through a caller- or account-override-supplied hint. http.Header.Del only
|
||||
// removes the canonical map key; inbound maps can contain raw lowercase keys.
|
||||
deleteOpenAIHeaderEqualFold(headers, openAICodexRoutingHintHeader)
|
||||
if account == nil || !account.IsOpenAIOAuth() {
|
||||
return
|
||||
}
|
||||
|
||||
model = strings.TrimSpace(model)
|
||||
if model == "" || strings.ContainsAny(model, ";=") {
|
||||
return
|
||||
}
|
||||
|
||||
// Codex treats "default" as an explicit standard-routing sentinel, not as a
|
||||
// service tier sent to the backend. Fast follows the gateway's existing
|
||||
// canonicalization and therefore becomes "priority"; flex stays "flex".
|
||||
canonicalTier := normalizedOpenAIServiceTierValue(serviceTier)
|
||||
// This backport has no Codex model-catalog snapshot with which to validate
|
||||
// arbitrary tier ids. Keep the hint to the two effective tiers Codex itself
|
||||
// selects; default, missing, and other gateway-compatible API values remain
|
||||
// model-only rather than expanding the ChatGPT routing protocol here.
|
||||
switch canonicalTier {
|
||||
case OpenAIFastTierPriority, OpenAIFastTierFlex:
|
||||
default:
|
||||
canonicalTier = ""
|
||||
}
|
||||
|
||||
hint := "model=" + model
|
||||
if canonicalTier != "" {
|
||||
hint += ";tier=" + canonicalTier
|
||||
}
|
||||
if !httpguts.ValidHeaderFieldValue(hint) {
|
||||
return
|
||||
}
|
||||
headers.Set(openAICodexRoutingHintHeader, hint)
|
||||
}
|
||||
|
||||
func deleteOpenAIHeaderEqualFold(headers http.Header, name string) {
|
||||
if headers == nil {
|
||||
return
|
||||
}
|
||||
name = strings.TrimSpace(name)
|
||||
for key := range headers {
|
||||
if strings.EqualFold(strings.TrimSpace(key), name) {
|
||||
delete(headers, key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func setOpenAICodexRoutingHintFromBody(headers http.Header, account *Account, body []byte) {
|
||||
fields := gjson.GetManyBytes(body, "model", "service_tier")
|
||||
setOpenAICodexRoutingHint(headers, account, fields[0].String(), fields[1].String())
|
||||
}
|
||||
|
||||
// logOpenAIRoutingDiagnostics records only gateway-derived routing state. In
|
||||
// particular, it deliberately does not include any header values, tokens, or
|
||||
// credentials because these diagnostics run on authentication-bearing paths.
|
||||
func logOpenAIRoutingDiagnostics(
|
||||
ctx context.Context,
|
||||
account *Account,
|
||||
transport string,
|
||||
model string,
|
||||
serviceTier string,
|
||||
hintGenerated bool,
|
||||
wsAffinityDecision string,
|
||||
) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
accountID := int64(0)
|
||||
if account != nil {
|
||||
accountID = account.ID
|
||||
}
|
||||
|
||||
logger.FromContext(ctx).Debug("openai routing decision",
|
||||
zap.String("component", "service.openai_routing"),
|
||||
zap.String("transport", strings.TrimSpace(transport)),
|
||||
zap.Int64("account_id", accountID),
|
||||
zap.String("final_model", strings.TrimSpace(model)),
|
||||
zap.String("final_service_tier", normalizedOpenAIServiceTierValue(serviceTier)),
|
||||
zap.Bool("routing_hint_generated", hintGenerated),
|
||||
zap.String("ws_affinity_decision", strings.TrimSpace(wsAffinityDecision)),
|
||||
)
|
||||
}
|
||||
|
||||
func logOpenAIRoutingDiagnosticsFromBody(
|
||||
ctx context.Context,
|
||||
account *Account,
|
||||
transport string,
|
||||
headers http.Header,
|
||||
body []byte,
|
||||
wsAffinityDecision string,
|
||||
) {
|
||||
fields := gjson.GetManyBytes(body, "model", "service_tier")
|
||||
logOpenAIRoutingDiagnostics(
|
||||
ctx,
|
||||
account,
|
||||
transport,
|
||||
fields[0].String(),
|
||||
fields[1].String(),
|
||||
strings.TrimSpace(headers.Get(openAICodexRoutingHintHeader)) != "",
|
||||
wsAffinityDecision,
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,369 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSetOpenAICodexRoutingHintCanonicalizesOfficialServiceTiers(t *testing.T) {
|
||||
oauthAccount := &Account{Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||||
tests := []struct {
|
||||
name string
|
||||
model string
|
||||
serviceTier string
|
||||
want string
|
||||
}{
|
||||
{name: "fast alias", model: "gpt-5.6", serviceTier: " fast ", want: "model=gpt-5.6;tier=priority"},
|
||||
{name: "priority", model: "gpt-5.6", serviceTier: "priority", want: "model=gpt-5.6;tier=priority"},
|
||||
{name: "flex", model: "gpt-5.6", serviceTier: "flex", want: "model=gpt-5.6;tier=flex"},
|
||||
{name: "explicit default sentinel", model: "gpt-5.6", serviceTier: "default", want: "model=gpt-5.6"},
|
||||
{name: "omitted tier", model: "gpt-5.6", want: "model=gpt-5.6"},
|
||||
{name: "auto is not expanded without catalog support", model: "gpt-5.6", serviceTier: "auto", want: "model=gpt-5.6"},
|
||||
{name: "scale is not expanded without catalog support", model: "gpt-5.6", serviceTier: "scale", want: "model=gpt-5.6"},
|
||||
{name: "unknown tier does not expand protocol", model: "gpt-5.6", serviceTier: "turbo", want: "model=gpt-5.6"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
headers := make(http.Header)
|
||||
setOpenAICodexRoutingHint(headers, oauthAccount, tt.model, tt.serviceTier)
|
||||
require.Equal(t, tt.want, headers.Get(openAICodexRoutingHintHeader))
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("invalid header value is omitted", func(t *testing.T) {
|
||||
headers := make(http.Header)
|
||||
setOpenAICodexRoutingHint(headers, oauthAccount, "gpt-5.6\ninvalid", "priority")
|
||||
require.Empty(t, headers.Get(openAICodexRoutingHintHeader))
|
||||
})
|
||||
|
||||
for _, model := range []string{"gpt-5.6;evil", "gpt=5.6"} {
|
||||
t.Run("delimiter in model is omitted: "+model, func(t *testing.T) {
|
||||
headers := make(http.Header)
|
||||
setOpenAICodexRoutingHint(headers, oauthAccount, model, "priority")
|
||||
require.Empty(t, headers.Get(openAICodexRoutingHintHeader))
|
||||
})
|
||||
}
|
||||
|
||||
t.Run("api key strips gateway-owned hint in every key casing", func(t *testing.T) {
|
||||
headers := make(http.Header)
|
||||
headers[openAICodexRoutingHintHeader] = []string{"lowercase-spoof"}
|
||||
headers["X-Codex-Routing-Hint"] = []string{"canonical-spoof"}
|
||||
setOpenAICodexRoutingHint(headers, &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, "gpt-5.6", "priority")
|
||||
for key := range headers {
|
||||
require.False(t, strings.EqualFold(key, openAICodexRoutingHintHeader))
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("oauth replaces spoofed lowercase hint", func(t *testing.T) {
|
||||
headers := make(http.Header)
|
||||
headers[openAICodexRoutingHintHeader] = []string{"model=spoof;tier=flex"}
|
||||
setOpenAICodexRoutingHint(headers, oauthAccount, "gpt-5.6", "priority")
|
||||
require.Equal(t, "model=gpt-5.6;tier=priority", headers.Get(openAICodexRoutingHintHeader))
|
||||
require.Len(t, headers, 1)
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIOAuthHTTPBuildersSendRoutingHintFromFinalBody(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
oauthAccount := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"chatgpt_account_id": "test-account",
|
||||
},
|
||||
}
|
||||
svc := &OpenAIGatewayService{}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
body []byte
|
||||
want string
|
||||
}{
|
||||
{name: "fast", body: []byte(`{"model":"gpt-5.6-codex","service_tier":"fast"}`), want: "model=gpt-5.6-codex;tier=priority"},
|
||||
{name: "flex", body: []byte(`{"model":"gpt-5.6-codex","service_tier":"flex"}`), want: "model=gpt-5.6-codex;tier=flex"},
|
||||
{name: "default", body: []byte(`{"model":"gpt-5.6-codex","service_tier":"default"}`), want: "model=gpt-5.6-codex"},
|
||||
{name: "omitted", body: []byte(`{"model":"gpt-5.6-codex"}`), want: "model=gpt-5.6-codex"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
for _, passthrough := range []bool{false, true} {
|
||||
mode := "ordinary"
|
||||
if passthrough {
|
||||
mode = "passthrough"
|
||||
}
|
||||
t.Run(mode, func(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(tt.body))
|
||||
|
||||
var req *http.Request
|
||||
var err error
|
||||
if passthrough {
|
||||
req, err = svc.buildUpstreamRequestOpenAIPassthrough(context.Background(), c, oauthAccount, tt.body, "test-token")
|
||||
} else {
|
||||
req, err = svc.buildUpstreamRequest(context.Background(), c, oauthAccount, tt.body, "test-token", false, "", true)
|
||||
}
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tt.want, req.Header.Get(openAICodexRoutingHintHeader))
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIHTTPPassthroughStripsOnlyOAuthLegacyResponsesBeta(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{
|
||||
Security: config.SecurityConfig{
|
||||
URLAllowlist: config.URLAllowlistConfig{Enabled: false},
|
||||
},
|
||||
}}
|
||||
body := []byte(`{"model":"gpt-5.6-codex","service_tier":"priority"}`)
|
||||
|
||||
build := func(t *testing.T, account *Account, betaValues []string, rawLowercaseKey bool) http.Header {
|
||||
t.Helper()
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
|
||||
if rawLowercaseKey {
|
||||
c.Request.Header["openai-beta"] = append([]string(nil), betaValues...)
|
||||
} else {
|
||||
for _, value := range betaValues {
|
||||
c.Request.Header.Add("OpenAI-Beta", value)
|
||||
}
|
||||
}
|
||||
|
||||
req, err := svc.buildUpstreamRequestOpenAIPassthrough(context.Background(), c, account, body, "test-token")
|
||||
require.NoError(t, err)
|
||||
return req.Header
|
||||
}
|
||||
|
||||
oauth := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"chatgpt_account_id": "test-account",
|
||||
},
|
||||
}
|
||||
apiKey := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{
|
||||
"api_key": "test-api-key",
|
||||
},
|
||||
}
|
||||
|
||||
t.Run("oauth legacy only is removed including raw lowercase key", func(t *testing.T) {
|
||||
headers := build(t, oauth, []string{"responses=experimental"}, true)
|
||||
require.Empty(t, headers.Values("OpenAI-Beta"))
|
||||
})
|
||||
|
||||
t.Run("oauth mixed beta preserves independent tokens", func(t *testing.T) {
|
||||
headers := build(t, oauth, []string{
|
||||
"responses=experimental, future_feature=v1",
|
||||
"another_feature=v2, RESPONSES=EXPERIMENTAL",
|
||||
}, false)
|
||||
require.Equal(t, []string{"future_feature=v1", "another_feature=v2"}, headers.Values("OpenAI-Beta"))
|
||||
})
|
||||
|
||||
t.Run("api key explicit beta remains caller controlled", func(t *testing.T) {
|
||||
headers := build(t, apiKey, []string{"responses=experimental, future_feature=v1"}, false)
|
||||
require.Equal(t, []string{"responses=experimental, future_feature=v1"}, headers.Values("OpenAI-Beta"))
|
||||
})
|
||||
}
|
||||
|
||||
func TestBuildOpenAIWSHeadersSendsOAuthRoutingHintOnly(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
svc := &OpenAIGatewayService{}
|
||||
decision := OpenAIWSProtocolDecision{Transport: OpenAIUpstreamTransportResponsesWebsocketV2}
|
||||
|
||||
build := func(t *testing.T, account *Account, tier string) http.Header {
|
||||
headers, _, err := svc.buildOpenAIWSHeaders(
|
||||
context.Background(),
|
||||
c,
|
||||
account,
|
||||
"test-token",
|
||||
decision,
|
||||
true,
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
"gpt-5.6-codex",
|
||||
tier,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
return headers
|
||||
}
|
||||
|
||||
oauthAccount := &Account{
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"chatgpt_account_id": "test-account",
|
||||
},
|
||||
}
|
||||
require.Equal(t, "model=gpt-5.6-codex;tier=priority", build(t, oauthAccount, "fast").Get(openAICodexRoutingHintHeader))
|
||||
require.Equal(t, "model=gpt-5.6-codex", build(t, oauthAccount, "default").Get(openAICodexRoutingHintHeader))
|
||||
require.Empty(t, build(t, &Account{Platform: PlatformOpenAI, Type: AccountTypeAPIKey}, "priority").Get(openAICodexRoutingHintHeader))
|
||||
}
|
||||
|
||||
func TestOpenAIRoutingDiagnosticsUseFinalDerivedValuesOnly(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
logSink, restore := captureStructuredLog(t)
|
||||
defer restore()
|
||||
|
||||
account := &Account{
|
||||
ID: 917,
|
||||
Platform: PlatformOpenAI,
|
||||
Type: AccountTypeOAuth,
|
||||
Credentials: map[string]any{
|
||||
"chatgpt_account_id": "chatgpt-account",
|
||||
},
|
||||
}
|
||||
body := []byte(`{"model":"gpt-5.6-codex","service_tier":"fast"}`)
|
||||
svc := &OpenAIGatewayService{}
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
|
||||
c.Request.Header.Set("Authorization", "Bearer caller-secret")
|
||||
c.Request.Header.Set(openAICodexRoutingHintHeader, "model=caller-secret")
|
||||
_, err := svc.buildUpstreamRequest(context.Background(), c, account, body, "oauth-secret", false, "", true)
|
||||
require.NoError(t, err)
|
||||
|
||||
decision := OpenAIWSProtocolDecision{Transport: OpenAIUpstreamTransportResponsesWebsocketV2}
|
||||
_, _, err = svc.buildOpenAIWSHeaders(
|
||||
context.Background(), c, account, "oauth-secret", decision, true,
|
||||
"", "", "", "gpt-5.6-codex", "fast",
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.True(t, logSink.ContainsMessageAtLevel("openai routing decision", "debug"))
|
||||
require.True(t, logSink.ContainsFieldValue("account_id", "917"))
|
||||
require.True(t, logSink.ContainsFieldValue("final_model", "gpt-5.6-codex"))
|
||||
require.True(t, logSink.ContainsFieldValue("final_service_tier", "priority"))
|
||||
require.True(t, logSink.ContainsFieldValue("routing_hint_generated", "true"))
|
||||
require.True(t, logSink.ContainsFieldValue("transport", "http"))
|
||||
require.True(t, logSink.ContainsFieldValue("transport", string(OpenAIUpstreamTransportResponsesWebsocketV2)))
|
||||
require.True(t, logSink.ContainsFieldValue("ws_affinity_decision", "not_applicable"))
|
||||
require.True(t, logSink.ContainsFieldValue("ws_affinity_decision", "soft_routing_hint"))
|
||||
require.False(t, logSink.ContainsFieldValue("authorization", "caller-secret"))
|
||||
require.False(t, logSink.ContainsFieldValue("credentials", "oauth-secret"))
|
||||
require.False(t, logSink.ContainsFieldValue("routing_hint", "caller-secret"))
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnPoolPreferredContinuationIgnoresRoutingHintChanges(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
||||
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
||||
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 2
|
||||
|
||||
pool := newOpenAIWSConnPool(cfg)
|
||||
dialer := &openAIWSCountingDialer{}
|
||||
pool.setClientDialerForTest(dialer)
|
||||
account := &Account{ID: 913, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||||
|
||||
acquire := func(t *testing.T, hint, preferred string, forcePreferred bool) *openAIWSConnLease {
|
||||
t.Helper()
|
||||
headers := make(http.Header)
|
||||
if hint != "" {
|
||||
headers.Set(openAICodexRoutingHintHeader, hint)
|
||||
}
|
||||
lease, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
||||
Account: account,
|
||||
WSURL: "wss://example.com/v1/responses",
|
||||
Headers: headers,
|
||||
PreferredConnID: preferred,
|
||||
ForcePreferredConn: forcePreferred,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, lease)
|
||||
return lease
|
||||
}
|
||||
|
||||
standard := acquire(t, "model=gpt-5.6-codex", "", false)
|
||||
connID := standard.ConnID()
|
||||
standard.Release()
|
||||
|
||||
priority := acquire(t, "model=gpt-5.6-codex;tier=priority", connID, true)
|
||||
require.True(t, priority.Reused())
|
||||
require.Equal(t, connID, priority.ConnID())
|
||||
priority.Release()
|
||||
|
||||
standardAgain := acquire(t, "model=gpt-5.6-codex", connID, true)
|
||||
require.True(t, standardAgain.Reused())
|
||||
require.Equal(t, connID, standardAgain.ConnID())
|
||||
standardAgain.Release()
|
||||
|
||||
require.Equal(t, 1, dialer.DialCount(), "routing hint is dial-time advisory, not continuation compatibility")
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnPoolUsesRoutingHintAsSoftDialAffinity(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 4
|
||||
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
||||
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 4
|
||||
|
||||
pool := newOpenAIWSConnPool(cfg)
|
||||
dialer := &openAIWSCountingDialer{}
|
||||
pool.setClientDialerForTest(dialer)
|
||||
account := &Account{ID: 913, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||||
|
||||
acquire := func(t *testing.T, hint string) *openAIWSConnLease {
|
||||
headers := make(http.Header)
|
||||
headers.Set(openAICodexRoutingHintHeader, hint)
|
||||
lease, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{
|
||||
Account: account,
|
||||
WSURL: "wss://example.com/v1/responses",
|
||||
Headers: headers,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, lease)
|
||||
return lease
|
||||
}
|
||||
|
||||
priority := acquire(t, "model=gpt-5.6-codex;tier=priority")
|
||||
priorityConnID := priority.ConnID()
|
||||
priority.Release()
|
||||
|
||||
priorityAgain := acquire(t, "model=gpt-5.6-codex;tier=priority")
|
||||
require.True(t, priorityAgain.Reused())
|
||||
require.Equal(t, priorityConnID, priorityAgain.ConnID())
|
||||
priorityAgain.Release()
|
||||
|
||||
flex := acquire(t, "model=gpt-5.6-codex;tier=flex")
|
||||
require.False(t, flex.Reused())
|
||||
require.NotEqual(t, priorityConnID, flex.ConnID())
|
||||
flex.Release()
|
||||
|
||||
otherModel := acquire(t, "model=gpt-5.5-codex;tier=priority")
|
||||
require.False(t, otherModel.Reused())
|
||||
require.NotEqual(t, priorityConnID, otherModel.ConnID())
|
||||
otherModel.Release()
|
||||
|
||||
defaultTier := acquire(t, "model=gpt-5.6-codex")
|
||||
require.False(t, defaultTier.Reused())
|
||||
require.NotEqual(t, priorityConnID, defaultTier.ConnID())
|
||||
defaultConnID := defaultTier.ConnID()
|
||||
defaultTier.Release()
|
||||
|
||||
defaultAgain := acquire(t, "model=gpt-5.6-codex")
|
||||
require.True(t, defaultAgain.Reused())
|
||||
require.Equal(t, defaultConnID, defaultAgain.ConnID())
|
||||
defaultAgain.Release()
|
||||
|
||||
require.Equal(t, 4, dialer.DialCount())
|
||||
}
|
||||
@@ -617,7 +617,20 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
}
|
||||
}
|
||||
|
||||
wsHeaders, _, buildHdrErr := s.buildOpenAIWSHeaders(ctx, c, account, token, wsDecision, isCodexCLI, turnState, strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader)), firstPayload.promptCacheKey)
|
||||
firstRoutingFields := gjson.GetManyBytes(firstPayload.payloadRaw, "model", "service_tier")
|
||||
wsHeaders, _, buildHdrErr := s.buildOpenAIWSHeaders(
|
||||
ctx,
|
||||
c,
|
||||
account,
|
||||
token,
|
||||
wsDecision,
|
||||
isCodexCLI,
|
||||
turnState,
|
||||
strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader)),
|
||||
firstPayload.promptCacheKey,
|
||||
firstRoutingFields[0].String(),
|
||||
firstRoutingFields[1].String(),
|
||||
)
|
||||
if buildHdrErr != nil {
|
||||
return fmt.Errorf("build ws headers: %w", buildHdrErr)
|
||||
}
|
||||
@@ -785,6 +798,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
}
|
||||
|
||||
sendAndRelay := func(turn int, lease *openAIWSConnLease, payload []byte, payloadBytes int, originalModel string, imageBillingModel string, imageSizeTier string, imageInputSize string) (*OpenAIForwardResult, error) {
|
||||
responseModelObserver := &upstreamResponseModelObserver{}
|
||||
if lease == nil {
|
||||
return nil, errors.New("upstream websocket lease is nil")
|
||||
}
|
||||
@@ -852,6 +866,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
}
|
||||
|
||||
eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(upstreamMessage)
|
||||
responseModelObserver.ObserveOpenAI(upstreamMessage, eventType)
|
||||
if responseID == "" && eventResponseID != "" {
|
||||
responseID = eventResponseID
|
||||
}
|
||||
@@ -1026,18 +1041,20 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
}
|
||||
imageCount := imageCounter.Count()
|
||||
result := &OpenAIForwardResult{
|
||||
RequestID: responseID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: mappedModel,
|
||||
ServiceTier: extractOpenAIServiceTierFromBody(payload),
|
||||
ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(payload, mappedModel, originalModel), payload, mappedModel),
|
||||
Stream: reqStream,
|
||||
OpenAIWSMode: true,
|
||||
UpstreamTerminalEvent: terminalEvent,
|
||||
ResponseHeaders: lease.HandshakeHeaders(),
|
||||
Duration: time.Since(turnStart),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
RequestID: responseID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: mappedModel,
|
||||
UpstreamResponseModel: responseModelObserver.Model(),
|
||||
UpstreamResponseModelConflict: responseModelObserver.Conflict(),
|
||||
ServiceTier: extractOpenAIServiceTierFromBody(payload),
|
||||
ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(payload, mappedModel, originalModel), payload, mappedModel),
|
||||
Stream: reqStream,
|
||||
OpenAIWSMode: true,
|
||||
UpstreamTerminalEvent: terminalEvent,
|
||||
ResponseHeaders: lease.HandshakeHeaders(),
|
||||
Duration: time.Since(turnStart),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
}
|
||||
if replayInput := replayCollector.Items(); len(replayInput) > 0 {
|
||||
result.wsReplayInput = replayInput
|
||||
@@ -1632,16 +1649,30 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient(
|
||||
if parseErr != nil {
|
||||
return parseErr
|
||||
}
|
||||
nextRoutingFields := gjson.GetManyBytes(nextPayload.payloadRaw, "model", "service_tier")
|
||||
if nextPayload.promptCacheKey != "" {
|
||||
// ingress 会话在整个客户端 WS 生命周期内复用同一上游连接;
|
||||
// prompt_cache_key 对握手头的更新仅在未来需要重新建连时生效。
|
||||
updatedHeaders, _, updHdrErr := s.buildOpenAIWSHeaders(ctx, c, account, token, wsDecision, isCodexCLI, turnState, strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader)), nextPayload.promptCacheKey)
|
||||
updatedHeaders, _, updHdrErr := s.buildOpenAIWSHeaders(
|
||||
ctx,
|
||||
c,
|
||||
account,
|
||||
token,
|
||||
wsDecision,
|
||||
isCodexCLI,
|
||||
turnState,
|
||||
strings.TrimSpace(c.GetHeader(openAIWSTurnMetadataHeader)),
|
||||
nextPayload.promptCacheKey,
|
||||
nextRoutingFields[0].String(),
|
||||
nextRoutingFields[1].String(),
|
||||
)
|
||||
if updHdrErr != nil {
|
||||
logOpenAIWSModeInfo("ingress_ws_update_headers_failed account_id=%d err=%v", account.ID, updHdrErr)
|
||||
} else {
|
||||
baseAcquireReq.Headers = updatedHeaders
|
||||
}
|
||||
}
|
||||
setOpenAICodexRoutingHint(baseAcquireReq.Headers, account, nextRoutingFields[0].String(), nextRoutingFields[1].String())
|
||||
if nextPayload.previousResponseID != "" {
|
||||
expectedPrev := strings.TrimSpace(lastTurnResponseID)
|
||||
chainedFromLast := expectedPrev != "" && nextPayload.previousResponseID == expectedPrev
|
||||
|
||||
@@ -551,13 +551,22 @@ func TestNormalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(t *testing.T)
|
||||
t.Parallel()
|
||||
|
||||
normalized, err := normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(
|
||||
[]byte(`{"model":"gpt-5.1","input":[1],"previous_response_id":"resp_x","metadata":{"b":2,"a":1}}`),
|
||||
[]byte(`{"model":"gpt-5.1","input":[1],"previous_response_id":"resp_x","client_metadata":{"request_start_ms":"1"},"stream_options":{"include_usage":true},"generate":false,"metadata":{"b":2,"a":1}}`),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.False(t, gjson.GetBytes(normalized, "input").Exists())
|
||||
require.False(t, gjson.GetBytes(normalized, "previous_response_id").Exists())
|
||||
require.False(t, gjson.GetBytes(normalized, "client_metadata").Exists())
|
||||
require.False(t, gjson.GetBytes(normalized, "stream_options").Exists())
|
||||
require.False(t, gjson.GetBytes(normalized, "generate").Exists())
|
||||
require.Equal(t, float64(1), gjson.GetBytes(normalized, "metadata.a").Float())
|
||||
|
||||
normalized, err = normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(
|
||||
[]byte(`{"model":"gpt-5.1","generate":true}`),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.True(t, gjson.GetBytes(normalized, "generate").Bool())
|
||||
|
||||
_, err = normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(nil)
|
||||
require.Error(t, err)
|
||||
|
||||
@@ -662,6 +671,36 @@ func TestShouldKeepIngressPreviousResponseID(t *testing.T) {
|
||||
require.Equal(t, "strict_incremental_ok", reason)
|
||||
})
|
||||
|
||||
t.Run("codex_prewarm_to_business_keep", func(t *testing.T) {
|
||||
prewarmPayload := []byte(`{
|
||||
"type":"response.create",
|
||||
"model":"gpt-5.1",
|
||||
"store":false,
|
||||
"generate":false,
|
||||
"client_metadata":{"x-codex-ws-stream-request-start-ms":"100"},
|
||||
"stream_options":{"include_usage":true},
|
||||
"input":[{"type":"input_text","text":"hello"}]
|
||||
}`)
|
||||
businessPayload := []byte(`{
|
||||
"type":"response.create",
|
||||
"model":"gpt-5.1",
|
||||
"store":false,
|
||||
"client_metadata":{"x-codex-ws-stream-request-start-ms":"200"},
|
||||
"previous_response_id":"resp_prewarm",
|
||||
"input":[{"type":"input_text","text":"hello"}]
|
||||
}`)
|
||||
|
||||
keep, reason, err := shouldKeepIngressPreviousResponseID(
|
||||
prewarmPayload,
|
||||
businessPayload,
|
||||
"resp_prewarm",
|
||||
false,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.True(t, keep)
|
||||
require.Equal(t, "strict_incremental_ok", reason)
|
||||
})
|
||||
|
||||
t.Run("missing_previous_response_id", func(t *testing.T) {
|
||||
payload := []byte(`{"type":"response.create","model":"gpt-5.1","input":[]}`)
|
||||
keep, reason, err := shouldKeepIngressPreviousResponseID(previousPayload, payload, "resp_turn_1", false)
|
||||
|
||||
@@ -75,6 +75,8 @@ func (s *OpenAIGatewayService) buildOpenAIWSHeaders(
|
||||
turnState string,
|
||||
turnMetadata string,
|
||||
promptCacheKey string,
|
||||
routingModel string,
|
||||
routingServiceTier string,
|
||||
) (http.Header, openAIWSSessionHeaderResolution, error) {
|
||||
headers := make(http.Header)
|
||||
if account == nil || !account.IsOpenAIAgentIdentity() {
|
||||
@@ -157,6 +159,16 @@ func (s *OpenAIGatewayService) buildOpenAIWSHeaders(
|
||||
// 账号级请求头覆写(仅 openai api_key 账号启用时生效;OAuth 路径 no-op)。
|
||||
// 覆盖所有 WS 模式(ctx_pool/dedicated/passthrough)的握手头。
|
||||
account.ApplyHeaderOverrides(headers)
|
||||
setOpenAICodexRoutingHint(headers, account, routingModel, routingServiceTier)
|
||||
logOpenAIRoutingDiagnostics(
|
||||
ctx,
|
||||
account,
|
||||
string(decision.Transport),
|
||||
routingModel,
|
||||
routingServiceTier,
|
||||
strings.TrimSpace(headers.Get(openAICodexRoutingHintHeader)) != "",
|
||||
"soft_routing_hint",
|
||||
)
|
||||
|
||||
return headers, sessionResolution, nil
|
||||
}
|
||||
@@ -449,6 +461,17 @@ func normalizeOpenAIWSPayloadWithoutInputAndPreviousResponseID(payload []byte) (
|
||||
}
|
||||
delete(decoded, "input")
|
||||
delete(decoded, "previous_response_id")
|
||||
// Codex changes transport-only metadata for every response.create. These fields
|
||||
// do not alter the context referenced by previous_response_id and are excluded
|
||||
// from Codex's own websocket reuse comparison.
|
||||
delete(decoded, "client_metadata")
|
||||
delete(decoded, "stream_options")
|
||||
// Official Codex prewarms a connection with generate=false, then omits the
|
||||
// field on the business request that continues from the prewarm response.
|
||||
// Only normalize false so a meaningful generate=true change remains visible.
|
||||
if generate, ok := decoded["generate"].(bool); ok && !generate {
|
||||
delete(decoded, "generate")
|
||||
}
|
||||
return json.Marshal(decoded)
|
||||
}
|
||||
|
||||
|
||||
@@ -414,6 +414,8 @@ func TestOpenAIGatewayService_BuildOpenAIWSHeadersPreservesCodexIdentity(t *test
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
"",
|
||||
)
|
||||
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -36,6 +36,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
|
||||
if s == nil || account == nil {
|
||||
return nil, wrapOpenAIWSFallback("invalid_state", errors.New("service or account is nil"))
|
||||
}
|
||||
responseModelObserver := &upstreamResponseModelObserver{}
|
||||
|
||||
wsURL, err := s.buildOpenAIResponsesWSURL(account)
|
||||
if err != nil {
|
||||
@@ -136,7 +137,19 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
|
||||
storeDisabledConnMode := s.openAIWSStoreDisabledConnMode()
|
||||
forceNewConnByPolicy := shouldForceNewConnOnStoreDisabled(storeDisabledConnMode, lastFailureReason)
|
||||
forceNewConn := forceNewConnByPolicy && storeDisabled && previousResponseID == "" && sessionHash != "" && preferredConnID == ""
|
||||
wsHeaders, sessionResolution, buildHdrErr := s.buildOpenAIWSHeaders(ctx, c, account, token, decision, isCodexCLI, turnState, turnMetadata, promptCacheKey)
|
||||
wsHeaders, sessionResolution, buildHdrErr := s.buildOpenAIWSHeaders(
|
||||
ctx,
|
||||
c,
|
||||
account,
|
||||
token,
|
||||
decision,
|
||||
isCodexCLI,
|
||||
turnState,
|
||||
turnMetadata,
|
||||
promptCacheKey,
|
||||
openAIWSPayloadString(payload, "model"),
|
||||
openAIWSPayloadString(payload, "service_tier"),
|
||||
)
|
||||
if buildHdrErr != nil {
|
||||
return nil, fmt.Errorf("build ws headers: %w", buildHdrErr)
|
||||
}
|
||||
@@ -512,6 +525,7 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
|
||||
if eventType == "" {
|
||||
continue
|
||||
}
|
||||
responseModelObserver.ObserveOpenAI(message, eventType)
|
||||
eventCount++
|
||||
if firstEventType == "" {
|
||||
firstEventType = eventType
|
||||
@@ -748,20 +762,22 @@ func (s *OpenAIGatewayService) forwardOpenAIWSV2(
|
||||
)
|
||||
|
||||
return &OpenAIForwardResult{
|
||||
RequestID: responseID,
|
||||
Usage: *usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: mappedModel,
|
||||
ImageCount: imageCounter.Count(),
|
||||
ImageOutputSizes: imageCounter.Sizes(),
|
||||
ServiceTier: extractOpenAIServiceTier(reqBody),
|
||||
ReasoningEffort: extractOpenAIReasoningEffort(reqBody, mappedModel, originalModel),
|
||||
Stream: reqStream,
|
||||
OpenAIWSMode: true,
|
||||
UpstreamTerminalEvent: upstreamTerminalEvent,
|
||||
ResponseHeaders: lease.HandshakeHeaders(),
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
RequestID: responseID,
|
||||
Usage: *usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: mappedModel,
|
||||
UpstreamResponseModel: responseModelObserver.Model(),
|
||||
UpstreamResponseModelConflict: responseModelObserver.Conflict(),
|
||||
ImageCount: imageCounter.Count(),
|
||||
ImageOutputSizes: imageCounter.Sizes(),
|
||||
ServiceTier: extractOpenAIServiceTier(reqBody),
|
||||
ReasoningEffort: extractOpenAIReasoningEffort(reqBody, mappedModel, originalModel),
|
||||
Stream: reqStream,
|
||||
OpenAIWSMode: true,
|
||||
UpstreamTerminalEvent: upstreamTerminalEvent,
|
||||
ResponseHeaders: lease.HandshakeHeaders(),
|
||||
Duration: time.Since(startTime),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -180,6 +180,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
||||
if writeClientMessage == nil {
|
||||
return nil, errors.New("client websocket writer is nil")
|
||||
}
|
||||
responseModelObserver := &upstreamResponseModelObserver{}
|
||||
|
||||
body, err := prepareOpenAIWSHTTPBridgeBody(payload)
|
||||
if err != nil {
|
||||
@@ -299,18 +300,20 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
||||
resultWithUsage := func() *OpenAIForwardResult {
|
||||
imageCount := imageCounter.Count()
|
||||
result := &OpenAIForwardResult{
|
||||
RequestID: responseID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: mappedModel,
|
||||
ServiceTier: extractOpenAIServiceTierFromBody(body),
|
||||
ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(body, mappedModel, originalModel), body, mappedModel),
|
||||
Stream: reqStream,
|
||||
OpenAIWSMode: true,
|
||||
UpstreamTerminalEvent: upstreamTerminalEvent,
|
||||
ResponseHeaders: cloneHeader(resp.Header),
|
||||
Duration: time.Since(turnStart),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
RequestID: responseID,
|
||||
Usage: usage,
|
||||
Model: originalModel,
|
||||
UpstreamModel: mappedModel,
|
||||
UpstreamResponseModel: responseModelObserver.Model(),
|
||||
UpstreamResponseModelConflict: responseModelObserver.Conflict(),
|
||||
ServiceTier: extractOpenAIServiceTierFromBody(body),
|
||||
ReasoningEffort: ApplyThinkingEnabledFallback(extractOpenAIReasoningEffortFromBody(body, mappedModel, originalModel), body, mappedModel),
|
||||
Stream: reqStream,
|
||||
OpenAIWSMode: true,
|
||||
UpstreamTerminalEvent: upstreamTerminalEvent,
|
||||
ResponseHeaders: cloneHeader(resp.Header),
|
||||
Duration: time.Since(turnStart),
|
||||
FirstTokenMs: firstTokenMs,
|
||||
}
|
||||
if replayInput := replayCollector.Items(); len(replayInput) > 0 {
|
||||
result.wsReplayInput = replayInput
|
||||
@@ -355,6 +358,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
||||
upstreamMessage = normalized
|
||||
}
|
||||
eventType, eventResponseID, _ := parseOpenAIWSEventEnvelope(upstreamMessage)
|
||||
responseModelObserver.ObserveOpenAI(upstreamMessage, eventType)
|
||||
if responseID == "" && eventResponseID != "" {
|
||||
responseID = eventResponseID
|
||||
}
|
||||
@@ -421,8 +425,18 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn(
|
||||
upstreamEventErr = errors.New(errMessage)
|
||||
}
|
||||
|
||||
// 客户端写出副本改写容量降载码:Codex 对 error/response.failed 中的
|
||||
// server_is_overloaded / slow_down 判致命并终止会话,改写后走客户端内置
|
||||
// 重试。账号状态与终止事件判定(下方 handleOpenAIWSTerminalTransientFailure)
|
||||
// 仍使用未改写的 upstreamMessage。
|
||||
clientMessage := upstreamMessage
|
||||
if eventType == "error" || eventType == "response.failed" {
|
||||
if rewritten, changed := sanitizeOpenAICapacityShedErrorCodeForClient(clientMessage); changed {
|
||||
clientMessage = rewritten
|
||||
}
|
||||
}
|
||||
if !clientDisconnected {
|
||||
if err := writeClientMessage(upstreamMessage); err != nil {
|
||||
if err := writeClientMessage(clientMessage); err != nil {
|
||||
if isOpenAIWSClientDisconnectError(err) {
|
||||
clientDisconnected = true
|
||||
closeStatus, closeReason := summarizeOpenAIWSReadCloseError(err)
|
||||
|
||||
@@ -220,6 +220,69 @@ func TestProxyOpenAIWSHTTPBridgeTurnSSEErrorFailoverSafety(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// 桥接转发 error / response.failed 给 WS 客户端前必须把容量降载码改写为可重试
|
||||
// 的 server_error:Codex 对 server_is_overloaded/slow_down 判致命并终止会话。
|
||||
// 账号状态判定使用改写前的原始事件,不受影响。
|
||||
func TestProxyOpenAIWSHTTPBridgeTurnRewritesCapacityShedCodeForClient(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
turn int
|
||||
body string
|
||||
wantErr bool
|
||||
}{
|
||||
{
|
||||
name: "turn2_error_frame",
|
||||
turn: 2,
|
||||
body: "data: {\"type\":\"error\",\"error\":{\"type\":\"service_unavailable_error\",\"code\":\"server_is_overloaded\",\"message\":\"Our servers are currently overloaded. Please try again later.\"}}\n\n",
|
||||
wantErr: true,
|
||||
},
|
||||
{
|
||||
// response.failed 不走 error 事件分支:即便 turn 1 也会被当终止事件
|
||||
// 原样转发(不 failover),因此改写必须在这里同样生效。
|
||||
name: "turn1_bare_response_failed",
|
||||
turn: 1,
|
||||
body: "data: {\"type\":\"response.failed\",\"response\":{\"id\":\"resp_shed\",\"status\":\"failed\",\"error\":{\"code\":\"server_is_overloaded\",\"message\":\"Our servers are currently overloaded. Please try again later.\"}}}\n\n",
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
upstream := &httpUpstreamRecorder{resp: &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(tt.body)),
|
||||
}}
|
||||
svc := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream}
|
||||
account := &Account{ID: 11, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Concurrency: 1}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
|
||||
var writes [][]byte
|
||||
|
||||
_, err := svc.proxyOpenAIWSHTTPBridgeTurn(
|
||||
context.Background(), c, account, "sk-test", payload, len(payload),
|
||||
"gpt-5", "", "", "", "", tt.turn,
|
||||
func(message []byte) error {
|
||||
writes = append(writes, append([]byte(nil), message...))
|
||||
return nil
|
||||
},
|
||||
)
|
||||
|
||||
if tt.wantErr {
|
||||
require.Error(t, err)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
require.Len(t, writes, 1)
|
||||
require.Contains(t, string(writes[0]), `"code":"server_error"`)
|
||||
require.NotContains(t, string(writes[0]), "server_is_overloaded")
|
||||
require.Contains(t, string(writes[0]), "Our servers are currently overloaded")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProxyOpenAIWSHTTPBridgeTurnRequiresTerminalEvent(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
|
||||
@@ -76,6 +76,10 @@ type openAIWSAcquireRequest struct {
|
||||
ForcePreferredConn bool
|
||||
}
|
||||
|
||||
type openAIWSHandshakeCompatibilityKey struct {
|
||||
betaFeatures string
|
||||
}
|
||||
|
||||
type openAIWSConnLease struct {
|
||||
pool *openAIWSConnPool
|
||||
accountID int64
|
||||
@@ -240,8 +244,9 @@ type openAIWSConn struct {
|
||||
id string
|
||||
ws openAIWSClientConn
|
||||
|
||||
handshakeHeaders http.Header
|
||||
betaFeatures string
|
||||
handshakeHeaders http.Header
|
||||
handshakeCompatibility openAIWSHandshakeCompatibilityKey
|
||||
routingAffinity string
|
||||
|
||||
leaseCh chan struct{}
|
||||
closedCh chan struct{}
|
||||
@@ -305,6 +310,14 @@ func (c *openAIWSConn) acquire(ctx context.Context) error {
|
||||
case <-c.closedCh:
|
||||
return errOpenAIWSConnClosed
|
||||
case <-c.leaseCh:
|
||||
// A cancellation and a lease delivery can become ready together. Once
|
||||
// the semaphore token has been consumed, check the context again and
|
||||
// return it before reporting cancellation so a canceled waiter cannot
|
||||
// strand a pooled connection.
|
||||
if err := ctx.Err(); err != nil {
|
||||
c.release()
|
||||
return err
|
||||
}
|
||||
select {
|
||||
case <-c.closedCh:
|
||||
c.release()
|
||||
@@ -525,8 +538,12 @@ func (c *openAIWSConn) handshakeHeader(name string) string {
|
||||
return strings.TrimSpace(c.handshakeHeaders.Get(strings.TrimSpace(name)))
|
||||
}
|
||||
|
||||
func (c *openAIWSConn) matchesBetaFeatures(betaFeatures string) bool {
|
||||
return c != nil && c.betaFeatures == betaFeatures
|
||||
func (c *openAIWSConn) matchesHandshakeCompatibility(compatibility openAIWSHandshakeCompatibilityKey) bool {
|
||||
return c != nil && c.handshakeCompatibility == compatibility
|
||||
}
|
||||
|
||||
func (c *openAIWSConn) matchesRoutingAffinity(routingAffinity string) bool {
|
||||
return c != nil && c.routingAffinity == routingAffinity
|
||||
}
|
||||
|
||||
func (c *openAIWSConn) isPrewarmed() bool {
|
||||
@@ -838,7 +855,8 @@ func (p *openAIWSConnPool) acquire(ctx context.Context, req openAIWSAcquireReque
|
||||
|
||||
retryAcquire:
|
||||
accountID := req.Account.ID
|
||||
betaFeatures := normalizeOpenAIWSBetaFeatures(req.Headers)
|
||||
compatibility := normalizeOpenAIWSHandshakeCompatibility(req.Headers)
|
||||
routingAffinity := normalizeOpenAIWSRoutingAffinity(req.Headers)
|
||||
effectiveMaxConns := p.effectiveMaxConnsByAccount(req.Account)
|
||||
if effectiveMaxConns <= 0 {
|
||||
return nil, errOpenAIWSConnQueueFull
|
||||
@@ -846,7 +864,7 @@ retryAcquire:
|
||||
var evicted []*openAIWSConn
|
||||
ap := p.getOrCreateAccountPool(accountID)
|
||||
ap.mu.Lock()
|
||||
ap.lastAcquire = cloneOpenAIWSAcquireRequestPtr(&req)
|
||||
acquireGeneration := ap.generation
|
||||
now := time.Now()
|
||||
if ap.lastCleanupAt.IsZero() || now.Sub(ap.lastCleanupAt) >= openAIWSAcquireCleanupInterval {
|
||||
evicted = p.cleanupAccountLocked(ap, now, effectiveMaxConns)
|
||||
@@ -866,7 +884,7 @@ retryAcquire:
|
||||
return nil, errOpenAIWSPreferredConnUnavailable
|
||||
}
|
||||
preferredConn, ok := ap.conns[preferredConnID]
|
||||
if !ok || !preferredConn.matchesBetaFeatures(betaFeatures) {
|
||||
if !ok || !preferredConn.matchesHandshakeCompatibility(compatibility) {
|
||||
p.recordConnPickDuration(time.Since(pickStartedAt))
|
||||
ap.mu.Unlock()
|
||||
closeOpenAIWSConns(evicted)
|
||||
@@ -895,6 +913,7 @@ retryAcquire:
|
||||
reused: true,
|
||||
}
|
||||
p.metrics.acquireReuseTotal.Add(1)
|
||||
p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req)
|
||||
p.ensureTargetIdleAsync(accountID)
|
||||
return lease, nil
|
||||
}
|
||||
@@ -942,12 +961,13 @@ retryAcquire:
|
||||
reused: true,
|
||||
}
|
||||
p.metrics.acquireReuseTotal.Add(1)
|
||||
p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req)
|
||||
p.ensureTargetIdleAsync(accountID)
|
||||
return lease, nil
|
||||
}
|
||||
|
||||
if preferredConnID != "" {
|
||||
if conn, ok := ap.conns[preferredConnID]; ok && conn.matchesBetaFeatures(betaFeatures) && conn.tryAcquire() {
|
||||
if conn, ok := ap.conns[preferredConnID]; ok && conn.matchesHandshakeCompatibility(compatibility) && conn.tryAcquire() {
|
||||
connPick := time.Since(pickStartedAt)
|
||||
p.recordConnPickDuration(connPick)
|
||||
ap.mu.Unlock()
|
||||
@@ -964,12 +984,16 @@ retryAcquire:
|
||||
}
|
||||
lease := &openAIWSConnLease{pool: p, accountID: accountID, conn: conn, connPick: connPick, reused: true}
|
||||
p.metrics.acquireReuseTotal.Add(1)
|
||||
p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req)
|
||||
p.ensureTargetIdleAsync(accountID)
|
||||
return lease, nil
|
||||
}
|
||||
}
|
||||
|
||||
best := p.pickLeastBusyConnLocked(ap, "", betaFeatures)
|
||||
// A routing hint is advisory at WebSocket dial time. Prefer a pooled
|
||||
// connection whose handshake used the same hint, but do not make that
|
||||
// preference a continuation compatibility requirement.
|
||||
best := p.pickLeastBusyConnWithRoutingAffinityLocked(ap, compatibility, routingAffinity)
|
||||
if best != nil && best.tryAcquire() {
|
||||
connPick := time.Since(pickStartedAt)
|
||||
p.recordConnPickDuration(connPick)
|
||||
@@ -987,11 +1011,12 @@ retryAcquire:
|
||||
}
|
||||
lease := &openAIWSConnLease{pool: p, accountID: accountID, conn: best, connPick: connPick, reused: true}
|
||||
p.metrics.acquireReuseTotal.Add(1)
|
||||
p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req)
|
||||
p.ensureTargetIdleAsync(accountID)
|
||||
return lease, nil
|
||||
}
|
||||
for _, conn := range ap.conns {
|
||||
if conn == nil || conn == best || !conn.matchesBetaFeatures(betaFeatures) {
|
||||
if conn == nil || conn == best || !conn.matchesHandshakeCompatibility(compatibility) || !conn.matchesRoutingAffinity(routingAffinity) {
|
||||
continue
|
||||
}
|
||||
if conn.tryAcquire() {
|
||||
@@ -1011,6 +1036,7 @@ retryAcquire:
|
||||
}
|
||||
lease := &openAIWSConnLease{pool: p, accountID: accountID, conn: conn, connPick: connPick, reused: true}
|
||||
p.metrics.acquireReuseTotal.Add(1)
|
||||
p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req)
|
||||
p.ensureTargetIdleAsync(accountID)
|
||||
return lease, nil
|
||||
}
|
||||
@@ -1018,12 +1044,18 @@ retryAcquire:
|
||||
}
|
||||
|
||||
if !req.ForceNewConn && len(ap.conns)+ap.creating >= effectiveMaxConns {
|
||||
compatible := p.pickLeastBusyConnLocked(ap, "", betaFeatures)
|
||||
if idle := p.pickOldestIdleConnWithDifferentBetaFeaturesLocked(ap, betaFeatures); idle != nil {
|
||||
affine := p.pickLeastBusyConnWithRoutingAffinityLocked(ap, compatibility, routingAffinity)
|
||||
if idle := p.pickOldestIdleConnWithoutHandshakeCompatibilityOrRoutingAffinityLocked(ap, compatibility, routingAffinity); idle != nil {
|
||||
delete(ap.conns, idle.id)
|
||||
evicted = append(evicted, idle)
|
||||
p.metrics.scaleDownTotal.Add(1)
|
||||
} else if compatible == nil {
|
||||
} else if affine == nil {
|
||||
compatible := p.pickLeastBusyConnLocked(ap, "", compatibility)
|
||||
if compatible != nil {
|
||||
// Capacity is full and every compatible connection is busy. The
|
||||
// hint remains soft here: queue on a compatible connection below.
|
||||
goto acquireAtCapacity
|
||||
}
|
||||
hasConnection := false
|
||||
for _, conn := range ap.conns {
|
||||
if conn != nil {
|
||||
@@ -1068,6 +1100,17 @@ retryAcquire:
|
||||
ap = p.getOrCreateAccountPool(accountID)
|
||||
ap.mu.Lock()
|
||||
ap.creating--
|
||||
if ap.generation != acquireGeneration {
|
||||
ap.signalChangedLocked()
|
||||
ap.mu.Unlock()
|
||||
if conn != nil {
|
||||
conn.close()
|
||||
}
|
||||
if retry < 1 {
|
||||
return p.acquire(ctx, req, retry+1)
|
||||
}
|
||||
return nil, errOpenAIWSConnClosed
|
||||
}
|
||||
if dialErr != nil {
|
||||
ap.prewarmFails++
|
||||
ap.prewarmFailAt = time.Now()
|
||||
@@ -1075,20 +1118,26 @@ retryAcquire:
|
||||
ap.mu.Unlock()
|
||||
return nil, dialErr
|
||||
}
|
||||
// Claim the freshly dialed connection before publishing it. Otherwise a
|
||||
// topology waiter awakened below can take the free semaphore first and
|
||||
// make the caller that paid for the dial queue behind it.
|
||||
if !conn.tryAcquire() {
|
||||
ap.signalChangedLocked()
|
||||
ap.mu.Unlock()
|
||||
conn.close()
|
||||
return nil, errOpenAIWSConnClosed
|
||||
}
|
||||
ap.conns[conn.id] = conn
|
||||
ap.prewarmFails = 0
|
||||
ap.prewarmFailAt = time.Time{}
|
||||
// Wake acquires that observed creating>0 with no compatible connection.
|
||||
// Without this signal they can remain asleep until the new lease is
|
||||
// released, even though the pool topology already changed.
|
||||
ap.signalChangedLocked()
|
||||
ap.mu.Unlock()
|
||||
p.metrics.acquireCreateTotal.Add(1)
|
||||
|
||||
if !conn.tryAcquire() {
|
||||
if err := conn.acquire(ctx); err != nil {
|
||||
conn.close()
|
||||
p.evictConn(accountID, conn.id)
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
lease := &openAIWSConnLease{pool: p, accountID: accountID, conn: conn, connPick: connPick}
|
||||
p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req)
|
||||
p.ensureTargetIdleAsync(accountID)
|
||||
return lease, nil
|
||||
}
|
||||
@@ -1100,7 +1149,8 @@ retryAcquire:
|
||||
return nil, errOpenAIWSConnQueueFull
|
||||
}
|
||||
|
||||
target := p.pickLeastBusyConnLocked(ap, req.PreferredConnID, betaFeatures)
|
||||
acquireAtCapacity:
|
||||
target := p.pickLeastBusyConnLocked(ap, req.PreferredConnID, compatibility)
|
||||
connPick := time.Since(pickStartedAt)
|
||||
p.recordConnPickDuration(connPick)
|
||||
if target == nil {
|
||||
@@ -1142,6 +1192,7 @@ retryAcquire:
|
||||
p.metrics.acquireQueueWaitMs.Add(queueWait.Milliseconds())
|
||||
lease := &openAIWSConnLease{pool: p, accountID: accountID, conn: target, queueWait: queueWait, connPick: connPick, reused: true}
|
||||
p.metrics.acquireReuseTotal.Add(1)
|
||||
p.recordLastSuccessfulAcquire(accountID, acquireGeneration, req)
|
||||
p.ensureTargetIdleAsync(accountID)
|
||||
return lease, nil
|
||||
}
|
||||
@@ -1157,6 +1208,23 @@ func (p *openAIWSConnPool) recordConnPickDuration(duration time.Duration) {
|
||||
p.metrics.connPickMs.Add(duration.Milliseconds())
|
||||
}
|
||||
|
||||
func (p *openAIWSConnPool) recordLastSuccessfulAcquire(accountID int64, generation uint64, req openAIWSAcquireRequest) {
|
||||
if p == nil || accountID <= 0 {
|
||||
return
|
||||
}
|
||||
ap, ok := p.getAccountPool(accountID)
|
||||
if !ok || ap == nil {
|
||||
return
|
||||
}
|
||||
ap.mu.Lock()
|
||||
if ap.generation != generation {
|
||||
ap.mu.Unlock()
|
||||
return
|
||||
}
|
||||
ap.lastAcquire = cloneOpenAIWSAcquireRequestPtr(&req)
|
||||
ap.mu.Unlock()
|
||||
}
|
||||
|
||||
func (p *openAIWSConnPool) pickOldestIdleConnLocked(ap *openAIWSAccountPool) *openAIWSConn {
|
||||
if ap == nil || len(ap.conns) == 0 {
|
||||
return nil
|
||||
@@ -1173,13 +1241,19 @@ func (p *openAIWSConnPool) pickOldestIdleConnLocked(ap *openAIWSAccountPool) *op
|
||||
return oldest
|
||||
}
|
||||
|
||||
func (p *openAIWSConnPool) pickOldestIdleConnWithDifferentBetaFeaturesLocked(ap *openAIWSAccountPool, betaFeatures string) *openAIWSConn {
|
||||
func (p *openAIWSConnPool) pickOldestIdleConnWithoutHandshakeCompatibilityOrRoutingAffinityLocked(
|
||||
ap *openAIWSAccountPool,
|
||||
compatibility openAIWSHandshakeCompatibilityKey,
|
||||
routingAffinity string,
|
||||
) *openAIWSConn {
|
||||
if ap == nil || len(ap.conns) == 0 {
|
||||
return nil
|
||||
}
|
||||
var oldest *openAIWSConn
|
||||
for _, conn := range ap.conns {
|
||||
if conn == nil || conn.matchesBetaFeatures(betaFeatures) || conn.isLeased() || conn.waiters.Load() > 0 || p.isConnPinnedLocked(ap, conn.id) {
|
||||
if conn == nil ||
|
||||
(conn.matchesHandshakeCompatibility(compatibility) && conn.matchesRoutingAffinity(routingAffinity)) ||
|
||||
conn.isLeased() || conn.waiters.Load() > 0 || p.isConnPinnedLocked(ap, conn.id) {
|
||||
continue
|
||||
}
|
||||
if oldest == nil || conn.lastUsedAt().Before(oldest.lastUsedAt()) {
|
||||
@@ -1330,13 +1404,17 @@ func (p *openAIWSConnPool) cleanupAccountLocked(ap *openAIWSAccountPool, now tim
|
||||
return evicted
|
||||
}
|
||||
|
||||
func (p *openAIWSConnPool) pickLeastBusyConnLocked(ap *openAIWSAccountPool, preferredConnID, betaFeatures string) *openAIWSConn {
|
||||
func (p *openAIWSConnPool) pickLeastBusyConnLocked(
|
||||
ap *openAIWSAccountPool,
|
||||
preferredConnID string,
|
||||
compatibility openAIWSHandshakeCompatibilityKey,
|
||||
) *openAIWSConn {
|
||||
if ap == nil || len(ap.conns) == 0 {
|
||||
return nil
|
||||
}
|
||||
preferredConnID = stringsTrim(preferredConnID)
|
||||
if preferredConnID != "" {
|
||||
if conn, ok := ap.conns[preferredConnID]; ok && conn.matchesBetaFeatures(betaFeatures) {
|
||||
if conn, ok := ap.conns[preferredConnID]; ok && conn.matchesHandshakeCompatibility(compatibility) {
|
||||
return conn
|
||||
}
|
||||
}
|
||||
@@ -1344,7 +1422,37 @@ func (p *openAIWSConnPool) pickLeastBusyConnLocked(ap *openAIWSAccountPool, pref
|
||||
var bestWaiters int32
|
||||
var bestLastUsed time.Time
|
||||
for _, conn := range ap.conns {
|
||||
if conn == nil || !conn.matchesBetaFeatures(betaFeatures) {
|
||||
if conn == nil || !conn.matchesHandshakeCompatibility(compatibility) {
|
||||
continue
|
||||
}
|
||||
waiters := conn.waiters.Load()
|
||||
lastUsed := conn.lastUsedAt()
|
||||
if best == nil ||
|
||||
waiters < bestWaiters ||
|
||||
(waiters == bestWaiters && lastUsed.Before(bestLastUsed)) {
|
||||
best = conn
|
||||
bestWaiters = waiters
|
||||
bestLastUsed = lastUsed
|
||||
}
|
||||
}
|
||||
return best
|
||||
}
|
||||
|
||||
func (p *openAIWSConnPool) pickLeastBusyConnWithRoutingAffinityLocked(
|
||||
ap *openAIWSAccountPool,
|
||||
compatibility openAIWSHandshakeCompatibilityKey,
|
||||
routingAffinity string,
|
||||
) *openAIWSConn {
|
||||
if ap == nil || len(ap.conns) == 0 {
|
||||
return nil
|
||||
}
|
||||
var best *openAIWSConn
|
||||
var bestWaiters int32
|
||||
var bestLastUsed time.Time
|
||||
for _, conn := range ap.conns {
|
||||
if conn == nil ||
|
||||
!conn.matchesHandshakeCompatibility(compatibility) ||
|
||||
!conn.matchesRoutingAffinity(routingAffinity) {
|
||||
continue
|
||||
}
|
||||
waiters := conn.waiters.Load()
|
||||
@@ -1488,12 +1596,20 @@ func (p *openAIWSConnPool) prewarmConns(accountID int64, req openAIWSAcquireRequ
|
||||
if len(generations) > 0 {
|
||||
generation = generations[0]
|
||||
}
|
||||
staleTarget := false
|
||||
defer func() {
|
||||
if ap, ok := p.getAccountPool(accountID); ok && ap != nil {
|
||||
ap.mu.Lock()
|
||||
ap.prewarmActive = false
|
||||
ap.signalChangedLocked()
|
||||
ap.mu.Unlock()
|
||||
}
|
||||
if staleTarget {
|
||||
// A newer acquire arrived while the old dial was in flight. Re-run
|
||||
// target selection only after clearing prewarmActive so the latest
|
||||
// beta/hint target can fill the idle budget.
|
||||
p.ensureTargetIdleAsync(accountID)
|
||||
}
|
||||
}()
|
||||
|
||||
for i := 0; i < total; i++ {
|
||||
@@ -1524,6 +1640,13 @@ func (p *openAIWSConnPool) prewarmConns(accountID int64, req openAIWSAcquireRequ
|
||||
conn.close()
|
||||
continue
|
||||
}
|
||||
if !sameOpenAIWSPrewarmTarget(req, *ap.lastAcquire) {
|
||||
staleTarget = true
|
||||
ap.signalChangedLocked()
|
||||
ap.mu.Unlock()
|
||||
conn.close()
|
||||
continue
|
||||
}
|
||||
if len(ap.conns) >= p.effectiveMaxConnsByAccount(req.Account) {
|
||||
ap.signalChangedLocked()
|
||||
ap.mu.Unlock()
|
||||
@@ -1563,6 +1686,7 @@ func (p *openAIWSConnPool) ClearAccount(accountID int64) {
|
||||
ap.prewarmUntil = time.Time{}
|
||||
ap.prewarmFails = 0
|
||||
ap.prewarmFailAt = time.Time{}
|
||||
ap.signalChangedLocked()
|
||||
ap.mu.Unlock()
|
||||
closeOpenAIWSConns(conns)
|
||||
}
|
||||
@@ -1676,7 +1800,8 @@ func (p *openAIWSConnPool) dialConn(ctx context.Context, req openAIWSAcquireRequ
|
||||
}
|
||||
id := p.nextConnID(req.Account.ID)
|
||||
pooledConn := newOpenAIWSConn(id, req.Account.ID, conn, handshakeHeaders)
|
||||
pooledConn.betaFeatures = normalizeOpenAIWSBetaFeatures(req.Headers)
|
||||
pooledConn.handshakeCompatibility = normalizeOpenAIWSHandshakeCompatibility(req.Headers)
|
||||
pooledConn.routingAffinity = normalizeOpenAIWSRoutingAffinity(req.Headers)
|
||||
return pooledConn, nil
|
||||
}
|
||||
|
||||
@@ -1855,6 +1980,12 @@ func cloneOpenAIWSAcquireRequestPtr(req *openAIWSAcquireRequest) *openAIWSAcquir
|
||||
return &copied
|
||||
}
|
||||
|
||||
func sameOpenAIWSPrewarmTarget(a, b openAIWSAcquireRequest) bool {
|
||||
return stringsTrim(a.WSURL) == stringsTrim(b.WSURL) &&
|
||||
stringsTrim(a.ProxyURL) == stringsTrim(b.ProxyURL) &&
|
||||
normalizeOpenAIWSHandshakeCompatibility(a.Headers) == normalizeOpenAIWSHandshakeCompatibility(b.Headers)
|
||||
}
|
||||
|
||||
func normalizeOpenAIWSBetaFeatures(headers http.Header) string {
|
||||
features := make(map[string]struct{})
|
||||
for name, values := range headers {
|
||||
@@ -1880,6 +2011,39 @@ func normalizeOpenAIWSBetaFeatures(headers http.Header) string {
|
||||
return strings.Join(normalized, ",")
|
||||
}
|
||||
|
||||
func normalizeOpenAIWSHandshakeCompatibility(headers http.Header) openAIWSHandshakeCompatibilityKey {
|
||||
return openAIWSHandshakeCompatibilityKey{
|
||||
betaFeatures: normalizeOpenAIWSBetaFeatures(headers),
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeOpenAIWSRoutingAffinity(headers http.Header) string {
|
||||
canonicalName := http.CanonicalHeaderKey(openAICodexRoutingHintHeader)
|
||||
if values, ok := headers[canonicalName]; ok {
|
||||
for _, value := range values {
|
||||
if value = strings.TrimSpace(value); value != "" {
|
||||
return value
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
variantNames := make([]string, 0)
|
||||
for name := range headers {
|
||||
if name != canonicalName && strings.EqualFold(strings.TrimSpace(name), openAICodexRoutingHintHeader) {
|
||||
variantNames = append(variantNames, name)
|
||||
}
|
||||
}
|
||||
sort.Strings(variantNames)
|
||||
for _, name := range variantNames {
|
||||
for _, value := range headers[name] {
|
||||
if value = strings.TrimSpace(value); value != "" {
|
||||
return value
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func cloneHeader(src http.Header) http.Header {
|
||||
if src == nil {
|
||||
return nil
|
||||
|
||||
@@ -61,6 +61,18 @@ func TestOpenAIWSConnPool_AcquireCleanupInterval(t *testing.T) {
|
||||
require.Less(t, openAIWSAcquireCleanupInterval, openAIWSBackgroundSweepTicker)
|
||||
}
|
||||
|
||||
func TestNormalizeOpenAIWSRoutingAffinityPrefersCanonicalAndSortsVariants(t *testing.T) {
|
||||
headers := http.Header{
|
||||
"X-CODEX-ROUTING-HINT": []string{" variant-uppercase "},
|
||||
"X-Codex-Routing-Hint": []string{" ", " canonical "},
|
||||
}
|
||||
|
||||
require.Equal(t, "canonical", normalizeOpenAIWSRoutingAffinity(headers))
|
||||
|
||||
delete(headers, "X-Codex-Routing-Hint")
|
||||
require.Equal(t, "variant-uppercase", normalizeOpenAIWSRoutingAffinity(headers))
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnLease_WriteJSONAndGuards(t *testing.T) {
|
||||
conn := newOpenAIWSConn("lease_write", 1, &openAIWSFakeConn{}, nil)
|
||||
lease := &openAIWSConnLease{conn: conn}
|
||||
@@ -310,6 +322,219 @@ func TestOpenAIWSConnPool_AcquireQueueWaitMetrics(t *testing.T) {
|
||||
require.GreaterOrEqual(t, metrics.ConnPickTotal, int64(1))
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnPool_DialSuccessWakesTopologyWaiterAndCanceledWaiterDoesNotLoseLease(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
||||
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
||||
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
|
||||
cfg.Gateway.OpenAIWS.QueueLimitPerConn = 4
|
||||
|
||||
pool := newOpenAIWSConnPool(cfg)
|
||||
dialer := newOpenAIWSFirstDialBlockingCaptureDialer()
|
||||
pool.setClientDialerForTest(dialer)
|
||||
account := &Account{ID: 991, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||||
req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
|
||||
|
||||
type result struct {
|
||||
lease *openAIWSConnLease
|
||||
err error
|
||||
}
|
||||
firstCh := make(chan result, 1)
|
||||
go func() {
|
||||
lease, err := pool.Acquire(context.Background(), req)
|
||||
firstCh <- result{lease: lease, err: err}
|
||||
}()
|
||||
<-dialer.firstStarted
|
||||
|
||||
waitCtx, cancelWait := context.WithCancel(context.Background())
|
||||
secondCh := make(chan result, 1)
|
||||
waitReq := req
|
||||
waitReq.Headers = http.Header{openAICodexRoutingHintHeader: {"model=gpt-5.6-codex;tier=priority"}}
|
||||
go func() {
|
||||
lease, err := pool.Acquire(waitCtx, waitReq)
|
||||
secondCh <- result{lease: lease, err: err}
|
||||
}()
|
||||
|
||||
close(dialer.releaseFirst)
|
||||
first := <-firstCh
|
||||
require.NoError(t, first.err)
|
||||
require.NotNil(t, first.lease)
|
||||
|
||||
// The second acquire initially waits on the account topology channel while
|
||||
// the first dial is in flight. Dial success must wake it immediately so it
|
||||
// can queue on the newly-created (still leased) connection.
|
||||
require.Eventually(t, func() bool {
|
||||
ap, ok := pool.getAccountPool(account.ID)
|
||||
if !ok || ap == nil {
|
||||
return false
|
||||
}
|
||||
ap.mu.Lock()
|
||||
defer ap.mu.Unlock()
|
||||
for _, conn := range ap.conns {
|
||||
if conn != nil && conn.waiters.Load() == 1 {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}, time.Second, 5*time.Millisecond)
|
||||
|
||||
cancelWait()
|
||||
second := <-secondCh
|
||||
require.ErrorIs(t, second.err, context.Canceled)
|
||||
require.Nil(t, second.lease)
|
||||
ap, ok := pool.getAccountPool(account.ID)
|
||||
require.True(t, ok)
|
||||
ap.mu.Lock()
|
||||
require.NotNil(t, ap.lastAcquire)
|
||||
require.Empty(t, normalizeOpenAIWSRoutingAffinity(ap.lastAcquire.Headers), "a canceled acquire must not replace the successful prewarm target")
|
||||
ap.mu.Unlock()
|
||||
first.lease.Release()
|
||||
|
||||
third, err := pool.Acquire(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.True(t, third.Reused(), "a canceled waiter must not consume the released semaphore token")
|
||||
require.Equal(t, first.lease.ConnID(), third.ConnID())
|
||||
third.Release()
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnPool_PrewarmHintChangeDoesNotInvalidateHealthyDial(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
||||
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 1
|
||||
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
|
||||
cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 2
|
||||
|
||||
pool := newOpenAIWSConnPool(cfg)
|
||||
dialer := newOpenAIWSFirstDialBlockingCaptureDialer()
|
||||
pool.setClientDialerForTest(dialer)
|
||||
account := &Account{ID: 992, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||||
oldHeaders := make(http.Header)
|
||||
oldHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex")
|
||||
newHeaders := make(http.Header)
|
||||
newHeaders.Set(openAICodexRoutingHintHeader, "model=gpt-5.6-codex;tier=priority")
|
||||
ap := pool.getOrCreateAccountPool(account.ID)
|
||||
ap.mu.Lock()
|
||||
ap.lastAcquire = &openAIWSAcquireRequest{
|
||||
Account: account,
|
||||
WSURL: "wss://example.com/v1/responses",
|
||||
Headers: oldHeaders,
|
||||
}
|
||||
ap.mu.Unlock()
|
||||
|
||||
pool.ensureTargetIdleAsync(account.ID)
|
||||
<-dialer.firstStarted
|
||||
|
||||
// Simulate a newer priority target arriving while the old model-only
|
||||
// prewarm dial is in flight. Routing hints are advisory, so this alone must
|
||||
// not discard an otherwise compatible connection.
|
||||
ap.mu.Lock()
|
||||
ap.lastAcquire = &openAIWSAcquireRequest{
|
||||
Account: account,
|
||||
WSURL: "wss://example.com/v1/responses",
|
||||
Headers: newHeaders,
|
||||
}
|
||||
ap.mu.Unlock()
|
||||
close(dialer.releaseFirst)
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
ap.mu.Lock()
|
||||
defer ap.mu.Unlock()
|
||||
if ap.prewarmActive || len(ap.conns) != 1 {
|
||||
return false
|
||||
}
|
||||
for _, conn := range ap.conns {
|
||||
return conn != nil && conn.routingAffinity == "model=gpt-5.6-codex"
|
||||
}
|
||||
return false
|
||||
}, 2*time.Second, 10*time.Millisecond)
|
||||
require.Equal(t, 1, dialer.DialCount(), "routing-hint-only changes must not turn advisory metadata into hard reconnects")
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnPool_ClearAccountWakesIncompatibleTopologyWaiter(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
||||
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
||||
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
|
||||
|
||||
pool := newOpenAIWSConnPool(cfg)
|
||||
dialer := &openAIWSCountingDialer{}
|
||||
pool.setClientDialerForTest(dialer)
|
||||
account := &Account{ID: 993, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||||
baseReq := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
|
||||
betaAReq := baseReq
|
||||
betaAReq.Headers = http.Header{"X-Codex-Beta-Features": {"feature_a"}}
|
||||
betaBReq := baseReq
|
||||
betaBReq.Headers = http.Header{"X-Codex-Beta-Features": {"feature_b"}}
|
||||
|
||||
busy, err := pool.Acquire(context.Background(), betaAReq)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, dialer.DialCount())
|
||||
|
||||
type result struct {
|
||||
lease *openAIWSConnLease
|
||||
err error
|
||||
}
|
||||
resultCh := make(chan result, 1)
|
||||
waitCtx, cancelWait := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancelWait()
|
||||
go func() {
|
||||
lease, acquireErr := pool.Acquire(waitCtx, betaBReq)
|
||||
resultCh <- result{lease: lease, err: acquireErr}
|
||||
}()
|
||||
|
||||
require.Never(t, func() bool { return dialer.DialCount() > 1 }, 50*time.Millisecond, 5*time.Millisecond)
|
||||
pool.ClearAccount(account.ID)
|
||||
|
||||
resultB := <-resultCh
|
||||
require.NoError(t, resultB.err)
|
||||
require.NotNil(t, resultB.lease)
|
||||
require.False(t, resultB.lease.Reused())
|
||||
require.Equal(t, 2, dialer.DialCount(), "ClearAccount must wake the waiter to redial immediately")
|
||||
resultB.lease.Release()
|
||||
busy.Release()
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnPool_ClearAccountDoesNotReviveInFlightDialGeneration(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 1
|
||||
cfg.Gateway.OpenAIWS.MinIdlePerAccount = 0
|
||||
cfg.Gateway.OpenAIWS.MaxIdlePerAccount = 1
|
||||
|
||||
pool := newOpenAIWSConnPool(cfg)
|
||||
dialer := newOpenAIWSFirstDialBlockingCaptureDialer()
|
||||
pool.setClientDialerForTest(dialer)
|
||||
account := &Account{ID: 994, Platform: PlatformOpenAI, Type: AccountTypeOAuth}
|
||||
req := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"}
|
||||
|
||||
type result struct {
|
||||
lease *openAIWSConnLease
|
||||
err error
|
||||
}
|
||||
resultCh := make(chan result, 1)
|
||||
go func() {
|
||||
lease, err := pool.Acquire(context.Background(), req)
|
||||
resultCh <- result{lease: lease, err: err}
|
||||
}()
|
||||
<-dialer.firstStarted
|
||||
|
||||
pool.ClearAccount(account.ID)
|
||||
close(dialer.releaseFirst)
|
||||
got := <-resultCh
|
||||
require.NoError(t, got.err)
|
||||
require.NotNil(t, got.lease)
|
||||
require.Equal(t, 2, dialer.DialCount(), "the pre-clear dial must be discarded and retried in the new generation")
|
||||
require.True(t, strings.HasSuffix(got.lease.ConnID(), "_2"))
|
||||
|
||||
ap, ok := pool.getAccountPool(account.ID)
|
||||
require.True(t, ok)
|
||||
ap.mu.Lock()
|
||||
require.Equal(t, uint64(1), ap.generation)
|
||||
require.Len(t, ap.conns, 1)
|
||||
require.NotNil(t, ap.lastAcquire, "only the post-clear successful acquire may restore the prewarm target")
|
||||
ap.mu.Unlock()
|
||||
got.lease.Release()
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnPool_ForceNewConnSkipsReuse(t *testing.T) {
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2
|
||||
@@ -1468,6 +1693,21 @@ func TestOpenAIWSConn_AdditionalGuardBranches(t *testing.T) {
|
||||
closeOpenAIWSConns([]*openAIWSConn{nil, connOK})
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnPool_CanceledWaiterReturnsDeliveredLease(t *testing.T) {
|
||||
conn := newOpenAIWSConn("cancelled_delivery", 1, &openAIWSFakeConn{}, nil)
|
||||
|
||||
// Both branches of acquire's select are ready. Before the post-delivery
|
||||
// cancellation check this intermittently returned nil after consuming the
|
||||
// only lease token, which made the next pool acquire block forever.
|
||||
for range 64 {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
require.ErrorIs(t, conn.acquire(ctx), context.Canceled)
|
||||
require.True(t, conn.tryAcquire(), "a canceled waiter must return a delivered lease token")
|
||||
conn.release()
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIWSConnLease_MarkBrokenEvictsConn(t *testing.T) {
|
||||
pool := newOpenAIWSConnPool(&config.Config{})
|
||||
accountID := int64(5001)
|
||||
@@ -1602,6 +1842,21 @@ type openAIWSCountingDialer struct {
|
||||
dialCount int
|
||||
}
|
||||
|
||||
type openAIWSFirstDialBlockingCaptureDialer struct {
|
||||
mu sync.Mutex
|
||||
dialCount int
|
||||
headers []http.Header
|
||||
firstStarted chan struct{}
|
||||
releaseFirst chan struct{}
|
||||
}
|
||||
|
||||
func newOpenAIWSFirstDialBlockingCaptureDialer() *openAIWSFirstDialBlockingCaptureDialer {
|
||||
return &openAIWSFirstDialBlockingCaptureDialer{
|
||||
firstStarted: make(chan struct{}),
|
||||
releaseFirst: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
type openAIWSAlwaysFailDialer struct {
|
||||
mu sync.Mutex
|
||||
dialCount int
|
||||
@@ -1671,6 +1926,36 @@ func (d *openAIWSCountingDialer) Dial(
|
||||
return &openAIWSFakeConn{}, 0, nil, nil
|
||||
}
|
||||
|
||||
func (d *openAIWSFirstDialBlockingCaptureDialer) Dial(
|
||||
ctx context.Context,
|
||||
wsURL string,
|
||||
headers http.Header,
|
||||
proxyURL string,
|
||||
) (openAIWSClientConn, int, http.Header, error) {
|
||||
_ = wsURL
|
||||
_ = proxyURL
|
||||
d.mu.Lock()
|
||||
d.dialCount++
|
||||
dialNumber := d.dialCount
|
||||
d.headers = append(d.headers, cloneHeader(headers))
|
||||
d.mu.Unlock()
|
||||
if dialNumber == 1 {
|
||||
close(d.firstStarted)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, 0, nil, ctx.Err()
|
||||
case <-d.releaseFirst:
|
||||
}
|
||||
}
|
||||
return &openAIWSFakeConn{}, 0, nil, nil
|
||||
}
|
||||
|
||||
func (d *openAIWSFirstDialBlockingCaptureDialer) DialCount() int {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
return d.dialCount
|
||||
}
|
||||
|
||||
func (d *openAIWSCountingDialer) DialCount() int {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
|
||||
Loaded 100 of 141 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user