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:
IanShaw027 committed 2026-08-08 11:12:56 +08:00
commit 1f58e25ab3
141 files changed
+5942 -680

No files matched your search

+1 -1
View File
@@ -1 +1 @@
0.1.171
0.1.172
+16 -14
View File
@@ -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
View File
@@ -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
+34 -30
View File
@@ -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()
+11
View File
@@ -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
View File
@@ -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))
+18
View File
@@ -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()
+105
View File
@@ -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))
+161
View File
@@ -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) {
+114
View File
@@ -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)
}
+6
View File
@@ -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
}
+48 -26
View File
@@ -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",
+2
View File
@@ -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) {
+4
View File
@@ -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"`
+1 -1
View File
@@ -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
+6
View File
@@ -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"
+20 -1
View File
@@ -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)
+2 -2
View File
@@ -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")
}
+5
View File
@@ -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 {
+23 -8
View File
@@ -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
+12 -9
View File
@@ -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
}
+10 -6
View File
@@ -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())
+1 -1
View File
@@ -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)
+192 -28
View File
@@ -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