diff --git a/README.md b/README.md index 661cd8e7eb..3e43fd3ae4 100644 --- a/README.md +++ b/README.md @@ -637,19 +637,19 @@ Simple Mode is designed for individual developers or internal teams who want qui --- -## Grok / xAI OAuth Support +## Grok / xAI Support -Sub2API supports Grok subscription accounts through xAI OAuth and forwards OpenAI-compatible Responses traffic to xAI. +Sub2API supports both Grok subscription accounts through xAI OAuth and standard xAI API-key accounts. Both account types forward OpenAI-compatible Responses traffic to xAI. ### Supported Scope - Platform name: `grok` -- Account type: OAuth subscription accounts -- Public Responses targets: `/v1/responses`, `/responses`, and `/backend-api/codex/responses`, forwarded to `${XAI_BASE_URL:-https://api.x.ai/v1}/responses` +- Account types: OAuth subscription accounts and xAI API-key accounts +- Public Responses targets: `/v1/responses`, `/responses`, and `/backend-api/codex/responses`, forwarded to the Grok subscription proxy for OAuth accounts or `https://api.x.ai/v1/responses` for API-key accounts - Public Claude-compatible target: `/v1/messages`, converted to xAI Responses and returned as Anthropic Messages output for Claude CLI style clients -- Public Chat Completions targets: `/v1/chat/completions` and `/chat/completions`, forwarded to `${XAI_BASE_URL:-https://api.x.ai/v1}/chat/completions` +- Public Chat Completions targets: `/v1/chat/completions` and `/chat/completions`, forwarded to the account-type-specific xAI upstream - Codex CLI style Responses WebSocket ingress is accepted on the Responses targets and bridged to xAI HTTP/SSE Responses upstream -- Initial text models: `grok-4.3`, `grok-build-0.1`, `grok-4.20-0309-reasoning`, `grok-4.20-0309-non-reasoning`, and `grok-4.20-multi-agent-0309` +- Text models: `grok-4.5`, `grok-4.3`, `grok-build-0.1`, `grok-composer-2.5-fast`, `grok-4.20-0309-reasoning`, `grok-4.20-0309-non-reasoning`, and `grok-4.20-multi-agent-0309` - Media targets for Grok groups: `/v1/images/generations`, `/images/generations`, `/v1/images/edits`, `/images/edits`, `/v1/videos/generations`, `/videos/generations`, `/v1/videos/{request_id}`, and `/videos/{request_id}`. Generation requests require the group image-generation permission. - Media models: `grok-imagine`, `grok-imagine-image-quality`, `grok-imagine-image`, `grok-imagine-edit`, `grok-imagine-video`, and `grok-imagine-video-1.5` - Out of scope for this provider: TTS, transcription, browser automation, cookies, and Grok web scraping @@ -665,9 +665,10 @@ The Grok OAuth flow uses PKCE and does not require committing private secrets. T | `XAI_OAUTH_REDIRECT_URI` | `http://127.0.0.1:56121/callback` | | `XAI_OAUTH_AUTHORIZE_URL` | `https://auth.x.ai/oauth2/authorize` | | `XAI_OAUTH_TOKEN_URL` | `https://auth.x.ai/oauth2/token` | -| `XAI_BASE_URL` | `https://api.x.ai/v1` | +| `XAI_BASE_URL` | `https://api.x.ai/v1`; runtime-diagnostics override (account `base_url` controls request forwarding) | +| `XAI_GROK_CLI_VERSION` | `0.2.93`; optional override for the client identity sent to `cli-chat-proxy.grok.com` | -Administrators can create or reauthorize Grok accounts from the dashboard, or use the admin API: +Administrators can create Grok OAuth or API-key accounts from the dashboard. OAuth authorization and reauthorization are also available through the admin API: | Endpoint | Purpose | |----------|---------| @@ -676,13 +677,47 @@ Administrators can create or reauthorize Grok accounts from the dashboard, or us | `POST /api/v1/admin/grok/oauth/refresh-token` | Validate or refresh a Grok refresh token | | `POST /api/v1/admin/grok/accounts/:id/refresh` | Refresh an existing Grok account | -Credential storage reuses the existing account JSON fields: `access_token`, `refresh_token`, `token_type`, `expires_at`, optional `email`, optional `subscription_tier`, and `entitlement_status`. +OAuth credential storage reuses the existing account JSON fields: `access_token`, `refresh_token`, `token_type`, `expires_at`, `base_url`, optional `email`, optional `subscription_tier`, and `entitlement_status`. OAuth inference defaults to `https://cli-chat-proxy.grok.com/v1`; existing OAuth accounts that stored the old `https://api.x.ai/v1` default are redirected to the subscription proxy at runtime. Explicit custom upstreams remain unchanged. + +For API-key accounts, select **Grok → API Key** in the create-account dialog. The official base URL defaults to `https://api.x.ai/v1`; credentials use the existing `base_url` and `api_key` account fields. OAuth accounts continue to use the subscription flow above. + +### Grok Build CLI Configuration + +1. In the Sub2API admin dashboard, add either a `grok` OAuth account and complete xAI authorization, or add a Grok API-key account. +2. Create a Grok group, attach the account to it, then create a Sub2API API key assigned to that group. +3. In the user API-key page, click **Use Key** and select **Grok CLI**. The modal generates the correct file and base URL for macOS/Linux or Windows. It also provides an OpenCode configuration on the **OpenCode** tab. +4. If configuring manually, save the following as `~/.grok/config.toml` (Windows: `%USERPROFILE%\.grok\config.toml`): + +```toml +[models] +default = "sub2api-grok" +web_search = "sub2api-grok" + +[model."sub2api-grok"] +model = "grok-4.5" +base_url = "https://your-sub2api.example.com/v1" +name = "Grok 4.5 via Sub2API" +description = "Grok 4.5 through a Sub2API Grok group" +api_key = "sk-your-sub2api-key" +api_backend = "responses" +context_window = 1000000 +supports_backend_search = true +``` + +Back up an existing `config.toml` before merging the entry. The file contains a Sub2API API key, so keep it private and restrict its permissions where supported. Verify the effective configuration and make a smoke request: + +```bash +grok inspect +grok -p "Reply with sub2api-ok" -m sub2api-grok +``` + +The `base_url` above is the public Sub2API URL ending in `/v1`, not `api.x.ai` or the internal xAI OAuth proxy URL. ### Usage And Quota Display xAI quota is passive. Sub2API does not invent subscription quota values; it records whitelisted xAI rate-limit headers from successful or rate-limited upstream responses when xAI sends them. Before the first usable upstream response, the dashboard shows quota as unknown and still displays local Sub2API usage stats. -`401` responses mark the account as needing reauthorization. `403` responses are treated as entitlement or subscription-tier failures instead of token-refresh loops. `429` responses use `Retry-After` or a short cooldown to temporarily remove the account from scheduling. +`401` responses temporarily remove accounts with invalid credentials from scheduling. `403` responses are treated as access or entitlement failures instead of token-refresh loops. `429` responses use `Retry-After` or a short cooldown to temporarily remove the account from scheduling. --- diff --git a/backend/cmd/server/VERSION b/backend/cmd/server/VERSION index 6c015fd0ca..611234586a 100644 --- a/backend/cmd/server/VERSION +++ b/backend/cmd/server/VERSION @@ -1 +1 @@ -0.1.151 +0.1.152 diff --git a/backend/ent/group.go b/backend/ent/group.go index 5bec594977..088da069d6 100644 --- a/backend/ent/group.go +++ b/backend/ent/group.go @@ -83,6 +83,8 @@ type Group struct { VideoPrice720p *float64 `json:"video_price_720p,omitempty"` // VideoPrice1080p holds the value of the "video_price_1080p" field. VideoPrice1080p *float64 `json:"video_price_1080p,omitempty"` + // Codex alpha/search 网页搜索单次价格（USD/次）；nil 表示使用默认价 0.01（官方 $10/1000 次） + WebSearchPricePerCall *float64 `json:"web_search_price_per_call,omitempty"` // 是否仅允许 Claude Code 客户端 ClaudeCodeOnly bool `json:"claude_code_only,omitempty"` // 非 Claude Code 请求降级使用的分组 ID @@ -223,7 +225,7 @@ func (*Group) scanValues(columns []string) ([]any, error) { values[i] = new([]byte) case group.FieldPeakRateEnabled, group.FieldIsExclusive, group.FieldAllowImageGeneration, group.FieldAllowBatchImageGeneration, group.FieldImageRateIndependent, group.FieldVideoRateIndependent, group.FieldClaudeCodeOnly, group.FieldModelRoutingEnabled, group.FieldMcpXMLInject, group.FieldAllowMessagesDispatch, group.FieldRequireOauthOnly, group.FieldRequirePrivacySet: values[i] = new(sql.NullBool) - case group.FieldRateMultiplier, group.FieldPeakRateMultiplier, group.FieldDailyLimitUsd, group.FieldWeeklyLimitUsd, group.FieldMonthlyLimitUsd, group.FieldImageRateMultiplier, group.FieldImagePrice1k, group.FieldImagePrice2k, group.FieldImagePrice4k, group.FieldBatchImageDiscountMultiplier, group.FieldBatchImageHoldMultiplier, group.FieldVideoRateMultiplier, group.FieldVideoPrice480p, group.FieldVideoPrice720p, group.FieldVideoPrice1080p: + case group.FieldRateMultiplier, group.FieldPeakRateMultiplier, group.FieldDailyLimitUsd, group.FieldWeeklyLimitUsd, group.FieldMonthlyLimitUsd, group.FieldImageRateMultiplier, group.FieldImagePrice1k, group.FieldImagePrice2k, group.FieldImagePrice4k, group.FieldBatchImageDiscountMultiplier, group.FieldBatchImageHoldMultiplier, group.FieldVideoRateMultiplier, group.FieldVideoPrice480p, group.FieldVideoPrice720p, group.FieldVideoPrice1080p, group.FieldWebSearchPricePerCall: values[i] = new(sql.NullFloat64) case group.FieldID, group.FieldDefaultValidityDays, group.FieldFallbackGroupID, group.FieldFallbackGroupIDOnInvalidRequest, group.FieldSortOrder, group.FieldRpmLimit: values[i] = new(sql.NullInt64) @@ -455,6 +457,13 @@ func (_m *Group) assignValues(columns []string, values []any) error { _m.VideoPrice1080p = new(float64) *_m.VideoPrice1080p = value.Float64 } + case group.FieldWebSearchPricePerCall: + if value, ok := values[i].(*sql.NullFloat64); !ok { + return fmt.Errorf("unexpected type %T for field web_search_price_per_call", values[i]) + } else if value.Valid { + _m.WebSearchPricePerCall = new(float64) + *_m.WebSearchPricePerCall = value.Float64 + } case group.FieldClaudeCodeOnly: if value, ok := values[i].(*sql.NullBool); !ok { return fmt.Errorf("unexpected type %T for field claude_code_only", values[i]) @@ -749,6 +758,11 @@ func (_m *Group) String() string { builder.WriteString(fmt.Sprintf("%v", *v)) } builder.WriteString(", ") + if v := _m.WebSearchPricePerCall; v != nil { + builder.WriteString("web_search_price_per_call=") + builder.WriteString(fmt.Sprintf("%v", *v)) + } + builder.WriteString(", ") builder.WriteString("claude_code_only=") builder.WriteString(fmt.Sprintf("%v", _m.ClaudeCodeOnly)) builder.WriteString(", ") diff --git a/backend/ent/group/group.go b/backend/ent/group/group.go index 769c63e6b1..61d7a21d67 100644 --- a/backend/ent/group/group.go +++ b/backend/ent/group/group.go @@ -80,6 +80,8 @@ const ( FieldVideoPrice720p = "video_price_720p" // FieldVideoPrice1080p holds the string denoting the video_price_1080p field in the database. FieldVideoPrice1080p = "video_price_1080p" + // FieldWebSearchPricePerCall holds the string denoting the web_search_price_per_call field in the database. + FieldWebSearchPricePerCall = "web_search_price_per_call" // FieldClaudeCodeOnly holds the string denoting the claude_code_only field in the database. FieldClaudeCodeOnly = "claude_code_only" // FieldFallbackGroupID holds the string denoting the fallback_group_id field in the database. @@ -217,6 +219,7 @@ var Columns = []string{ FieldVideoPrice480p, FieldVideoPrice720p, FieldVideoPrice1080p, + FieldWebSearchPricePerCall, FieldClaudeCodeOnly, FieldFallbackGroupID, FieldFallbackGroupIDOnInvalidRequest, @@ -511,6 +514,11 @@ func ByVideoPrice1080p(opts ...sql.OrderTermOption) OrderOption { return sql.OrderByField(FieldVideoPrice1080p, opts...).ToFunc() } +// ByWebSearchPricePerCall orders the results by the web_search_price_per_call field. +func ByWebSearchPricePerCall(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldWebSearchPricePerCall, opts...).ToFunc() +} + // ByClaudeCodeOnly orders the results by the claude_code_only field. func ByClaudeCodeOnly(opts ...sql.OrderTermOption) OrderOption { return sql.OrderByField(FieldClaudeCodeOnly, opts...).ToFunc() diff --git a/backend/ent/group/where.go b/backend/ent/group/where.go index 5a9d92d0f4..b9a52a2eb6 100644 --- a/backend/ent/group/where.go +++ b/backend/ent/group/where.go @@ -215,6 +215,11 @@ func VideoPrice1080p(v float64) predicate.Group { return predicate.Group(sql.FieldEQ(FieldVideoPrice1080p, v)) } +// WebSearchPricePerCall applies equality check predicate on the "web_search_price_per_call" field. It's identical to WebSearchPricePerCallEQ. +func WebSearchPricePerCall(v float64) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldWebSearchPricePerCall, v)) +} + // ClaudeCodeOnly applies equality check predicate on the "claude_code_only" field. It's identical to ClaudeCodeOnlyEQ. func ClaudeCodeOnly(v bool) predicate.Group { return predicate.Group(sql.FieldEQ(FieldClaudeCodeOnly, v)) @@ -1655,6 +1660,56 @@ func VideoPrice1080pNotNil() predicate.Group { return predicate.Group(sql.FieldNotNull(FieldVideoPrice1080p)) } +// WebSearchPricePerCallEQ applies the EQ predicate on the "web_search_price_per_call" field. +func WebSearchPricePerCallEQ(v float64) predicate.Group { + return predicate.Group(sql.FieldEQ(FieldWebSearchPricePerCall, v)) +} + +// WebSearchPricePerCallNEQ applies the NEQ predicate on the "web_search_price_per_call" field. +func WebSearchPricePerCallNEQ(v float64) predicate.Group { + return predicate.Group(sql.FieldNEQ(FieldWebSearchPricePerCall, v)) +} + +// WebSearchPricePerCallIn applies the In predicate on the "web_search_price_per_call" field. +func WebSearchPricePerCallIn(vs ...float64) predicate.Group { + return predicate.Group(sql.FieldIn(FieldWebSearchPricePerCall, vs...)) +} + +// WebSearchPricePerCallNotIn applies the NotIn predicate on the "web_search_price_per_call" field. +func WebSearchPricePerCallNotIn(vs ...float64) predicate.Group { + return predicate.Group(sql.FieldNotIn(FieldWebSearchPricePerCall, vs...)) +} + +// WebSearchPricePerCallGT applies the GT predicate on the "web_search_price_per_call" field. +func WebSearchPricePerCallGT(v float64) predicate.Group { + return predicate.Group(sql.FieldGT(FieldWebSearchPricePerCall, v)) +} + +// WebSearchPricePerCallGTE applies the GTE predicate on the "web_search_price_per_call" field. +func WebSearchPricePerCallGTE(v float64) predicate.Group { + return predicate.Group(sql.FieldGTE(FieldWebSearchPricePerCall, v)) +} + +// WebSearchPricePerCallLT applies the LT predicate on the "web_search_price_per_call" field. +func WebSearchPricePerCallLT(v float64) predicate.Group { + return predicate.Group(sql.FieldLT(FieldWebSearchPricePerCall, v)) +} + +// WebSearchPricePerCallLTE applies the LTE predicate on the "web_search_price_per_call" field. +func WebSearchPricePerCallLTE(v float64) predicate.Group { + return predicate.Group(sql.FieldLTE(FieldWebSearchPricePerCall, v)) +} + +// WebSearchPricePerCallIsNil applies the IsNil predicate on the "web_search_price_per_call" field. +func WebSearchPricePerCallIsNil() predicate.Group { + return predicate.Group(sql.FieldIsNull(FieldWebSearchPricePerCall)) +} + +// WebSearchPricePerCallNotNil applies the NotNil predicate on the "web_search_price_per_call" field. +func WebSearchPricePerCallNotNil() predicate.Group { + return predicate.Group(sql.FieldNotNull(FieldWebSearchPricePerCall)) +} + // ClaudeCodeOnlyEQ applies the EQ predicate on the "claude_code_only" field. func ClaudeCodeOnlyEQ(v bool) predicate.Group { return predicate.Group(sql.FieldEQ(FieldClaudeCodeOnly, v)) diff --git a/backend/ent/group_create.go b/backend/ent/group_create.go index 2a6c18e67d..53fd733c1f 100644 --- a/backend/ent/group_create.go +++ b/backend/ent/group_create.go @@ -469,6 +469,20 @@ func (_c *GroupCreate) SetNillableVideoPrice1080p(v *float64) *GroupCreate { return _c } +// SetWebSearchPricePerCall sets the "web_search_price_per_call" field. +func (_c *GroupCreate) SetWebSearchPricePerCall(v float64) *GroupCreate { + _c.mutation.SetWebSearchPricePerCall(v) + return _c +} + +// SetNillableWebSearchPricePerCall sets the "web_search_price_per_call" field if the given value is not nil. +func (_c *GroupCreate) SetNillableWebSearchPricePerCall(v *float64) *GroupCreate { + if v != nil { + _c.SetWebSearchPricePerCall(*v) + } + return _c +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (_c *GroupCreate) SetClaudeCodeOnly(v bool) *GroupCreate { _c.mutation.SetClaudeCodeOnly(v) @@ -1218,6 +1232,10 @@ func (_c *GroupCreate) createSpec() (*Group, *sqlgraph.CreateSpec) { _spec.SetField(group.FieldVideoPrice1080p, field.TypeFloat64, value) _node.VideoPrice1080p = &value } + if value, ok := _c.mutation.WebSearchPricePerCall(); ok { + _spec.SetField(group.FieldWebSearchPricePerCall, field.TypeFloat64, value) + _node.WebSearchPricePerCall = &value + } if value, ok := _c.mutation.ClaudeCodeOnly(); ok { _spec.SetField(group.FieldClaudeCodeOnly, field.TypeBool, value) _node.ClaudeCodeOnly = value @@ -1968,6 +1986,30 @@ func (u *GroupUpsert) ClearVideoPrice1080p() *GroupUpsert { return u } +// SetWebSearchPricePerCall sets the "web_search_price_per_call" field. +func (u *GroupUpsert) SetWebSearchPricePerCall(v float64) *GroupUpsert { + u.Set(group.FieldWebSearchPricePerCall, v) + return u +} + +// UpdateWebSearchPricePerCall sets the "web_search_price_per_call" field to the value that was provided on create. +func (u *GroupUpsert) UpdateWebSearchPricePerCall() *GroupUpsert { + u.SetExcluded(group.FieldWebSearchPricePerCall) + return u +} + +// AddWebSearchPricePerCall adds v to the "web_search_price_per_call" field. +func (u *GroupUpsert) AddWebSearchPricePerCall(v float64) *GroupUpsert { + u.Add(group.FieldWebSearchPricePerCall, v) + return u +} + +// ClearWebSearchPricePerCall clears the value of the "web_search_price_per_call" field. +func (u *GroupUpsert) ClearWebSearchPricePerCall() *GroupUpsert { + u.SetNull(group.FieldWebSearchPricePerCall) + return u +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (u *GroupUpsert) SetClaudeCodeOnly(v bool) *GroupUpsert { u.Set(group.FieldClaudeCodeOnly, v) @@ -2858,6 +2900,34 @@ func (u *GroupUpsertOne) ClearVideoPrice1080p() *GroupUpsertOne { }) } +// SetWebSearchPricePerCall sets the "web_search_price_per_call" field. +func (u *GroupUpsertOne) SetWebSearchPricePerCall(v float64) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.SetWebSearchPricePerCall(v) + }) +} + +// AddWebSearchPricePerCall adds v to the "web_search_price_per_call" field. +func (u *GroupUpsertOne) AddWebSearchPricePerCall(v float64) *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.AddWebSearchPricePerCall(v) + }) +} + +// UpdateWebSearchPricePerCall sets the "web_search_price_per_call" field to the value that was provided on create. +func (u *GroupUpsertOne) UpdateWebSearchPricePerCall() *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.UpdateWebSearchPricePerCall() + }) +} + +// ClearWebSearchPricePerCall clears the value of the "web_search_price_per_call" field. +func (u *GroupUpsertOne) ClearWebSearchPricePerCall() *GroupUpsertOne { + return u.Update(func(s *GroupUpsert) { + s.ClearWebSearchPricePerCall() + }) +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (u *GroupUpsertOne) SetClaudeCodeOnly(v bool) *GroupUpsertOne { return u.Update(func(s *GroupUpsert) { @@ -3951,6 +4021,34 @@ func (u *GroupUpsertBulk) ClearVideoPrice1080p() *GroupUpsertBulk { }) } +// SetWebSearchPricePerCall sets the "web_search_price_per_call" field. +func (u *GroupUpsertBulk) SetWebSearchPricePerCall(v float64) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.SetWebSearchPricePerCall(v) + }) +} + +// AddWebSearchPricePerCall adds v to the "web_search_price_per_call" field. +func (u *GroupUpsertBulk) AddWebSearchPricePerCall(v float64) *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.AddWebSearchPricePerCall(v) + }) +} + +// UpdateWebSearchPricePerCall sets the "web_search_price_per_call" field to the value that was provided on create. +func (u *GroupUpsertBulk) UpdateWebSearchPricePerCall() *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.UpdateWebSearchPricePerCall() + }) +} + +// ClearWebSearchPricePerCall clears the value of the "web_search_price_per_call" field. +func (u *GroupUpsertBulk) ClearWebSearchPricePerCall() *GroupUpsertBulk { + return u.Update(func(s *GroupUpsert) { + s.ClearWebSearchPricePerCall() + }) +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (u *GroupUpsertBulk) SetClaudeCodeOnly(v bool) *GroupUpsertBulk { return u.Update(func(s *GroupUpsert) { diff --git a/backend/ent/group_update.go b/backend/ent/group_update.go index 3bb18d3e1a..1e767a139b 100644 --- a/backend/ent/group_update.go +++ b/backend/ent/group_update.go @@ -640,6 +640,33 @@ func (_u *GroupUpdate) ClearVideoPrice1080p() *GroupUpdate { return _u } +// SetWebSearchPricePerCall sets the "web_search_price_per_call" field. +func (_u *GroupUpdate) SetWebSearchPricePerCall(v float64) *GroupUpdate { + _u.mutation.ResetWebSearchPricePerCall() + _u.mutation.SetWebSearchPricePerCall(v) + return _u +} + +// SetNillableWebSearchPricePerCall sets the "web_search_price_per_call" field if the given value is not nil. +func (_u *GroupUpdate) SetNillableWebSearchPricePerCall(v *float64) *GroupUpdate { + if v != nil { + _u.SetWebSearchPricePerCall(*v) + } + return _u +} + +// AddWebSearchPricePerCall adds value to the "web_search_price_per_call" field. +func (_u *GroupUpdate) AddWebSearchPricePerCall(v float64) *GroupUpdate { + _u.mutation.AddWebSearchPricePerCall(v) + return _u +} + +// ClearWebSearchPricePerCall clears the value of the "web_search_price_per_call" field. +func (_u *GroupUpdate) ClearWebSearchPricePerCall() *GroupUpdate { + _u.mutation.ClearWebSearchPricePerCall() + return _u +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (_u *GroupUpdate) SetClaudeCodeOnly(v bool) *GroupUpdate { _u.mutation.SetClaudeCodeOnly(v) @@ -1375,6 +1402,15 @@ func (_u *GroupUpdate) sqlSave(ctx context.Context) (_node int, err error) { if _u.mutation.VideoPrice1080pCleared() { _spec.ClearField(group.FieldVideoPrice1080p, field.TypeFloat64) } + if value, ok := _u.mutation.WebSearchPricePerCall(); ok { + _spec.SetField(group.FieldWebSearchPricePerCall, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedWebSearchPricePerCall(); ok { + _spec.AddField(group.FieldWebSearchPricePerCall, field.TypeFloat64, value) + } + if _u.mutation.WebSearchPricePerCallCleared() { + _spec.ClearField(group.FieldWebSearchPricePerCall, field.TypeFloat64) + } if value, ok := _u.mutation.ClaudeCodeOnly(); ok { _spec.SetField(group.FieldClaudeCodeOnly, field.TypeBool, value) } @@ -2364,6 +2400,33 @@ func (_u *GroupUpdateOne) ClearVideoPrice1080p() *GroupUpdateOne { return _u } +// SetWebSearchPricePerCall sets the "web_search_price_per_call" field. +func (_u *GroupUpdateOne) SetWebSearchPricePerCall(v float64) *GroupUpdateOne { + _u.mutation.ResetWebSearchPricePerCall() + _u.mutation.SetWebSearchPricePerCall(v) + return _u +} + +// SetNillableWebSearchPricePerCall sets the "web_search_price_per_call" field if the given value is not nil. +func (_u *GroupUpdateOne) SetNillableWebSearchPricePerCall(v *float64) *GroupUpdateOne { + if v != nil { + _u.SetWebSearchPricePerCall(*v) + } + return _u +} + +// AddWebSearchPricePerCall adds value to the "web_search_price_per_call" field. +func (_u *GroupUpdateOne) AddWebSearchPricePerCall(v float64) *GroupUpdateOne { + _u.mutation.AddWebSearchPricePerCall(v) + return _u +} + +// ClearWebSearchPricePerCall clears the value of the "web_search_price_per_call" field. +func (_u *GroupUpdateOne) ClearWebSearchPricePerCall() *GroupUpdateOne { + _u.mutation.ClearWebSearchPricePerCall() + return _u +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (_u *GroupUpdateOne) SetClaudeCodeOnly(v bool) *GroupUpdateOne { _u.mutation.SetClaudeCodeOnly(v) @@ -3129,6 +3192,15 @@ func (_u *GroupUpdateOne) sqlSave(ctx context.Context) (_node *Group, err error) if _u.mutation.VideoPrice1080pCleared() { _spec.ClearField(group.FieldVideoPrice1080p, field.TypeFloat64) } + if value, ok := _u.mutation.WebSearchPricePerCall(); ok { + _spec.SetField(group.FieldWebSearchPricePerCall, field.TypeFloat64, value) + } + if value, ok := _u.mutation.AddedWebSearchPricePerCall(); ok { + _spec.AddField(group.FieldWebSearchPricePerCall, field.TypeFloat64, value) + } + if _u.mutation.WebSearchPricePerCallCleared() { + _spec.ClearField(group.FieldWebSearchPricePerCall, field.TypeFloat64) + } if value, ok := _u.mutation.ClaudeCodeOnly(); ok { _spec.SetField(group.FieldClaudeCodeOnly, field.TypeBool, value) } diff --git a/backend/ent/migrate/schema.go b/backend/ent/migrate/schema.go index 1dc6c611c5..cac559d535 100644 --- a/backend/ent/migrate/schema.go +++ b/backend/ent/migrate/schema.go @@ -865,6 +865,7 @@ var ( {Name: "video_price_480p", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}}, {Name: "video_price_720p", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}}, {Name: "video_price_1080p", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}}, + {Name: "web_search_price_per_call", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}}, {Name: "claude_code_only", Type: field.TypeBool, Default: false}, {Name: "fallback_group_id", Type: field.TypeInt64, Nullable: true}, {Name: "fallback_group_id_on_invalid_request", Type: field.TypeInt64, Nullable: true}, @@ -915,7 +916,7 @@ var ( { Name: "group_sort_order", Unique: false, - Columns: []*schema.Column{GroupsColumns[40]}, + Columns: []*schema.Column{GroupsColumns[41]}, }, }, } diff --git a/backend/ent/mutation.go b/backend/ent/mutation.go index bbd3663b3a..fb35531878 100644 --- a/backend/ent/mutation.go +++ b/backend/ent/mutation.go @@ -20842,6 +20842,8 @@ type GroupMutation struct { addvideo_price_720p *float64 video_price_1080p *float64 addvideo_price_1080p *float64 + web_search_price_per_call *float64 + addweb_search_price_per_call *float64 claude_code_only *bool fallback_group_id *int64 addfallback_group_id *int64 @@ -22608,6 +22610,76 @@ func (m *GroupMutation) ResetVideoPrice1080p() { delete(m.clearedFields, group.FieldVideoPrice1080p) } +// SetWebSearchPricePerCall sets the "web_search_price_per_call" field. +func (m *GroupMutation) SetWebSearchPricePerCall(f float64) { + m.web_search_price_per_call = &f + m.addweb_search_price_per_call = nil +} + +// WebSearchPricePerCall returns the value of the "web_search_price_per_call" field in the mutation. +func (m *GroupMutation) WebSearchPricePerCall() (r float64, exists bool) { + v := m.web_search_price_per_call + if v == nil { + return + } + return *v, true +} + +// OldWebSearchPricePerCall returns the old "web_search_price_per_call" field's value of the Group entity. +// If the Group object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *GroupMutation) OldWebSearchPricePerCall(ctx context.Context) (v *float64, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldWebSearchPricePerCall is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldWebSearchPricePerCall requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldWebSearchPricePerCall: %w", err) + } + return oldValue.WebSearchPricePerCall, nil +} + +// AddWebSearchPricePerCall adds f to the "web_search_price_per_call" field. +func (m *GroupMutation) AddWebSearchPricePerCall(f float64) { + if m.addweb_search_price_per_call != nil { + *m.addweb_search_price_per_call += f + } else { + m.addweb_search_price_per_call = &f + } +} + +// AddedWebSearchPricePerCall returns the value that was added to the "web_search_price_per_call" field in this mutation. +func (m *GroupMutation) AddedWebSearchPricePerCall() (r float64, exists bool) { + v := m.addweb_search_price_per_call + if v == nil { + return + } + return *v, true +} + +// ClearWebSearchPricePerCall clears the value of the "web_search_price_per_call" field. +func (m *GroupMutation) ClearWebSearchPricePerCall() { + m.web_search_price_per_call = nil + m.addweb_search_price_per_call = nil + m.clearedFields[group.FieldWebSearchPricePerCall] = struct{}{} +} + +// WebSearchPricePerCallCleared returns if the "web_search_price_per_call" field was cleared in this mutation. +func (m *GroupMutation) WebSearchPricePerCallCleared() bool { + _, ok := m.clearedFields[group.FieldWebSearchPricePerCall] + return ok +} + +// ResetWebSearchPricePerCall resets all changes to the "web_search_price_per_call" field. +func (m *GroupMutation) ResetWebSearchPricePerCall() { + m.web_search_price_per_call = nil + m.addweb_search_price_per_call = nil + delete(m.clearedFields, group.FieldWebSearchPricePerCall) +} + // SetClaudeCodeOnly sets the "claude_code_only" field. func (m *GroupMutation) SetClaudeCodeOnly(b bool) { m.claude_code_only = &b @@ -23642,7 +23714,7 @@ func (m *GroupMutation) Type() string { // order to get all numeric fields that were incremented/decremented, call // AddedFields(). func (m *GroupMutation) Fields() []string { - fields := make([]string, 0, 47) + fields := make([]string, 0, 48) if m.created_at != nil { fields = append(fields, group.FieldCreatedAt) } @@ -23739,6 +23811,9 @@ func (m *GroupMutation) Fields() []string { if m.video_price_1080p != nil { fields = append(fields, group.FieldVideoPrice1080p) } + if m.web_search_price_per_call != nil { + fields = append(fields, group.FieldWebSearchPricePerCall) + } if m.claude_code_only != nil { fields = append(fields, group.FieldClaudeCodeOnly) } @@ -23856,6 +23931,8 @@ func (m *GroupMutation) Field(name string) (ent.Value, bool) { return m.VideoPrice720p() case group.FieldVideoPrice1080p: return m.VideoPrice1080p() + case group.FieldWebSearchPricePerCall: + return m.WebSearchPricePerCall() case group.FieldClaudeCodeOnly: return m.ClaudeCodeOnly() case group.FieldFallbackGroupID: @@ -23959,6 +24036,8 @@ func (m *GroupMutation) OldField(ctx context.Context, name string) (ent.Value, e return m.OldVideoPrice720p(ctx) case group.FieldVideoPrice1080p: return m.OldVideoPrice1080p(ctx) + case group.FieldWebSearchPricePerCall: + return m.OldWebSearchPricePerCall(ctx) case group.FieldClaudeCodeOnly: return m.OldClaudeCodeOnly(ctx) case group.FieldFallbackGroupID: @@ -24222,6 +24301,13 @@ func (m *GroupMutation) SetField(name string, value ent.Value) error { } m.SetVideoPrice1080p(v) return nil + case group.FieldWebSearchPricePerCall: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetWebSearchPricePerCall(v) + return nil case group.FieldClaudeCodeOnly: v, ok := value.(bool) if !ok { @@ -24383,6 +24469,9 @@ func (m *GroupMutation) AddedFields() []string { if m.addvideo_price_1080p != nil { fields = append(fields, group.FieldVideoPrice1080p) } + if m.addweb_search_price_per_call != nil { + fields = append(fields, group.FieldWebSearchPricePerCall) + } if m.addfallback_group_id != nil { fields = append(fields, group.FieldFallbackGroupID) } @@ -24435,6 +24524,8 @@ func (m *GroupMutation) AddedField(name string) (ent.Value, bool) { return m.AddedVideoPrice720p() case group.FieldVideoPrice1080p: return m.AddedVideoPrice1080p() + case group.FieldWebSearchPricePerCall: + return m.AddedWebSearchPricePerCall() case group.FieldFallbackGroupID: return m.AddedFallbackGroupID() case group.FieldFallbackGroupIDOnInvalidRequest: @@ -24564,6 +24655,13 @@ func (m *GroupMutation) AddField(name string, value ent.Value) error { } m.AddVideoPrice1080p(v) return nil + case group.FieldWebSearchPricePerCall: + v, ok := value.(float64) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.AddWebSearchPricePerCall(v) + return nil case group.FieldFallbackGroupID: v, ok := value.(int64) if !ok { @@ -24633,6 +24731,9 @@ func (m *GroupMutation) ClearedFields() []string { if m.FieldCleared(group.FieldVideoPrice1080p) { fields = append(fields, group.FieldVideoPrice1080p) } + if m.FieldCleared(group.FieldWebSearchPricePerCall) { + fields = append(fields, group.FieldWebSearchPricePerCall) + } if m.FieldCleared(group.FieldFallbackGroupID) { fields = append(fields, group.FieldFallbackGroupID) } @@ -24689,6 +24790,9 @@ func (m *GroupMutation) ClearField(name string) error { case group.FieldVideoPrice1080p: m.ClearVideoPrice1080p() return nil + case group.FieldWebSearchPricePerCall: + m.ClearWebSearchPricePerCall() + return nil case group.FieldFallbackGroupID: m.ClearFallbackGroupID() return nil @@ -24802,6 +24906,9 @@ func (m *GroupMutation) ResetField(name string) error { case group.FieldVideoPrice1080p: m.ResetVideoPrice1080p() return nil + case group.FieldWebSearchPricePerCall: + m.ResetWebSearchPricePerCall() + return nil case group.FieldClaudeCodeOnly: m.ResetClaudeCodeOnly() return nil diff --git a/backend/ent/runtime/runtime.go b/backend/ent/runtime/runtime.go index 2b4f477bb1..867f1cbdbd 100644 --- a/backend/ent/runtime/runtime.go +++ b/backend/ent/runtime/runtime.go @@ -1044,53 +1044,53 @@ func init() { // group.DefaultVideoRateMultiplier holds the default value on creation for the video_rate_multiplier field. group.DefaultVideoRateMultiplier = groupDescVideoRateMultiplier.Default.(float64) // groupDescClaudeCodeOnly is the schema descriptor for claude_code_only field. - groupDescClaudeCodeOnly := groupFields[29].Descriptor() + groupDescClaudeCodeOnly := groupFields[30].Descriptor() // group.DefaultClaudeCodeOnly holds the default value on creation for the claude_code_only field. group.DefaultClaudeCodeOnly = groupDescClaudeCodeOnly.Default.(bool) // groupDescModelRoutingEnabled is the schema descriptor for model_routing_enabled field. - groupDescModelRoutingEnabled := groupFields[33].Descriptor() + groupDescModelRoutingEnabled := groupFields[34].Descriptor() // group.DefaultModelRoutingEnabled holds the default value on creation for the model_routing_enabled field. group.DefaultModelRoutingEnabled = groupDescModelRoutingEnabled.Default.(bool) // groupDescMcpXMLInject is the schema descriptor for mcp_xml_inject field. - groupDescMcpXMLInject := groupFields[34].Descriptor() + groupDescMcpXMLInject := groupFields[35].Descriptor() // group.DefaultMcpXMLInject holds the default value on creation for the mcp_xml_inject field. group.DefaultMcpXMLInject = groupDescMcpXMLInject.Default.(bool) // groupDescSupportedModelScopes is the schema descriptor for supported_model_scopes field. - groupDescSupportedModelScopes := groupFields[35].Descriptor() + groupDescSupportedModelScopes := groupFields[36].Descriptor() // group.DefaultSupportedModelScopes holds the default value on creation for the supported_model_scopes field. group.DefaultSupportedModelScopes = groupDescSupportedModelScopes.Default.([]string) // groupDescSortOrder is the schema descriptor for sort_order field. - groupDescSortOrder := groupFields[36].Descriptor() + groupDescSortOrder := groupFields[37].Descriptor() // group.DefaultSortOrder holds the default value on creation for the sort_order field. group.DefaultSortOrder = groupDescSortOrder.Default.(int) // groupDescAllowMessagesDispatch is the schema descriptor for allow_messages_dispatch field. - groupDescAllowMessagesDispatch := groupFields[37].Descriptor() + groupDescAllowMessagesDispatch := groupFields[38].Descriptor() // group.DefaultAllowMessagesDispatch holds the default value on creation for the allow_messages_dispatch field. group.DefaultAllowMessagesDispatch = groupDescAllowMessagesDispatch.Default.(bool) // groupDescRequireOauthOnly is the schema descriptor for require_oauth_only field. - groupDescRequireOauthOnly := groupFields[38].Descriptor() + groupDescRequireOauthOnly := groupFields[39].Descriptor() // group.DefaultRequireOauthOnly holds the default value on creation for the require_oauth_only field. group.DefaultRequireOauthOnly = groupDescRequireOauthOnly.Default.(bool) // groupDescRequirePrivacySet is the schema descriptor for require_privacy_set field. - groupDescRequirePrivacySet := groupFields[39].Descriptor() + groupDescRequirePrivacySet := groupFields[40].Descriptor() // group.DefaultRequirePrivacySet holds the default value on creation for the require_privacy_set field. group.DefaultRequirePrivacySet = groupDescRequirePrivacySet.Default.(bool) // groupDescDefaultMappedModel is the schema descriptor for default_mapped_model field. - groupDescDefaultMappedModel := groupFields[40].Descriptor() + groupDescDefaultMappedModel := groupFields[41].Descriptor() // group.DefaultDefaultMappedModel holds the default value on creation for the default_mapped_model field. group.DefaultDefaultMappedModel = groupDescDefaultMappedModel.Default.(string) // group.DefaultMappedModelValidator is a validator for the "default_mapped_model" field. It is called by the builders before save. group.DefaultMappedModelValidator = groupDescDefaultMappedModel.Validators[0].(func(string) error) // groupDescMessagesDispatchModelConfig is the schema descriptor for messages_dispatch_model_config field. - groupDescMessagesDispatchModelConfig := groupFields[41].Descriptor() + groupDescMessagesDispatchModelConfig := groupFields[42].Descriptor() // group.DefaultMessagesDispatchModelConfig holds the default value on creation for the messages_dispatch_model_config field. group.DefaultMessagesDispatchModelConfig = groupDescMessagesDispatchModelConfig.Default.(domain.OpenAIMessagesDispatchModelConfig) // groupDescModelsListConfig is the schema descriptor for models_list_config field. - groupDescModelsListConfig := groupFields[42].Descriptor() + groupDescModelsListConfig := groupFields[43].Descriptor() // group.DefaultModelsListConfig holds the default value on creation for the models_list_config field. group.DefaultModelsListConfig = groupDescModelsListConfig.Default.(domain.GroupModelsListConfig) // groupDescRpmLimit is the schema descriptor for rpm_limit field. - groupDescRpmLimit := groupFields[43].Descriptor() + groupDescRpmLimit := groupFields[44].Descriptor() // group.DefaultRpmLimit holds the default value on creation for the rpm_limit field. group.DefaultRpmLimit = groupDescRpmLimit.Default.(int) idempotencyrecordMixin := schema.IdempotencyRecord{}.Mixin() diff --git a/backend/ent/schema/group.go b/backend/ent/schema/group.go index b104609a1b..70093a3cfa 100644 --- a/backend/ent/schema/group.go +++ b/backend/ent/schema/group.go @@ -142,6 +142,11 @@ func (Group) Fields() []ent.Field { Optional(). Nillable(). SchemaType(map[string]string{dialect.Postgres: "decimal(20,8)"}), + field.Float("web_search_price_per_call"). + Optional(). + Nillable(). + SchemaType(map[string]string{dialect.Postgres: "decimal(20,8)"}). + Comment("Codex alpha/search 网页搜索单次价格（USD/次）；nil 表示使用默认价 0.01（官方 $10/1000 次）"), // Claude Code 客户端限制 (added by migration 029) field.Bool("claude_code_only"). diff --git a/backend/go.mod b/backend/go.mod index a06f06437d..64b5c30d96 100644 --- a/backend/go.mod +++ b/backend/go.mod @@ -44,6 +44,7 @@ require ( go.uber.org/zap v1.24.0 golang.org/x/crypto v0.51.0 golang.org/x/image v0.39.0 + golang.org/x/mod v0.35.0 golang.org/x/net v0.55.0 golang.org/x/sync v0.20.0 golang.org/x/term v0.43.0 @@ -176,7 +177,6 @@ require ( go.uber.org/multierr v1.9.0 // indirect golang.org/x/arch v0.3.0 // indirect golang.org/x/exp v0.0.0-20251023183803-a4bb9ffd2546 // indirect - golang.org/x/mod v0.35.0 // indirect golang.org/x/sys v0.45.0 // indirect golang.org/x/text v0.37.0 // indirect golang.org/x/tools v0.44.0 // indirect diff --git a/backend/go.sum b/backend/go.sum index 4738443bb9..3d9989bb43 100644 --- a/backend/go.sum +++ b/backend/go.sum @@ -220,6 +220,8 @@ github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovk github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= +github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZgg3U= +github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= github.com/mattn/go-sqlite3 v1.14.17 h1:mCRHCLDUBXgpKAqIKsaAaAsrAlbkeomtRFKXh2L6YIM= github.com/mattn/go-sqlite3 v1.14.17/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg= github.com/mdelapenya/tlscert v0.2.0 h1:7H81W6Z/4weDvZBNOfQte5GpIMo0lGYEeWbkGp5LJHI= @@ -253,6 +255,8 @@ github.com/morikuni/aec v1.0.0 h1:nP9CBfwrvYnBRgY6qfDQkygYDmYwOilePFkwzv4dU8A= github.com/morikuni/aec v1.0.0/go.mod h1:BbKIizmSmc5MMPqRYbxO4ZU0S0+P200+tUnFx7PXmsc= github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= +github.com/olekukonko/tablewriter v0.0.5 h1:P2Ga83D34wi1o9J6Wh1mRuqd4mF/x/lgBS7N7AbDhec= +github.com/olekukonko/tablewriter v0.0.5/go.mod h1:hPp6KlRPjbx+hW8ykQs1w3UBbZlj6HuIJcUGPhkA7kY= github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U= github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM= github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040= @@ -282,6 +286,8 @@ github.com/refraction-networking/utls v1.8.2 h1:j4Q1gJj0xngdeH+Ox/qND11aEfhpgoEv github.com/refraction-networking/utls v1.8.2/go.mod h1:jkSOEkLqn+S/jtpEHPOsVv/4V4EVnelwbMQl4vCWXAM= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= +github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY= +github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= github.com/robfig/cron/v3 v3.0.1 h1:WdRxkvbJztn8LMz/QEvLN5sBU+xKpSqwwUO1Pjr4qDs= github.com/robfig/cron/v3 v3.0.1/go.mod h1:eQICP3HwyT7UooqI/z+Ov+PtYAWygg1TEWWzGIFLtro= github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII= @@ -314,6 +320,8 @@ github.com/spf13/afero v1.11.0 h1:WJQKhtpdm3v2IzqG8VMqrr6Rf3UYpEF239Jy9wNepM8= github.com/spf13/afero v1.11.0/go.mod h1:GH9Y3pIexgf1MTIWtNGyogA5MwRIDXGUr+hbWNoBjkY= github.com/spf13/cast v1.6.0 h1:GEiTHELF+vaR5dhz3VqZfFSzZjYbgeKDpBxQVS4GYJ0= github.com/spf13/cast v1.6.0/go.mod h1:ancEpBxwJDODSW/UG4rDrAqiKolqNNh2DX3mk86cAdo= +github.com/spf13/cobra v1.7.0 h1:hyqWnYt1ZQShIddO5kBpj3vu05/++x6tJ6dg8EC572I= +github.com/spf13/cobra v1.7.0/go.mod h1:uLxZILRyS/50WlhOIKD7W6V5bgeIt+4sICxh6uRMrb0= github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA= github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/spf13/viper v1.18.2 h1:LUXCnvUvSM6FXAsj6nnfc8Q2tp1dIgUfY9Kc8GsSOiQ= diff --git a/backend/internal/handler/admin/grok_oauth_handler_test.go b/backend/internal/handler/admin/grok_oauth_handler_test.go index 6ac77e0e56..0b0f0d1aba 100644 --- a/backend/internal/handler/admin/grok_oauth_handler_test.go +++ b/backend/internal/handler/admin/grok_oauth_handler_test.go @@ -98,7 +98,7 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) { require.Contains(t, rec.Body.String(), `"source":"active_probe"`) require.Contains(t, rec.Body.String(), `"headers_observed":true`) require.NotContains(t, rec.Body.String(), "access-token") - require.Equal(t, xai.DefaultBaseURL+"/responses", upstream.lastReq.URL.String()) + require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization")) require.Contains(t, string(upstream.lastBody), `"store":false`) require.NotNil(t, repo.updates[42]) diff --git a/backend/internal/handler/admin/group_handler.go b/backend/internal/handler/admin/group_handler.go index 56a0b29ed0..02a4446846 100644 --- a/backend/internal/handler/admin/group_handler.go +++ b/backend/internal/handler/admin/group_handler.go @@ -110,6 +110,7 @@ type CreateGroupRequest struct { VideoPrice480P *float64 `json:"video_price_480p"` VideoPrice720P *float64 `json:"video_price_720p"` VideoPrice1080P *float64 `json:"video_price_1080p"` + WebSearchPricePerCall *float64 `json:"web_search_price_per_call"` ClaudeCodeOnly bool `json:"claude_code_only"` FallbackGroupID *int64 `json:"fallback_group_id"` FallbackGroupIDOnInvalidRequest *int64 `json:"fallback_group_id_on_invalid_request"` @@ -163,6 +164,7 @@ type UpdateGroupRequest struct { VideoPrice480P *float64 `json:"video_price_480p"` VideoPrice720P *float64 `json:"video_price_720p"` VideoPrice1080P *float64 `json:"video_price_1080p"` + WebSearchPricePerCall *float64 `json:"web_search_price_per_call"` ClaudeCodeOnly *bool `json:"claude_code_only"` FallbackGroupID *int64 `json:"fallback_group_id"` FallbackGroupIDOnInvalidRequest *int64 `json:"fallback_group_id_on_invalid_request"` @@ -334,6 +336,7 @@ func (h *GroupHandler) Create(c *gin.Context) { VideoPrice480P: req.VideoPrice480P, VideoPrice720P: req.VideoPrice720P, VideoPrice1080P: req.VideoPrice1080P, + WebSearchPricePerCall: req.WebSearchPricePerCall, ClaudeCodeOnly: req.ClaudeCodeOnly, FallbackGroupID: req.FallbackGroupID, FallbackGroupIDOnInvalidRequest: req.FallbackGroupIDOnInvalidRequest, @@ -402,6 +405,7 @@ func (h *GroupHandler) Update(c *gin.Context) { VideoPrice480P: req.VideoPrice480P, VideoPrice720P: req.VideoPrice720P, VideoPrice1080P: req.VideoPrice1080P, + WebSearchPricePerCall: req.WebSearchPricePerCall, ClaudeCodeOnly: req.ClaudeCodeOnly, FallbackGroupID: req.FallbackGroupID, FallbackGroupIDOnInvalidRequest: req.FallbackGroupIDOnInvalidRequest, diff --git a/backend/internal/handler/dto/mappers.go b/backend/internal/handler/dto/mappers.go index c2dc13c4a5..3c45c3b95e 100644 --- a/backend/internal/handler/dto/mappers.go +++ b/backend/internal/handler/dto/mappers.go @@ -199,6 +199,7 @@ func groupFromServiceBase(g *service.Group) Group { VideoPrice480P: g.VideoPrice480P, VideoPrice720P: g.VideoPrice720P, VideoPrice1080P: g.VideoPrice1080P, + WebSearchPricePerCall: g.WebSearchPricePerCall, ClaudeCodeOnly: g.ClaudeCodeOnly, FallbackGroupID: g.FallbackGroupID, FallbackGroupIDOnInvalidRequest: g.FallbackGroupIDOnInvalidRequest, diff --git a/backend/internal/handler/dto/types.go b/backend/internal/handler/dto/types.go index 53b0ef4878..619926c1e4 100644 --- a/backend/internal/handler/dto/types.go +++ b/backend/internal/handler/dto/types.go @@ -120,6 +120,8 @@ type Group struct { VideoPrice480P *float64 `json:"video_price_480p"` VideoPrice720P *float64 `json:"video_price_720p"` VideoPrice1080P *float64 `json:"video_price_1080p"` + // Codex alpha/search 网页搜索单次价格（USD/次）；null 表示使用默认价 0.01 + WebSearchPricePerCall *float64 `json:"web_search_price_per_call"` // Claude Code 客户端限制 ClaudeCodeOnly bool `json:"claude_code_only"` diff --git a/backend/internal/handler/endpoint.go b/backend/internal/handler/endpoint.go index 0b9930c5cc..2871eef6ce 100644 --- a/backend/internal/handler/endpoint.go +++ b/backend/internal/handler/endpoint.go @@ -18,6 +18,7 @@ const ( EndpointMessages = "/v1/messages" EndpointChatCompletions = "/v1/chat/completions" EndpointEmbeddings = "/v1/embeddings" + EndpointAlphaSearch = "/v1/alpha/search" EndpointResponses = "/v1/responses" EndpointResponsesCompact = "/v1/responses/compact" EndpointImagesGenerations = "/v1/images/generations" @@ -75,6 +76,8 @@ func NormalizeInboundEndpoint(path string) string { switch { case strings.Contains(path, EndpointEmbeddings): return EndpointEmbeddings + case strings.Contains(path, EndpointAlphaSearch) || isBareOrSubpathOf(strings.TrimRight(path, "/"), "/alpha/search") || isBareOrSubpathOf(strings.TrimRight(path, "/"), "/backend-api/codex/alpha/search"): + return EndpointAlphaSearch case strings.Contains(path, EndpointChatCompletions): return EndpointChatCompletions case strings.Contains(path, EndpointMessages): @@ -155,8 +158,11 @@ func isBareOrSubpathOf(path, root string) bool { // account platform and the normalized inbound endpoint. // // Platform-specific rules: -// - OpenAI always forwards to /v1/responses (with optional subpath -// such as /v1/responses/compact preserved from the raw URL). +// - OpenAI and Grok text compatibility routes forward to /v1/responses +// (with optional subpath such as /v1/responses/compact preserved from +// the raw URL); native endpoints such as embeddings and alpha search +// retain their paths. Grok raw Chat requests override this through the +// forwarding result consumed by resolveOpenAIUpstreamEndpoint. // - Anthropic → /v1/messages // - Gemini → /v1beta/models // - Antigravity → /v1/messages (Claude) or gemini (Gemini) @@ -167,7 +173,7 @@ func DeriveUpstreamEndpoint(inbound, rawRequestPath, platform string) string { switch platform { case service.PlatformOpenAI, service.PlatformGrok: - if inbound == EndpointEmbeddings || inbound == EndpointImagesGenerations || inbound == EndpointImagesEdits || inbound == EndpointVideosGenerations || inbound == EndpointVideos { + if inbound == EndpointEmbeddings || inbound == EndpointAlphaSearch || inbound == EndpointImagesGenerations || inbound == EndpointImagesEdits || inbound == EndpointVideosGenerations || inbound == EndpointVideos { return inbound } // OpenAI forwards everything to the Responses API. diff --git a/backend/internal/handler/endpoint_test.go b/backend/internal/handler/endpoint_test.go index 96ed1292b3..dc8ad728c8 100644 --- a/backend/internal/handler/endpoint_test.go +++ b/backend/internal/handler/endpoint_test.go @@ -25,6 +25,7 @@ func TestNormalizeInboundEndpoint(t *testing.T) { {"/v1/messages", EndpointMessages}, {"/v1/chat/completions", EndpointChatCompletions}, {"/v1/embeddings", EndpointEmbeddings}, + {"/v1/alpha/search", EndpointAlphaSearch}, {"/v1/responses", EndpointResponses}, {"/v1/responses/compact", EndpointResponsesCompact}, {"/v1/responses/compact/detail", EndpointResponsesCompact}, @@ -50,11 +51,13 @@ func TestNormalizeInboundEndpoint(t *testing.T) { {"/responses", EndpointResponses}, {"/responses/compact", EndpointResponsesCompact}, {"/responses/compact/detail", EndpointResponsesCompact}, + {"/alpha/search", EndpointAlphaSearch}, // Bare Codex direct alias route — root vs. compact. {"/backend-api/codex/responses", EndpointResponses}, {"/backend-api/codex/responses/compact", EndpointResponsesCompact}, {"/backend-api/codex/responses/compact/detail", EndpointResponsesCompact}, + {"/backend-api/codex/alpha/search", EndpointAlphaSearch}, // Must NOT generalize to arbitrary paths merely ending in // "/responses" (or "/responses/compact") that are unrelated to @@ -119,8 +122,11 @@ func TestDeriveUpstreamEndpoint(t *testing.T) { {"openai from messages", EndpointMessages, "/v1/messages", service.PlatformOpenAI, EndpointResponses}, {"openai from completions", EndpointChatCompletions, "/v1/chat/completions", service.PlatformOpenAI, EndpointResponses}, {"openai embeddings", EndpointEmbeddings, "/v1/embeddings", service.PlatformOpenAI, EndpointEmbeddings}, + {"openai alpha search", EndpointAlphaSearch, "/backend-api/codex/alpha/search", service.PlatformOpenAI, EndpointAlphaSearch}, {"openai image generations", EndpointImagesGenerations, "/v1/images/generations", service.PlatformOpenAI, EndpointImagesGenerations}, {"openai image edits", EndpointImagesEdits, "/openai/v1/images/edits", service.PlatformOpenAI, EndpointImagesEdits}, + {"grok chat defaults to responses without runtime result", EndpointChatCompletions, "/v1/chat/completions", service.PlatformGrok, EndpointResponses}, + {"grok responses", EndpointResponses, "/v1/responses", service.PlatformGrok, EndpointResponses}, {"grok video generations", EndpointVideosGenerations, "/v1/videos/generations", service.PlatformGrok, EndpointVideosGenerations}, {"grok video status", EndpointVideos, "/videos/req_123", service.PlatformGrok, EndpointVideos}, @@ -138,6 +144,59 @@ func TestDeriveUpstreamEndpoint(t *testing.T) { } } +func TestResolveOpenAIUpstreamEndpointPrefersForwardResult(t *testing.T) { + tests := []struct { + name string + account *service.Account + result *service.OpenAIForwardResult + runtimeEndpoint string + want string + }{ + { + name: "grok raw chat result overrides stale context", + account: &service.Account{Platform: service.PlatformGrok, Type: service.AccountTypeOAuth}, + result: &service.OpenAIForwardResult{UpstreamEndpoint: EndpointChatCompletions}, + runtimeEndpoint: EndpointResponses, + want: EndpointChatCompletions, + }, + { + name: "grok chat bridged to responses", + account: &service.Account{Platform: service.PlatformGrok, Type: service.AccountTypeOAuth}, + result: &service.OpenAIForwardResult{UpstreamEndpoint: EndpointResponses}, + want: EndpointResponses, + }, + { + name: "grok empty result keeps responses default", + account: &service.Account{Platform: service.PlatformGrok, Type: service.AccountTypeOAuth}, + result: &service.OpenAIForwardResult{}, + want: EndpointResponses, + }, + { + name: "grok raw error uses runtime endpoint", + account: &service.Account{Platform: service.PlatformGrok, Type: service.AccountTypeOAuth}, + runtimeEndpoint: EndpointChatCompletions, + want: EndpointChatCompletions, + }, + { + name: "openai behavior remains responses", + account: &service.Account{Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth}, + result: &service.OpenAIForwardResult{}, + want: EndpointResponses, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, EndpointChatCompletions, nil) + c.Set(ctxKeyInboundEndpoint, EndpointChatCompletions) + service.SetActualOpenAIUpstreamEndpoint(c, tt.runtimeEndpoint) + require.Equal(t, tt.want, resolveOpenAIUpstreamEndpoint(c, tt.account, tt.result)) + }) + } +} + // ────────────────────────────────────────────────────────── // responsesSubpathSuffix // ────────────────────────────────────────────────────────── diff --git a/backend/internal/handler/no_account_error.go b/backend/internal/handler/no_account_error.go index a3bf3b049e..001cef611d 100644 --- a/backend/internal/handler/no_account_error.go +++ b/backend/internal/handler/no_account_error.go @@ -107,3 +107,31 @@ func classifyNoAccountErrorFromGin( } return classifyNoAccountError(ctx, diag, apiKey, routingModel, displayModel, platform) } + +func classifyOpenAICompatibleNoAccountErrorFromGin( + c *gin.Context, + diag service.ModelAvailabilityDiagnoser, + apiKey *service.APIKey, + routingModel string, + displayModel string, +) noAccountErrorClassification { + return classifyNoAccountErrorFromGin( + c, + diag, + apiKey, + routingModel, + displayModel, + openAICompatibleRequestPlatform(apiKey), + ) +} + +func openAICompatibleSelectionErrorForLog(err error, platform string) error { + if err == nil || platform != service.PlatformGrok { + return err + } + message := strings.ReplaceAll(err.Error(), "OpenAI accounts", "Grok accounts") + if message == err.Error() { + return err + } + return fmt.Errorf("%s", message) +} diff --git a/backend/internal/handler/no_account_error_test.go b/backend/internal/handler/no_account_error_test.go index cfe41bb34f..174da82cc7 100644 --- a/backend/internal/handler/no_account_error_test.go +++ b/backend/internal/handler/no_account_error_test.go @@ -4,6 +4,7 @@ package handler import ( "context" + "fmt" "net/http" "net/http/httptest" "testing" @@ -114,6 +115,33 @@ func TestClassifyNoAccountError_ModelNotSupported_Returns404(t *testing.T) { require.Equal(t, int64(42), *fd.calls[0].GroupID) } +func TestClassifyOpenAICompatibleNoAccountError_GrokUsesGrokPlatform(t *testing.T) { + c := newTestGinContextWithRequest() + fd := &fakeDiagnoser{resp: service.ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: false}} + groupID := int64(43) + apiKey := &service.APIKey{ + GroupID: &groupID, + Group: &service.Group{ + ID: groupID, + Platform: service.PlatformGrok, + }, + } + + cls := classifyOpenAICompatibleNoAccountErrorFromGin(c, fd, apiKey, "grok-4.5", "grok-4.5") + + require.Equal(t, http.StatusNotFound, cls.Status) + require.Equal(t, "model_not_found", cls.ErrType) + require.True(t, cls.ModelNotFound) + require.Len(t, fd.calls, 1) + require.Equal(t, service.PlatformGrok, fd.calls[0].Platform) + + logErr := openAICompatibleSelectionErrorForLog( + fmt.Errorf("no available OpenAI accounts supporting model: grok-4.5"), + service.PlatformGrok, + ) + require.EqualError(t, logErr, "no available Grok accounts supporting model: grok-4.5") +} + func TestClassifyNoAccountError_HasModelSupport_KeepsRoutingMessageGenerationToCaller(t *testing.T) { c := newTestGinContextWithRequest() fd := &fakeDiagnoser{resp: service.ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true}} diff --git a/backend/internal/handler/openai_alpha_search.go b/backend/internal/handler/openai_alpha_search.go new file mode 100644 index 0000000000..a808145235 --- /dev/null +++ b/backend/internal/handler/openai_alpha_search.go @@ -0,0 +1,250 @@ +package handler + +import ( + "context" + "errors" + "net/http" + "strconv" + "strings" + "time" + + pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil" + "github.com/Wei-Shaw/sub2api/internal/pkg/ip" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" + middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" + "go.uber.org/zap" +) + +// AlphaSearch proxies the standalone search endpoint used by Codex Responses Lite. +func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) { + streamStarted := false + defer h.recoverResponsesPanic(c, &streamStarted) + setOpenAIClientTransportHTTP(c) + requestStart := time.Now() + + apiKey, ok := middleware2.GetAPIKeyFromContext(c) + if !ok || apiKey.Group == nil { + h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key") + return + } + if apiKey.Group.Platform != service.PlatformOpenAI { + h.errorResponse(c, http.StatusNotFound, "not_found_error", "Codex alpha search is only available for OpenAI groups") + return + } + subject, ok := middleware2.GetAuthSubjectFromContext(c) + if !ok { + h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found") + return + } + reqLog := requestLogger( + c, + "handler.openai_gateway.alpha_search", + zap.Int64("user_id", subject.UserID), + zap.Int64("api_key_id", apiKey.ID), + zap.Any("group_id", apiKey.GroupID), + ) + if !h.ensureResponsesDependencies(c, reqLog) { + return + } + + body, err := pkghttputil.ReadRequestBodyWithPrealloc(c.Request) + if err != nil { + if maxErr, ok := extractMaxBytesError(err); ok { + h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit)) + return + } + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body") + return + } + if len(body) == 0 { + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty") + return + } + if !gjson.ValidBytes(body) { + logRequestBodyParseFailure(reqLog, body, nil) + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body") + return + } + + modelResult := gjson.GetBytes(body, "model") + if !modelResult.Exists() || modelResult.Type != gjson.String || strings.TrimSpace(modelResult.String()) == "" { + h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required") + return + } + requestedModel := strings.TrimSpace(modelResult.String()) + reqLog = reqLog.With(zap.String("model", requestedModel)) + setOpsRequestContext(c, requestedModel, false) + setOpsEndpointContext(c, "", int16(service.RequestTypeSync)) + + channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, requestedModel) + forwardBody := openAIModelMappedBody(body, channelMapping.Mapped, channelMapping.MappedModel, h.gatewayService.ReplaceModelInBody) + subscription, _ := middleware2.GetSubscriptionFromContext(c) + service.SetOpsLatencyMs(c, service.OpsAuthLatencyMsKey, time.Since(requestStart).Milliseconds()) + + userRelease, acquired := h.acquireResponsesUserSlot(c, subject.UserID, subject.Concurrency, false, &streamStarted, reqLog) + if !acquired { + return + } + if userRelease != nil { + defer userRelease() + } + + if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil { + status, code, message, retryAfter := billingErrorDetails(err) + if retryAfter > 0 { + c.Header("Retry-After", strconv.Itoa(retryAfter)) + } + h.errorResponse(c, status, code, message) + return + } + + searchID := strings.TrimSpace(gjson.GetBytes(body, "id").String()) + sessionHash := h.gatewayService.GenerateSessionHashWithFallback(c, nil, searchID) + failedAccountIDs := make(map[int64]struct{}) + var lastFailoverErr *service.UpstreamFailoverError + switchCount := 0 + routingStart := time.Now() + + for { + selection, _, err := h.gatewayService.SelectAccountWithSchedulerForCapability( + c.Request.Context(), + apiKey.GroupID, + "", + sessionHash, + requestedModel, + failedAccountIDs, + service.OpenAIUpstreamTransportHTTPSSE, + service.OpenAIEndpointCapabilityChatCompletions, + false, + false, + service.PlatformOpenAI, + ) + if err != nil || selection == nil || selection.Account == nil { + if len(failedAccountIDs) == 0 { + cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestedModel, requestedModel, service.PlatformOpenAI) + if !cls.ModelNotFound { + markOpsRoutingCapacityLimitedIfNoAvailable(c, err) + } + h.errorResponse(c, cls.Status, cls.ErrType, cls.Message) + return + } + if lastFailoverErr != nil { + h.handleFailoverExhausted(c, lastFailoverErr, false) + } else { + h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed") + } + return + } + + account := selection.Account + setOpsSelectedAccount(c, account.ID, account.Platform) + accountRelease, acquired := h.acquireResponsesAccountSlot(c, apiKey.GroupID, sessionHash, selection, false, &streamStarted, reqLog) + if !acquired { + return + } + service.SetOpsLatencyMs(c, service.OpsRoutingLatencyMsKey, time.Since(routingStart).Milliseconds()) + writerSizeBeforeForward := c.Writer.Size() + forwardStart := time.Now() + var result *service.OpenAIForwardResult + result, err = func() (*service.OpenAIForwardResult, error) { + if accountRelease != nil { + defer accountRelease() + } + return h.gatewayService.ForwardAlphaSearch(c.Request.Context(), c, account, forwardBody) + }() + service.SetOpsLatencyMs(c, service.OpsResponseLatencyMsKey, time.Since(forwardStart).Milliseconds()) + + if err == nil { + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil) + if result != nil { + h.recordAlphaSearchUsage(c, apiKey, account, subscription, channelMapping, requestedModel, body, result, subject.UserID) + } + return + } + + var failoverErr *service.UpstreamFailoverError + if !errors.As(err, &failoverErr) { + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + if c.Writer.Size() == writerSizeBeforeForward { + h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed") + } + reqLog.Warn("openai_alpha_search.forward_failed", zap.Int64("account_id", account.ID), zap.Error(err)) + return + } + + h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil) + if c.Writer.Size() != writerSizeBeforeForward { + h.handleFailoverExhausted(c, failoverErr, true) + return + } + h.gatewayService.RecordOpenAIAccountSwitch() + failedAccountIDs[account.ID] = struct{}{} + lastFailoverErr = failoverErr + if switchCount >= h.maxAccountSwitches { + h.handleFailoverExhausted(c, failoverErr, false) + return + } + switchCount++ + if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount) { + h.handleFailoverExhausted(c, failoverErr, false) + return + } + reqLog.Warn("openai_alpha_search.upstream_failover_switching", + zap.Int64("account_id", account.ID), + zap.Int("upstream_status", failoverErr.StatusCode), + zap.Int("switch_count", switchCount), + ) + } +} + +// recordAlphaSearchUsage 为一次成功的 alpha/search 网页搜索落按次计费用量行 +// （上游不返回 usage 字段，按 WebSearchCalls 走分组单价 × 倍率的按次口径）。 +// 与 images 一致使用 mandatory 池提交，池满时同步兜底执行，保证扣费不丢。 +func (h *OpenAIGatewayHandler) recordAlphaSearchUsage( + c *gin.Context, + apiKey *service.APIKey, + account *service.Account, + subscription *service.UserSubscription, + channelMapping service.ChannelMappingResult, + requestedModel string, + body []byte, + result *service.OpenAIForwardResult, + userID int64, +) { + userAgent := c.GetHeader("User-Agent") + clientIP := ip.GetClientIP(c) + requestPayloadHash := service.HashUsageRequestPayload(body) + inboundEndpoint := GetInboundEndpoint(c) + upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform) + quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) + + h.submitMandatoryUsageRecordTask(c.Request.Context(), func(ctx context.Context) { + if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{ + Result: result, + APIKey: apiKey, + User: apiKey.User, + Account: account, + Subscription: subscription, + InboundEndpoint: inboundEndpoint, + UpstreamEndpoint: upstreamEndpoint, + UserAgent: userAgent, + IPAddress: clientIP, + RequestPayloadHash: requestPayloadHash, + APIKeyService: h.apiKeyService, + QuotaPlatform: quotaPlatform, + ChannelUsageFields: channelMapping.ToUsageFields(requestedModel, result.UpstreamModel), + }); err != nil { + logger.L().With( + zap.String("component", "handler.openai_gateway.alpha_search"), + zap.Int64("user_id", userID), + zap.Int64("api_key_id", apiKey.ID), + zap.Any("group_id", apiKey.GroupID), + zap.String("model", requestedModel), + zap.Int64("account_id", account.ID), + ).Error("openai_alpha_search.record_usage_failed", zap.Error(err)) + } + }) +} diff --git a/backend/internal/handler/openai_chat_completions.go b/backend/internal/handler/openai_chat_completions.go index f5f2522e49..636e143740 100644 --- a/backend/internal/handler/openai_chat_completions.go +++ b/backend/internal/handler/openai_chat_completions.go @@ -5,6 +5,7 @@ import ( "errors" "net/http" "strconv" + "strings" "time" "github.com/Wei-Shaw/sub2api/internal/pkg/ip" @@ -150,11 +151,11 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { ) if err != nil { reqLog.Warn("openai_chat_completions.account_select_failed", - zap.Error(err), + zap.Error(openAICompatibleSelectionErrorForLog(err, requestPlatform)), zap.Int("excluded_account_count", len(failedAccountIDs)), ) if len(failedAccountIDs) == 0 { - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel, service.PlatformOpenAI) + cls := classifyOpenAICompatibleNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel) if !cls.ModelNotFound { markOpsRoutingCapacityLimitedIfNoAvailable(c, err) } @@ -170,7 +171,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { } } if selection == nil || selection.Account == nil { - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel, service.PlatformOpenAI) + cls := classifyOpenAICompatibleNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel) if !cls.ModelNotFound { markOpsRoutingCapacityLimited(c) } @@ -298,7 +299,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { userAgent := c.GetHeader("User-Agent") clientIP := ip.GetClientIP(c) inboundEndpoint := GetInboundEndpoint(c) - upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account) + upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) cyberBlocked := service.GetOpsCyberPolicy(c) != nil @@ -337,14 +338,22 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) { } // resolveOpenAIUpstreamEndpoint returns the actual upstream endpoint for an -// OpenAI account, used by every OpenAI usage-recording site. APIKey accounts -// whose upstream is forced or probed to not support the Responses API are -// served directly via /v1/chat/completions (the raw chat path) regardless of -// the inbound endpoint; everything else goes through the Responses API. -func resolveOpenAIUpstreamEndpoint(c *gin.Context, account *service.Account) string { +// OpenAI-compatible account. A forwarding result is authoritative because a +// single inbound route may choose raw Chat or a Responses bridge at runtime. +// The account-based derivation remains as a fallback for existing callers and +// forwarding paths that do not report their endpoint yet. +func resolveOpenAIUpstreamEndpoint(c *gin.Context, account *service.Account, result *service.OpenAIForwardResult) string { + if result != nil { + if endpoint := strings.TrimSpace(result.UpstreamEndpoint); endpoint != "" { + return endpoint + } + } + if endpoint := service.GetActualOpenAIUpstreamEndpoint(c); endpoint != "" { + return endpoint + } if account != nil && account.Type == service.AccountTypeAPIKey && !openai_compat.ShouldUseResponsesAPI(account.Extra) { - return "/v1/chat/completions" + return EndpointChatCompletions } return GetUpstreamEndpoint(c, account.Platform) } diff --git a/backend/internal/handler/openai_gateway_compact_body_signal_test.go b/backend/internal/handler/openai_gateway_compact_body_signal_test.go index a4bfb90466..a44d47c856 100644 --- a/backend/internal/handler/openai_gateway_compact_body_signal_test.go +++ b/backend/internal/handler/openai_gateway_compact_body_signal_test.go @@ -23,46 +23,61 @@ func newCompactBodySignalTestContext(t *testing.T, path string, body []byte) *gi return c } -// body-signal 提升后必须与 path-based compact 走同一条链路： -// path 改写、requireCompact 判定、stream/store/prompt_cache_key 归一化删除。 -// 回归防护：若 stream 字段存活，Forward 会用流式 handler 解析 compact 的 -// JSON 响应，导致 "stream ended before a terminal event" 的换号 failover 风暴。 -func TestNormalizeOpenAIResponsesCompactRequest_BodySignalPromoted(t *testing.T) { +func TestNormalizeOpenAIResponsesCompactRequest_RemoteV2StaysOnResponses(t *testing.T) { h := &OpenAIGatewayHandler{} body := []byte(`{ - "model":"gpt-5.5", + "model":"gpt-5.6-sol", "stream":true, "store":true, "prompt_cache_key":"pck-signal-1", + "reasoning":{"effort":"max","context":"all_turns"}, "input":[ {"type":"message","role":"user","content":"hello"}, {"type":"compaction_trigger"} ] }`) c := newCompactBodySignalTestContext(t, "/v1/responses", body) + c.Request.Header.Set("x-codex-beta-features", "responses_websockets_v2, remote_compaction_v2, another_feature") normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body) require.True(t, ok) - require.Equal(t, "/v1/responses/compact", c.Request.URL.Path) - require.True(t, isOpenAIRemoteCompactPath(c)) - - require.False(t, gjson.GetBytes(normalized, "stream").Exists()) - require.False(t, gjson.GetBytes(normalized, "store").Exists()) - require.False(t, gjson.GetBytes(normalized, "prompt_cache_key").Exists()) - require.Equal(t, "gpt-5.5", gjson.GetBytes(normalized, "model").String()) - require.True(t, gjson.GetBytes(normalized, "input").IsArray()) + require.Equal(t, "/v1/responses", c.Request.URL.Path) + require.False(t, isOpenAIRemoteCompactPath(c)) + require.Equal(t, body, normalized) + require.True(t, gjson.GetBytes(normalized, "stream").Bool()) + require.True(t, gjson.GetBytes(normalized, "store").Bool()) + require.Equal(t, "pck-signal-1", gjson.GetBytes(normalized, "prompt_cache_key").String()) + require.Equal(t, "max", gjson.GetBytes(normalized, "reasoning.effort").String()) + require.Equal(t, "all_turns", gjson.GetBytes(normalized, "reasoning.context").String()) reqStream, streamOK := parseOpenAICompatibleStream(normalized) require.True(t, streamOK) - require.False(t, reqStream) + require.True(t, reqStream) - seed, exists := c.Get(service.OpenAICompactSessionSeedKeyForTest()) - require.True(t, exists) - require.Equal(t, "pck-signal-1", seed) + _, seedExists := c.Get(service.OpenAICompactSessionSeedKeyForTest()) + require.False(t, seedExists) + _, streamMarkerExists := c.Get(service.OpenAICompactClientStreamKeyForTest()) + require.False(t, streamMarkerExists) } -func TestNormalizeOpenAIResponsesCompactRequest_BodySignalTrailingSlash(t *testing.T) { +func TestNormalizeOpenAIResponsesCompactRequest_RemoteV2PathAliasesStayOnResponses(t *testing.T) { + h := &OpenAIGatewayHandler{} + body := []byte(`{"model":"gpt-5.6-sol","stream":true,"input":[{"type":"compaction_trigger"}]}`) + for _, path := range []string{"/v1/responses/", "/backend-api/codex/responses"} { + t.Run(path, func(t *testing.T) { + c := newCompactBodySignalTestContext(t, path, body) + c.Request.Header.Set("x-codex-beta-features", "remote_compaction_v2") + + normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body) + require.True(t, ok) + require.Equal(t, path, c.Request.URL.Path) + require.Equal(t, body, normalized) + }) + } +} + +func TestNormalizeOpenAIResponsesCompactRequest_BodySignalTrailingSlashPromoted(t *testing.T) { h := &OpenAIGatewayHandler{} body := []byte(`{"model":"gpt-5.5","input":[{"type":"compaction_trigger"}]}`) c := newCompactBodySignalTestContext(t, "/v1/responses/", body) @@ -82,6 +97,64 @@ func TestNormalizeOpenAIResponsesCompactRequest_CodexDirectAliasPromoted(t *test require.Equal(t, "/backend-api/codex/responses/compact", c.Request.URL.Path) } +func TestNormalizeOpenAIResponsesCompactRequest_NonRemoteV2BodySignalPromoted(t *testing.T) { + h := &OpenAIGatewayHandler{} + tests := []struct { + name string + body []byte + betaHeader string + wantMarked bool + }{ + { + name: "no_header", + body: []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"compaction_trigger"}]}`), + wantMarked: true, + }, + { + name: "unrelated_header", + body: []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"compaction_trigger"}]}`), + betaHeader: "responses_websockets_v2", + wantMarked: true, + }, + { + name: "wrong_case_header", + body: []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"compaction_trigger"}]}`), + betaHeader: "REMOTE_COMPACTION_V2", + wantMarked: true, + }, + { + name: "stream_false", + body: []byte(`{"model":"gpt-5.5","stream":false,"input":[{"type":"compaction_trigger"}]}`), + betaHeader: "remote_compaction_v2", + }, + { + name: "stream_absent", + body: []byte(`{"model":"gpt-5.5","input":[{"type":"compaction_trigger"}]}`), + betaHeader: "remote_compaction_v2", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + c := newCompactBodySignalTestContext(t, "/v1/responses", tt.body) + if tt.betaHeader != "" { + c.Request.Header.Set("x-codex-beta-features", tt.betaHeader) + } + + normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), tt.body) + require.True(t, ok) + require.Equal(t, "/v1/responses/compact", c.Request.URL.Path) + require.False(t, gjson.GetBytes(normalized, "stream").Exists()) + + marked, exists := c.Get(service.OpenAICompactClientStreamKeyForTest()) + require.Equal(t, tt.wantMarked, exists) + if tt.wantMarked { + require.Equal(t, true, marked) + } + }) + } +} + func TestNormalizeOpenAIResponsesCompactRequest_NoTriggerUntouched(t *testing.T) { h := &OpenAIGatewayHandler{} body := []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"message","role":"user","content":"hello"}]}`) @@ -99,6 +172,7 @@ func TestNormalizeOpenAIResponsesCompactRequest_PathBasedNoDoubleSuffix(t *testi h := &OpenAIGatewayHandler{} body := []byte(`{"model":"gpt-5.5","stream":true,"store":true,"input":[{"type":"message","role":"user","content":"hello"}]}`) c := newCompactBodySignalTestContext(t, "/v1/responses/compact", body) + c.Request.Header.Set("x-codex-beta-features", "remote_compaction_v2") normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body) require.True(t, ok) @@ -118,36 +192,6 @@ func TestNormalizeOpenAIResponsesCompactRequest_SubpathNotPromoted(t *testing.T) require.Equal(t, body, normalized) } -// 回归 #3875：body-signal 原始请求 stream:true 时必须标记 client-stream， -// 供响应写回阶段把上游 unary JSON 合成回 Codex remote compact v2 所需的 SSE。 -func TestNormalizeOpenAIResponsesCompactRequest_BodySignalStreamTrueMarksClientStream(t *testing.T) { - h := &OpenAIGatewayHandler{} - body := []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"compaction_trigger"}]}`) - c := newCompactBodySignalTestContext(t, "/v1/responses", body) - - _, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body) - require.True(t, ok) - - marked, exists := c.Get(service.OpenAICompactClientStreamKeyForTest()) - require.True(t, exists) - require.Equal(t, true, marked) -} - -func TestNormalizeOpenAIResponsesCompactRequest_BodySignalStreamFalseNotMarked(t *testing.T) { - h := &OpenAIGatewayHandler{} - for name, body := range map[string][]byte{ - "stream_false": []byte(`{"model":"gpt-5.5","stream":false,"input":[{"type":"compaction_trigger"}]}`), - "stream_absent": []byte(`{"model":"gpt-5.5","input":[{"type":"compaction_trigger"}]}`), - } { - c := newCompactBodySignalTestContext(t, "/v1/responses", body) - _, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body) - require.True(t, ok, name) - require.Equal(t, "/v1/responses/compact", c.Request.URL.Path, name) - _, exists := c.Get(service.OpenAICompactClientStreamKeyForTest()) - require.False(t, exists, "case %s 不应标记 client-stream", name) - } -} - // path-based compact（Codex v1 unary 协议）即使 body 带 stream:true 也不标记， // 保持 JSON 写回行为不变。 func TestNormalizeOpenAIResponsesCompactRequest_PathBasedStreamTrueNotMarked(t *testing.T) { diff --git a/backend/internal/handler/openai_gateway_count_tokens.go b/backend/internal/handler/openai_gateway_count_tokens.go index 0461017067..0a010cc176 100644 --- a/backend/internal/handler/openai_gateway_count_tokens.go +++ b/backend/internal/handler/openai_gateway_count_tokens.go @@ -115,8 +115,9 @@ func (h *OpenAIGatewayHandler) CountTokens(c *gin.Context) { ) service.SetOpsLatencyMs(c, service.OpsAuthLatencyMsKey, time.Since(requestStart).Milliseconds()) if err != nil { - reqLog.Warn("openai_count_tokens.account_select_failed", zap.Error(err)) - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, currentRoutingModel, reqModel, service.PlatformOpenAI) + requestPlatform := openAICompatibleRequestPlatform(apiKey) + reqLog.Warn("openai_count_tokens.account_select_failed", zap.Error(openAICompatibleSelectionErrorForLog(err, requestPlatform))) + cls := classifyOpenAICompatibleNoAccountErrorFromGin(c, h.gatewayService, apiKey, currentRoutingModel, reqModel) if !cls.ModelNotFound { markOpsRoutingCapacityLimitedIfNoAvailable(c, err) } @@ -124,7 +125,7 @@ func (h *OpenAIGatewayHandler) CountTokens(c *gin.Context) { return } if selection == nil || selection.Account == nil { - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, currentRoutingModel, reqModel, service.PlatformOpenAI) + cls := classifyOpenAICompatibleNoAccountErrorFromGin(c, h.gatewayService, apiKey, currentRoutingModel, reqModel) if !cls.ModelNotFound { markOpsRoutingCapacityLimited(c) } diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index 83e644d857..afa2a5073a 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -351,7 +351,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { ) if err != nil { reqLog.Warn("openai.account_select_failed", - zap.Error(err), + zap.Error(openAICompatibleSelectionErrorForLog(err, requestPlatform)), zap.Int("excluded_account_count", len(failedAccountIDs)), ) if len(failedAccountIDs) == 0 { @@ -360,7 +360,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "compact_not_supported", "No available OpenAI accounts support /responses/compact", streamStarted) return } - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel, service.PlatformOpenAI) + cls := classifyOpenAICompatibleNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel) if !cls.ModelNotFound { markOpsRoutingCapacityLimitedIfNoAvailable(c, err) } @@ -375,7 +375,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { return } if selection == nil || selection.Account == nil { - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel, service.PlatformOpenAI) + cls := classifyOpenAICompatibleNoAccountErrorFromGin(c, h.gatewayService, apiKey, reqModel, reqModel) if !cls.ModelNotFound { markOpsRoutingCapacityLimited(c) } @@ -522,7 +522,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) { clientIP := ip.GetClientIP(c) requestPayloadHash := service.HashUsageRequestPayload(body) inboundEndpoint := GetInboundEndpoint(c) - upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account) + upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) // 使用量记录通过有界 worker 池提交，避免请求热路径创建无界 goroutine。 @@ -580,21 +580,33 @@ func isBareOpenAIResponsesPath(c *gin.Context) bool { return strings.HasSuffix(normalizedPath, "/responses") } -// normalizeOpenAIResponsesCompactRequest 统一处理两种入站 compact 形态： -// path-based（POST /v1/responses/compact）与 Codex remote compact v2 的 -// body-signal（普通 POST /v1/responses 的 input 中携带 type=compaction_trigger， -// 见 #3777）。body-signal 命中时在 stream 解析、compact body 归一化与 -// requireCompact 调度判定之前改写 URL path，使后续全部链路（含 passthrough -// 分支与上游 URL 构建）与 path-based 完全一致。 +func isOpenAIRemoteCompactionV2Request(c *gin.Context, body []byte) bool { + stream, valid := parseOpenAICompatibleStream(body) + if !valid || !stream || c == nil || c.Request == nil { + return false + } + for _, header := range c.Request.Header.Values("x-codex-beta-features") { + for _, feature := range strings.Split(header, ",") { + if strings.TrimSpace(feature) == "remote_compaction_v2" { + return true + } + } + } + return false +} + +// normalizeOpenAIResponsesCompactRequest keeps Codex remote compaction v2 on +// its native streaming /responses wire and preserves the legacy body-signal +// promotion for clients that do not explicitly advertise that protocol. // 返回归一化后的 body；ok=false 表示错误响应已写出，调用方应直接 return。 func (h *OpenAIGatewayHandler) normalizeOpenAIResponsesCompactRequest(c *gin.Context, reqLog *zap.Logger, body []byte) ([]byte, bool) { isCompactRequest := service.IsOpenAIResponsesCompactPathForTest(c) if !isCompactRequest && isBareOpenAIResponsesPath(c) && service.HasCompactionTriggerInInput(body) { + if isOpenAIRemoteCompactionV2Request(c, body) { + return body, true + } c.Request.URL.Path = strings.TrimRight(c.Request.URL.Path, "/") + "/compact" isCompactRequest = true - // Codex remote compact v2 的原始请求是流式 /responses：白名单归一化会删除 - // stream 并让上游走 unary JSON，但客户端仍按 SSE 消费响应。记录原始 - // stream 意图，响应写回阶段据此把 JSON 合成回 SSE（#3875）。 clientStream := gjson.GetBytes(body, "stream").Bool() if clientStream { service.MarkOpenAICompactClientStream(c) @@ -843,12 +855,12 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { ) if err != nil { reqLog.Warn("openai_messages.account_select_failed", - zap.Error(err), + zap.Error(openAICompatibleSelectionErrorForLog(err, requestPlatform)), zap.Int("excluded_account_count", len(failedAccountIDs)), ) if len(failedAccountIDs) == 0 { if err != nil { - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, currentRoutingModel, reqModel, service.PlatformOpenAI) + cls := classifyOpenAICompatibleNoAccountErrorFromGin(c, h.gatewayService, apiKey, currentRoutingModel, reqModel) if !cls.ModelNotFound { markOpsRoutingCapacityLimitedIfNoAvailable(c, err) } @@ -865,7 +877,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { } } if selection == nil || selection.Account == nil { - cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, currentRoutingModel, reqModel, service.PlatformOpenAI) + cls := classifyOpenAICompatibleNoAccountErrorFromGin(c, h.gatewayService, apiKey, currentRoutingModel, reqModel) if !cls.ModelNotFound { markOpsRoutingCapacityLimited(c) } @@ -994,7 +1006,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) { clientIP := ip.GetClientIP(c) requestPayloadHash := service.HashUsageRequestPayload(body) inboundEndpoint := GetInboundEndpoint(c) - upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account) + upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) cyberBlocked := service.GetOpsCyberPolicy(c) != nil @@ -1444,7 +1456,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { ) if err != nil { reqLog.Warn("openai.websocket_account_select_failed", - zap.Error(err), + zap.Error(openAICompatibleSelectionErrorForLog(err, requestPlatform)), zap.Int("excluded_account_count", len(failedAccountIDs)), ) if lastFailoverErr != nil { @@ -1601,7 +1613,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) { } h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, result.FirstTokenMs) inboundEndpoint := GetInboundEndpoint(c) - upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account) + upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result) quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey) cyberBlocked := service.GetOpsCyberPolicy(c) != nil h.submitOpenAIUsageRecordTask(ctx, result, func(taskCtx context.Context) { @@ -2441,7 +2453,7 @@ func (h *OpenAIGatewayHandler) recordCyberPolicyIfMarked(c *gin.Context, apiKey var accountID int64 if account != nil { accountID = account.ID - upstreamEndpoint = resolveOpenAIUpstreamEndpoint(c, account) + upstreamEndpoint = resolveOpenAIUpstreamEndpoint(c, account, nil) } stream := false if v, ok := c.Get(opsStreamKey); ok { diff --git a/backend/internal/handler/ops_capture_writer_nil_test.go b/backend/internal/handler/ops_capture_writer_nil_test.go index 4e96333f9d..88aa7c043f 100644 --- a/backend/internal/handler/ops_capture_writer_nil_test.go +++ b/backend/internal/handler/ops_capture_writer_nil_test.go @@ -1,9 +1,15 @@ package handler import ( + "net/http" + "net/http/httptest" "testing" + "time" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) func TestOpsCaptureWriter_NilInnerWriter_NoPanic(t *testing.T) { @@ -57,3 +63,28 @@ func TestOpsCaptureWriter_NilInnerWriter_NoPanic(t *testing.T) { assert.Nil(t, p) }) } + +func TestOpsCaptureWriter_CompactKeepaliveRestoresOriginalWriter(t *testing.T) { + gin.SetMode(gin.TestMode) + router := gin.New() + outerStatus := -1 + router.Use(func(c *gin.Context) { + c.Next() + outerStatus = c.Writer.Status() + }) + router.Use(OpsErrorLoggerMiddleware(nil)) + router.GET("/compact", func(c *gin.Context) { + service.MarkOpenAICompactClientStream(c) + stop := service.StartOpenAICompactSSEKeepalive(c, time.Hour) + defer stop() + c.Status(http.StatusOK) + }) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/compact", nil) + require.NotPanics(t, func() { + router.ServeHTTP(recorder, request) + }) + require.Equal(t, http.StatusOK, outerStatus) + require.Equal(t, http.StatusOK, recorder.Code) +} diff --git a/backend/internal/pkg/apicompat/responses_stream_event_wire_test.go b/backend/internal/pkg/apicompat/responses_stream_event_wire_test.go index f44f3e7770..fb138a1469 100644 --- a/backend/internal/pkg/apicompat/responses_stream_event_wire_test.go +++ b/backend/internal/pkg/apicompat/responses_stream_event_wire_test.go @@ -131,3 +131,56 @@ func TestWire_UnknownEventFallsBackToDefault(t *testing.T) { }) require.Contains(t, m, "response") } + +func TestResponsesOutputUnmarshal_ToolSearchObjectArguments(t *testing.T) { + var item ResponsesOutput + require.NoError(t, json.Unmarshal([]byte(`{ + "type":"tool_search_call", + "id":"item_1", + "call_id":"call_1", + "execution":"client", + "arguments":{"query":"gmail","limit":2} + }`), &item)) + require.Equal(t, "tool_search_call", item.Type) + require.Equal(t, `{"query":"gmail","limit":2}`, item.Arguments) + + wire, err := json.Marshal(item) + require.NoError(t, err) + var decoded map[string]any + require.NoError(t, json.Unmarshal(wire, &decoded)) + args, ok := decoded["arguments"].(map[string]any) + require.True(t, ok, "tool_search_call arguments must remain an object") + require.Equal(t, "gmail", args["query"]) +} + +func TestResponsesResponseUnmarshal_ToolSearchObjectArguments(t *testing.T) { + var response ResponsesResponse + require.NoError(t, json.Unmarshal([]byte(`{ + "id":"response_1", + "object":"response", + "status":"completed", + "output":[{ + "type":"tool_search_call", + "id":"item_1", + "call_id":"call_1", + "arguments":{"query":"gmail"} + }] + }`), &response)) + require.Len(t, response.Output, 1) + require.Equal(t, `{"query":"gmail"}`, response.Output[0].Arguments) +} + +func TestResponsesStreamEventUnmarshal_ToolSearchObjectArguments(t *testing.T) { + var event ResponsesStreamEvent + require.NoError(t, json.Unmarshal([]byte(`{ + "type":"response.output_item.done", + "item":{ + "type":"tool_search_call", + "id":"item_1", + "call_id":"call_1", + "arguments":{"query":"gmail"} + } + }`), &event)) + require.NotNil(t, event.Item) + require.Equal(t, `{"query":"gmail"}`, event.Item.Arguments) +} diff --git a/backend/internal/pkg/apicompat/types.go b/backend/internal/pkg/apicompat/types.go index 9f3f2daa66..6cf9a2be31 100644 --- a/backend/internal/pkg/apicompat/types.go +++ b/backend/internal/pkg/apicompat/types.go @@ -353,6 +353,56 @@ func (o ResponsesOutput) MarshalJSON() ([]byte, error) { return json.Marshal(m) } +// UnmarshalJSON accepts both the Responses function-call string form and the +// tool_search_call object form for arguments. The bridge stores arguments as a +// string internally, so object arguments are retained as their raw JSON. +func (o *ResponsesOutput) UnmarshalJSON(data []byte) error { + type responsesOutputAlias ResponsesOutput + + var kind struct { + Type string `json:"type"` + } + if err := json.Unmarshal(data, &kind); err != nil { + return err + } + if kind.Type != "tool_search_call" { + var decoded responsesOutputAlias + if err := json.Unmarshal(data, &decoded); err != nil { + return err + } + *o = ResponsesOutput(decoded) + return nil + } + + var fields map[string]json.RawMessage + if err := json.Unmarshal(data, &fields); err != nil { + return err + } + arguments, hasArguments := fields["arguments"] + delete(fields, "arguments") + normalized, err := json.Marshal(fields) + if err != nil { + return err + } + + var decoded responsesOutputAlias + if err := json.Unmarshal(normalized, &decoded); err != nil { + return err + } + *o = ResponsesOutput(decoded) + if !hasArguments || string(arguments) == "null" { + return nil + } + + var argumentString string + if err := json.Unmarshal(arguments, &argumentString); err == nil { + o.Arguments = argumentString + } else { + o.Arguments = string(arguments) + } + return nil +} + // WebSearchAction describes the search action in a web_search_call output item. type WebSearchAction struct { Type string `json:"type,omitempty"` // "search" diff --git a/backend/internal/repository/account_repo.go b/backend/internal/repository/account_repo.go index c3c2f6708a..261be7c14b 100644 --- a/backend/internal/repository/account_repo.go +++ b/backend/internal/repository/account_repo.go @@ -1286,6 +1286,38 @@ func (r *accountRepository) SetRateLimited(ctx context.Context, id int64, resetA return nil } +// SetRateLimitedIfLater atomically extends an account-level rate limit. Grok +// requests may finish concurrently, so an older response must not overwrite a +// later reset boundary observed by another request or instance. +func (r *accountRepository) SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error { + now := time.Now() + updated, err := r.client.Account.Update(). + Where( + dbaccount.IDEQ(id), + dbaccount.Or( + dbaccount.RateLimitResetAtIsNil(), + dbaccount.RateLimitResetAtLT(resetAt), + ), + ). + SetRateLimitedAt(now). + SetRateLimitResetAt(resetAt). + Save(ctx) + if err != nil { + return err + } + if updated == 0 { + // This instance may not have observed the later value written elsewhere. + // Refresh its local scheduler snapshot even though no outbox event is needed. + r.syncSchedulerAccountSnapshot(ctx, id) + return nil + } + if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountChanged, &id, nil, nil); err != nil { + logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue extended rate limit failed: account=%d err=%v", id, err) + } + r.syncSchedulerAccountSnapshot(ctx, id) + return nil +} + func (r *accountRepository) SetModelRateLimit(ctx context.Context, id int64, scope string, resetAt time.Time, reason ...string) error { if scope == "" { return nil diff --git a/backend/internal/repository/account_repo_integration_test.go b/backend/internal/repository/account_repo_integration_test.go index de7fd3a9f3..f82ad4ca28 100644 --- a/backend/internal/repository/account_repo_integration_test.go +++ b/backend/internal/repository/account_repo_integration_test.go @@ -703,6 +703,25 @@ func (s *AccountRepoSuite) TestSetRateLimited() { s.Require().WithinDuration(resetAt, *got.RateLimitResetAt, time.Second) } +func (s *AccountRepoSuite) TestSetRateLimitedIfLaterDoesNotShortenReset() { + account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-rl-monotonic"}) + later := time.Now().Add(30 * time.Minute).UTC().Truncate(time.Second) + earlier := time.Now().Add(5 * time.Minute).UTC().Truncate(time.Second) + cacheRecorder := &schedulerCacheRecorder{} + s.repo.schedulerCache = cacheRecorder + + s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, later)) + s.Require().NoError(s.repo.SetRateLimitedIfLater(s.ctx, account.ID, earlier)) + + got, err := s.repo.GetByID(s.ctx, account.ID) + s.Require().NoError(err) + s.Require().NotNil(got.RateLimitResetAt) + s.Require().WithinDuration(later, *got.RateLimitResetAt, time.Second) + s.Require().Len(cacheRecorder.setAccounts, 2) + s.Require().NotNil(cacheRecorder.setAccounts[1].RateLimitResetAt) + s.Require().WithinDuration(later, *cacheRecorder.setAccounts[1].RateLimitResetAt, time.Second) +} + func (s *AccountRepoSuite) TestClearRateLimit() { account := mustCreateAccount(s.T(), s.client, &service.Account{Name: "acc-clear"}) until := time.Now().Add(1 * time.Hour) diff --git a/backend/internal/repository/api_key_repo.go b/backend/internal/repository/api_key_repo.go index d348ef29b7..ee4ed4785f 100644 --- a/backend/internal/repository/api_key_repo.go +++ b/backend/internal/repository/api_key_repo.go @@ -190,6 +190,7 @@ func (r *apiKeyRepository) GetByKeyForAuth(ctx context.Context, key string) (*se group.FieldVideoPrice480p, group.FieldVideoPrice720p, group.FieldVideoPrice1080p, + group.FieldWebSearchPricePerCall, group.FieldClaudeCodeOnly, group.FieldFallbackGroupID, group.FieldFallbackGroupIDOnInvalidRequest, @@ -943,6 +944,7 @@ func groupEntityToService(g *dbent.Group) *service.Group { VideoPrice480P: g.VideoPrice480p, VideoPrice720P: g.VideoPrice720p, VideoPrice1080P: g.VideoPrice1080p, + WebSearchPricePerCall: g.WebSearchPricePerCall, DefaultValidityDays: g.DefaultValidityDays, ClaudeCodeOnly: g.ClaudeCodeOnly, FallbackGroupID: g.FallbackGroupID, diff --git a/backend/internal/repository/group_repo.go b/backend/internal/repository/group_repo.go index 37529c60be..47efb72dd8 100644 --- a/backend/internal/repository/group_repo.go +++ b/backend/internal/repository/group_repo.go @@ -63,6 +63,7 @@ func (r *groupRepository) Create(ctx context.Context, groupIn *service.Group) er SetNillableVideoPrice480p(groupIn.VideoPrice480P). SetNillableVideoPrice720p(groupIn.VideoPrice720P). SetNillableVideoPrice1080p(groupIn.VideoPrice1080P). + SetNillableWebSearchPricePerCall(groupIn.WebSearchPricePerCall). SetDefaultValidityDays(groupIn.DefaultValidityDays). SetClaudeCodeOnly(groupIn.ClaudeCodeOnly). SetNillableFallbackGroupID(groupIn.FallbackGroupID). @@ -215,6 +216,11 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er } else { builder = builder.ClearVideoPrice1080p() } + if groupIn.WebSearchPricePerCall != nil { + builder = builder.SetWebSearchPricePerCall(*groupIn.WebSearchPricePerCall) + } else { + builder = builder.ClearWebSearchPricePerCall() + } // 处理 FallbackGroupID：nil 时清除，否则设置 if groupIn.FallbackGroupID != nil { diff --git a/backend/internal/repository/http_upstream.go b/backend/internal/repository/http_upstream.go index eac60d6268..bb079b0789 100644 --- a/backend/internal/repository/http_upstream.go +++ b/backend/internal/repository/http_upstream.go @@ -13,6 +13,7 @@ import ( "net" "net/http" "net/url" + "os" "strings" "sync" "sync/atomic" @@ -27,6 +28,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/Wei-Shaw/sub2api/internal/util/urlvalidator" + "golang.org/x/mod/semver" ) // 默认配置常量 @@ -57,6 +59,13 @@ const ( defaultOpenAIHTTP2FallbackErrorThreshold = 2 defaultOpenAIHTTP2FallbackWindow = 60 * time.Second defaultOpenAIHTTP2FallbackTTL = 10 * time.Minute + + // The Grok CLI proxy rejects requests that do not identify a supported + // client version. Keep a known-good stable version in the binary while + // allowing operators to bump it without waiting for a Sub2API release. + grokCLIProxyHost = "cli-chat-proxy.grok.com" + grokCLIStableVersion = "0.2.93" + grokCLIVersionOverride = "XAI_GROK_CLI_VERSION" ) const ( @@ -161,6 +170,7 @@ func NewHTTPUpstream(cfg *config.Config) service.HTTPUpstream { // - 调用方必须关闭 resp.Body，否则会导致 inFlight 计数泄漏 // - inFlight > 0 的客户端不会被淘汰，确保活跃请求不被中断 func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) { + applyGrokCLIProxyHeaders(req) if err := s.validateRequestHost(req); err != nil { return nil, err } @@ -207,6 +217,7 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco if profile == nil { return s.Do(req, proxyURL, accountID, accountConcurrency) } + applyGrokCLIProxyHeaders(req) upstreamProfile := service.HTTPUpstreamProfileDefault if req != nil { upstreamProfile = service.HTTPUpstreamProfileFromContext(req.Context()) @@ -250,6 +261,34 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco return resp, nil } +// applyGrokCLIProxyHeaders applies the official Grok Build client identity at +// the final shared transport boundary. Keying this behavior to the exact CLI +// proxy host keeps direct api.x.ai traffic unchanged and automatically covers +// Responses, Chat Completions, media, quota probes, and account tests. +func applyGrokCLIProxyHeaders(req *http.Request) { + if req == nil || req.URL == nil || !strings.EqualFold(strings.TrimSpace(req.URL.Hostname()), grokCLIProxyHost) { + return + } + if req.Header == nil { + req.Header = make(http.Header) + } + version := strings.TrimSpace(os.Getenv(grokCLIVersionOverride)) + if !isSupportedGrokCLIVersion(version) { + version = grokCLIStableVersion + } + req.Header.Set("X-XAI-Token-Auth", "xai-grok-cli") + req.Header.Set("x-grok-client-version", version) + req.Header.Set("User-Agent", "xai-grok-workspace/"+version) +} + +func isSupportedGrokCLIVersion(version string) bool { + canonical := "v" + version + minimum := "v" + grokCLIStableVersion + return semver.IsValid(canonical) && + semver.Canonical(canonical) == canonical && + semver.Compare(canonical, minimum) >= 0 +} + // acquireClientWithTLS 获取或创建带 TLS 指纹的客户端 func (s *httpUpstreamService) acquireClientWithTLS(proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile, upstreamProfile service.HTTPUpstreamProfile) (*upstreamClientEntry, error) { return s.getClientEntryWithTLS(proxyURL, accountID, accountConcurrency, profile, upstreamProfile, true, true) diff --git a/backend/internal/repository/http_upstream_test.go b/backend/internal/repository/http_upstream_test.go index f331e00589..cb8fe70d5b 100644 --- a/backend/internal/repository/http_upstream_test.go +++ b/backend/internal/repository/http_upstream_test.go @@ -15,6 +15,151 @@ import ( "github.com/stretchr/testify/suite" ) +func TestHTTPUpstreamDoAppliesGrokCLIIdentityBeforeOAuthRoundTrip(t *testing.T) { + t.Setenv("XAI_GROK_CLI_VERSION", "") + + for _, endpoint := range []string{"responses", "chat/completions"} { + t.Run(endpoint, func(t *testing.T) { + upstream := NewHTTPUpstream(nil) + svc, ok := upstream.(*httpUpstreamService) + require.True(t, ok) + + const accountID int64 = 4084 + isolation := svc.getIsolationMode() + profile := service.HTTPUpstreamProfileDefault + proxyKey := directProxyKey + protocolMode := svc.resolveProtocolMode(profile, proxyKey, nil) + settings := svc.resolvePoolSettings(isolation, 1) + settings = svc.applyProfilePoolSettings(settings, profile) + cacheKey := buildCacheKey(isolation, proxyKey, accountID, protocolMode) + + var capturedHeaders http.Header + svc.clients[cacheKey] = &upstreamClientEntry{ + client: &http.Client{Transport: roundTripFunc(func(req *http.Request) (*http.Response, error) { + capturedHeaders = req.Header.Clone() + statusCode := http.StatusOK + if req.Header.Get("X-XAI-Token-Auth") != "xai-grok-cli" { + statusCode = http.StatusForbidden + } + return &http.Response{ + StatusCode: statusCode, + Header: make(http.Header), + Body: http.NoBody, + Request: req, + }, nil + })}, + proxyKey: proxyKey, + poolKey: buildPoolKey(settings, protocolMode), + protocolMode: protocolMode, + } + + req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/"+endpoint, nil) + require.NoError(t, err) + req.Header.Set("User-Agent", "sub2api-grok/1.0") + + resp, err := svc.Do(req, "", accountID, 1) + require.NoError(t, err) + require.Equal(t, http.StatusOK, resp.StatusCode) + require.NoError(t, resp.Body.Close()) + + require.Equal(t, "0.2.93", capturedHeaders.Get("x-grok-client-version")) + require.Equal(t, "xai-grok-cli", capturedHeaders.Get("X-XAI-Token-Auth")) + require.Equal(t, "xai-grok-workspace/0.2.93", capturedHeaders.Get("User-Agent")) + }) + } +} + +func TestApplyGrokCLIProxyHeaders(t *testing.T) { + t.Run("uses pinned stable version for the CLI proxy", func(t *testing.T) { + t.Setenv("XAI_GROK_CLI_VERSION", "") + req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", nil) + require.NoError(t, err) + req.Header.Set("User-Agent", "sub2api-grok/1.0") + + applyGrokCLIProxyHeaders(req) + + require.Equal(t, "0.2.93", req.Header.Get("x-grok-client-version")) + require.Equal(t, "xai-grok-cli", req.Header.Get("X-XAI-Token-Auth")) + require.Equal(t, "xai-grok-workspace/0.2.93", req.Header.Get("User-Agent")) + }) + + t.Run("accepts a valid operator override", func(t *testing.T) { + t.Setenv("XAI_GROK_CLI_VERSION", "0.2.95-alpha.1") + req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/chat/completions", nil) + require.NoError(t, err) + + applyGrokCLIProxyHeaders(req) + + require.Equal(t, "0.2.95-alpha.1", req.Header.Get("x-grok-client-version")) + require.Equal(t, "xai-grok-workspace/0.2.95-alpha.1", req.Header.Get("User-Agent")) + }) + + t.Run("rejects an unsafe override", func(t *testing.T) { + t.Setenv("XAI_GROK_CLI_VERSION", "0.2.95\r\nX-Injected: true") + req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", nil) + require.NoError(t, err) + + applyGrokCLIProxyHeaders(req) + + require.Equal(t, "0.2.93", req.Header.Get("x-grok-client-version")) + require.Empty(t, req.Header.Get("X-Injected")) + }) + + t.Run("rejects an override below the supported minimum", func(t *testing.T) { + t.Setenv("XAI_GROK_CLI_VERSION", "0.2.92") + req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", nil) + require.NoError(t, err) + + applyGrokCLIProxyHeaders(req) + + require.Equal(t, "0.2.93", req.Header.Get("x-grok-client-version")) + require.Equal(t, "xai-grok-workspace/0.2.93", req.Header.Get("User-Agent")) + }) + + t.Run("rejects a prerelease override at the minimum version", func(t *testing.T) { + t.Setenv("XAI_GROK_CLI_VERSION", "0.2.93-beta.1") + req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", nil) + require.NoError(t, err) + + applyGrokCLIProxyHeaders(req) + + require.Equal(t, "0.2.93", req.Header.Get("x-grok-client-version")) + require.Equal(t, "xai-grok-workspace/0.2.93", req.Header.Get("User-Agent")) + }) + + for _, version := range []string{ + "0.2.093", + "0.2.94-alpha..1", + "0.3", + "1", + "0.2.95+build.1", + } { + t.Run("rejects invalid semver "+version, func(t *testing.T) { + t.Setenv("XAI_GROK_CLI_VERSION", version) + req, err := http.NewRequest(http.MethodPost, "https://cli-chat-proxy.grok.com/v1/responses", nil) + require.NoError(t, err) + + applyGrokCLIProxyHeaders(req) + + require.Equal(t, "0.2.93", req.Header.Get("x-grok-client-version")) + require.Equal(t, "xai-grok-workspace/0.2.93", req.Header.Get("User-Agent")) + }) + } + + t.Run("leaves direct xAI API requests unchanged", func(t *testing.T) { + t.Setenv("XAI_GROK_CLI_VERSION", "0.2.95") + req, err := http.NewRequest(http.MethodPost, "https://api.x.ai/v1/responses", nil) + require.NoError(t, err) + req.Header.Set("User-Agent", "sub2api-grok/1.0") + + applyGrokCLIProxyHeaders(req) + + require.Empty(t, req.Header.Get("x-grok-client-version")) + require.Empty(t, req.Header.Get("X-XAI-Token-Auth")) + require.Equal(t, "sub2api-grok/1.0", req.Header.Get("User-Agent")) + }) +} + // HTTPUpstreamSuite HTTP 上游服务测试套件 // 使用 testify/suite 组织测试，支持 SetupTest 初始化 type HTTPUpstreamSuite struct { diff --git a/backend/internal/server/api_contract_test.go b/backend/internal/server/api_contract_test.go index e5c050c174..a5e3fde155 100644 --- a/backend/internal/server/api_contract_test.go +++ b/backend/internal/server/api_contract_test.go @@ -366,6 +366,7 @@ func TestAPIContracts(t *testing.T) { "video_price_480p": null, "video_price_720p": null, "video_price_1080p": null, + "web_search_price_per_call": null, "allow_image_generation": false, "allow_batch_image_generation": false, "batch_image_discount_multiplier": 0, diff --git a/backend/internal/server/routes/gateway.go b/backend/internal/server/routes/gateway.go index ba5b4f61d1..7960137604 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -147,6 +147,7 @@ func RegisterGatewayRoutes( } h.Gateway.Responses(c) }) + gateway.POST("/alpha/search", h.OpenAIGateway.AlphaSearch) gateway.GET("/responses", func(c *gin.Context) { h.OpenAIGateway.ResponsesWebSocket(c) }) @@ -212,6 +213,7 @@ func RegisterGatewayRoutes( } r.POST("/responses", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, responsesHandler) r.POST("/responses/*subpath", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, responsesHandler) + r.POST("/alpha/search", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.OpenAIGateway.AlphaSearch) r.GET("/responses", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, func(c *gin.Context) { h.OpenAIGateway.ResponsesWebSocket(c) }) @@ -220,6 +222,7 @@ func RegisterGatewayRoutes( { codexDirect.POST("/responses", responsesHandler) codexDirect.POST("/responses/*subpath", responsesHandler) + codexDirect.POST("/alpha/search", h.OpenAIGateway.AlphaSearch) codexDirect.GET("/responses", func(c *gin.Context) { h.OpenAIGateway.ResponsesWebSocket(c) }) diff --git a/backend/internal/server/routes/gateway_test.go b/backend/internal/server/routes/gateway_test.go index 2779dd8f01..6b15fbfa9b 100644 --- a/backend/internal/server/routes/gateway_test.go +++ b/backend/internal/server/routes/gateway_test.go @@ -65,6 +65,36 @@ func TestGatewayRoutesOpenAIResponsesCompactPathIsRegistered(t *testing.T) { } } +func TestGatewayRoutesOpenAIAlphaSearchPathsAreRegistered(t *testing.T) { + router := newGatewayRoutesTestRouter() + registered := make(map[string]bool) + for _, route := range router.Routes() { + if route.Method == http.MethodPost { + registered[route.Path] = true + } + } + + for _, path := range []string{ + "/v1/alpha/search", + "/alpha/search", + "/backend-api/codex/alpha/search", + } { + require.True(t, registered[path], "POST %s should be registered", path) + } +} + +func TestGatewayRoutesAlphaSearchRejectsNonOpenAIGroup(t *testing.T) { + router := newGatewayRoutesTestRouter(service.PlatformGrok) + req := httptest.NewRequest(http.MethodPost, "/v1/alpha/search", strings.NewReader(`{"model":"gpt-5.6-sol"}`)) + req.Header.Set("Content-Type", "application/json") + w := httptest.NewRecorder() + + router.ServeHTTP(w, req) + + require.Equal(t, http.StatusNotFound, w.Code) + require.Contains(t, w.Body.String(), "only available for OpenAI groups") +} + func TestGatewayRoutesOpenAIImagesPathsAreRegistered(t *testing.T) { router := newGatewayRoutesTestRouter() diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index 01f3340aef..27a86a60c6 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -6,6 +6,7 @@ import ( "errors" "hash/fnv" "log/slog" + "net/url" "reflect" "sort" "strconv" @@ -1263,12 +1264,39 @@ func (a *Account) GetGrokBaseURL() string { return "" } baseURL := a.GetCredential("base_url") + if a.IsGrokOAuth() { + if strings.TrimSpace(baseURL) == "" || isOfficialGrokAPIBaseURL(baseURL) { + return xai.DefaultCLIBaseURL + } + } if baseURL != "" { return baseURL } return xai.DefaultBaseURL } +func isOfficialGrokAPIBaseURL(raw string) bool { + parsed, err := url.Parse(strings.TrimSpace(raw)) + if err != nil || parsed == nil || parsed.Opaque != "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" { + return false + } + defaultURL, err := url.Parse(xai.DefaultBaseURL) + if err != nil { + return false + } + if !strings.EqualFold(parsed.Scheme, defaultURL.Scheme) || !strings.EqualFold(parsed.Hostname(), defaultURL.Hostname()) { + return false + } + if port := parsed.Port(); port != "" { + portNumber, err := strconv.Atoi(port) + if err != nil || portNumber != 443 { + return false + } + } + path := strings.TrimRight(parsed.Path, "/") + return path == "" || path == strings.TrimRight(defaultURL.Path, "/") +} + func (a *Account) GetGrokAccessToken() string { if !a.IsGrok() { return "" diff --git a/backend/internal/service/account_base_url_test.go b/backend/internal/service/account_base_url_test.go index a132219398..0ffaa21ae4 100644 --- a/backend/internal/service/account_base_url_test.go +++ b/backend/internal/service/account_base_url_test.go @@ -4,6 +4,9 @@ package service import ( "testing" + + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/stretchr/testify/require" ) func TestGetBaseURL(t *testing.T) { @@ -158,3 +161,135 @@ func TestGetGeminiBaseURL(t *testing.T) { }) } } + +func TestGetGrokBaseURLUsesSubscriptionProxyForOAuth(t *testing.T) { + tests := []struct { + name string + account Account + expected string + }{ + { + name: "oauth without base_url uses CLI subscription proxy", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{}, + }, + expected: xai.DefaultCLIBaseURL, + }, + { + name: "oauth legacy API default is migrated at runtime to CLI subscription proxy", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": xai.DefaultBaseURL, + }, + }, + expected: xai.DefaultCLIBaseURL, + }, + { + name: "oauth legacy API default with trailing slash is migrated at runtime", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": xai.DefaultBaseURL + "/", + }, + }, + expected: xai.DefaultCLIBaseURL, + }, + { + name: "oauth legacy API root is migrated at runtime", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": "https://api.x.ai", + }, + }, + expected: xai.DefaultCLIBaseURL, + }, + { + name: "oauth legacy API root with canonical HTTPS port is migrated at runtime", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": "HTTPS://API.X.AI:443/", + }, + }, + expected: xai.DefaultCLIBaseURL, + }, + { + name: "oauth legacy API canonical port with leading zeroes is migrated at runtime", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": "https://api.x.ai:0443/v1", + }, + }, + expected: xai.DefaultCLIBaseURL, + }, + { + name: "oauth legacy API encoded version path is migrated at runtime", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": "https://api.x.ai/%76%31", + }, + }, + expected: xai.DefaultCLIBaseURL, + }, + { + name: "oauth legacy API encoded trailing slash is migrated at runtime", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": "https://api.x.ai/v1%2F", + }, + }, + expected: xai.DefaultCLIBaseURL, + }, + { + name: "oauth non-default API port remains an explicit override", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": "https://api.x.ai:8443/v1", + }, + }, + expected: "https://api.x.ai:8443/v1", + }, + { + name: "oauth explicit custom base_url remains supported", + account: Account{ + Type: AccountTypeOAuth, + Platform: PlatformGrok, + Credentials: map[string]any{ + "base_url": "https://custom.example.com/v1", + }, + }, + expected: "https://custom.example.com/v1", + }, + { + name: "API key without base_url uses official credit-backed API", + account: Account{ + Type: AccountTypeAPIKey, + Platform: PlatformGrok, + Credentials: map[string]any{}, + }, + expected: xai.DefaultBaseURL, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + require.Equal(t, tt.expected, tt.account.GetGrokBaseURL()) + }) + } +} diff --git a/backend/internal/service/account_test_service.go b/backend/internal/service/account_test_service.go index b9dae84057..e62f459e08 100644 --- a/backend/internal/service/account_test_service.go +++ b/backend/internal/service/account_test_service.go @@ -654,16 +654,10 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account return s.processOpenAIStream(c, resp.Body) } -// testGrokAccountConnection tests a Grok OAuth account through xAI's Responses API. +// testGrokAccountConnection tests a Grok OAuth or API-key account through xAI's Responses API. func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account *Account, modelID string) error { ctx := c.Request.Context() - if account.Type != AccountTypeOAuth { - return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported Grok account type: %s", account.Type)) - } - if s.grokTokenProvider == nil { - return s.sendErrorAndEnd(c, "Grok token provider not configured") - } if s.httpUpstream == nil { return s.sendErrorAndEnd(c, "HTTP upstream not configured") } @@ -676,9 +670,24 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account * testModelID = mapped } - authToken, err := s.grokTokenProvider.GetAccessToken(ctx, account) - if err != nil { - return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to get Grok access token: %s", err.Error())) + var authToken string + switch account.Type { + case AccountTypeOAuth: + if s.grokTokenProvider == nil { + return s.sendErrorAndEnd(c, "Grok token provider not configured") + } + var err error + authToken, err = s.grokTokenProvider.GetAccessToken(ctx, account) + if err != nil { + return s.sendErrorAndEnd(c, fmt.Sprintf("Failed to get Grok access token: %s", err.Error())) + } + case AccountTypeAPIKey: + authToken = strings.TrimSpace(account.GetCredential("api_key")) + if authToken == "" { + return s.sendErrorAndEnd(c, "Grok API key is missing") + } + default: + return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported Grok account type: %s", account.Type)) } apiURL, err := xai.BuildResponsesURL(account.GetGrokBaseURL()) @@ -710,7 +719,7 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account * req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json, text/event-stream") req.Header.Set("Authorization", "Bearer "+authToken) - req.Header.Set("User-Agent", "sub2api-grok/1.0") + applyGrokCLIHeaders(req.Header) proxyURL := "" if account.ProxyID != nil && account.Proxy != nil { @@ -723,10 +732,19 @@ func (s *AccountTestService) testGrokAccountConnection(c *gin.Context, account * } defer func() { _ = resp.Body.Close() }() - if snapshot := xai.ParseQuotaHeaders(resp.Header, resp.StatusCode); snapshot != nil && s.accountRepo != nil { + now := time.Now() + snapshot := parseGrokQuotaSnapshot(resp.Header, resp.StatusCode, now) + if snapshot != nil && s.accountRepo != nil { + resetAt, limited := grokRateLimitResetAt(snapshot, now) + if limited { + normalizeGrokExhaustedWindowResets(snapshot, resetAt, now) + } _ = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ grokQuotaSnapshotExtraKey: snapshot, }) + if limited { + persistGrokRateLimit(ctx, s.accountRepo, account, resetAt) + } } if resp.StatusCode != http.StatusOK { diff --git a/backend/internal/service/account_test_service_grok_test.go b/backend/internal/service/account_test_service_grok_test.go index 356fc62fb7..497224b713 100644 --- a/backend/internal/service/account_test_service_grok_test.go +++ b/backend/internal/service/account_test_service_grok_test.go @@ -3,6 +3,7 @@ package service import ( + "context" "io" "net/http" "net/http/httptest" @@ -15,6 +16,18 @@ import ( "github.com/tidwall/gjson" ) +type grokAccountTestRateLimitRepo struct { + *mockAccountRepoForGemini + rateLimitedCalls int + resetAt time.Time +} + +func (r *grokAccountTestRateLimitRepo) SetRateLimited(_ context.Context, _ int64, resetAt time.Time) error { + r.rateLimitedCalls++ + r.resetAt = resetAt + return nil +} + func TestAccountTestService_TestAccountConnection_GrokUsesXAIResponses(t *testing.T) { gin.SetMode(gin.TestMode) @@ -58,10 +71,81 @@ func TestAccountTestService_TestAccountConnection_GrokUsesXAIResponses(t *testin err := svc.TestAccountConnection(c, account.ID, "grok", "", AccountTestModeDefault) require.NoError(t, err) - require.Equal(t, "https://api.x.ai/v1/responses", upstream.lastReq.URL.String()) + require.Equal(t, "https://cli-chat-proxy.grok.com/v1/responses", upstream.lastReq.URL.String()) require.Equal(t, "Bearer grok-access-token", upstream.lastReq.Header.Get("Authorization")) + require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version")) require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String()) require.NotContains(t, rec.Body.String(), "claude") require.Contains(t, rec.Body.String(), `"model":"grok-4.3"`) require.Contains(t, rec.Body.String(), `"type":"test_complete"`) } + +func TestAccountTestService_Grok429PersistsRateLimitReset(t *testing.T) { + gin.SetMode(gin.TestMode) + + account := &Account{ + ID: 14, + Name: "grok-oauth-limited", + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Status: StatusActive, + Schedulable: true, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "grok-access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + baseRepo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}} + repo := &grokAccountTestRateLimitRepo{mockAccountRepoForGemini: baseRepo} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Retry-After": []string{"45"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)), + }} + svc := &AccountTestService{ + accountRepo: repo, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + httpUpstream: upstream, + } + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/14/test", nil) + + err := svc.TestAccountConnection(c, account.ID, "grok", "", AccountTestModeDefault) + + require.Error(t, err) + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, time.Now().Add(45*time.Second), repo.resetAt, time.Second) +} + +func TestAccountTestService_Grok429WithoutQuotaHeadersUsesFallback(t *testing.T) { + gin.SetMode(gin.TestMode) + account := &Account{ + ID: 15, Name: "grok-oauth-limited-no-headers", Platform: PlatformGrok, + Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, + Credentials: map[string]any{ + "access_token": "grok-access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + }, + } + baseRepo := &mockAccountRepoForGemini{accountsByID: map[int64]*Account{account.ID: account}} + repo := &grokAccountTestRateLimitRepo{mockAccountRepoForGemini: baseRepo} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusTooManyRequests, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"quota exhausted"}}`)), + }} + svc := &AccountTestService{ + accountRepo: repo, grokTokenProvider: NewGrokTokenProvider(repo, nil), httpUpstream: upstream, + } + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/15/test", nil) + before := time.Now() + + err := svc.TestAccountConnection(c, account.ID, "grok", "", AccountTestModeDefault) + + require.Error(t, err) + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, before.Add(grokRateLimitFallbackCooldown), repo.resetAt, time.Second) +} diff --git a/backend/internal/service/admin_group.go b/backend/internal/service/admin_group.go index c85056d623..d622d547b0 100644 --- a/backend/internal/service/admin_group.go +++ b/backend/internal/service/admin_group.go @@ -156,6 +156,7 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn videoPrice480P := normalizePrice(input.VideoPrice480P) videoPrice720P := normalizePrice(input.VideoPrice720P) videoPrice1080P := normalizePrice(input.VideoPrice1080P) + webSearchPricePerCall := normalizePrice(input.WebSearchPricePerCall) imageRateMultiplier := 1.0 if input.ImageRateMultiplier != nil { if *input.ImageRateMultiplier < 0 { @@ -287,6 +288,7 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn VideoPrice480P: videoPrice480P, VideoPrice720P: videoPrice720P, VideoPrice1080P: videoPrice1080P, + WebSearchPricePerCall: webSearchPricePerCall, ClaudeCodeOnly: input.ClaudeCodeOnly, FallbackGroupID: input.FallbackGroupID, FallbackGroupIDOnInvalidRequest: fallbackOnInvalidRequest, @@ -543,6 +545,9 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd if input.VideoPrice1080P != nil { group.VideoPrice1080P = normalizePrice(input.VideoPrice1080P) } + if input.WebSearchPricePerCall != nil { + group.WebSearchPricePerCall = normalizePrice(input.WebSearchPricePerCall) + } // Claude Code 客户端限制 if input.ClaudeCodeOnly != nil { diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index 07b85ab827..2a7125f51f 100644 --- a/backend/internal/service/admin_service.go +++ b/backend/internal/service/admin_service.go @@ -220,8 +220,10 @@ type CreateGroupInput struct { VideoPrice480P *float64 VideoPrice720P *float64 VideoPrice1080P *float64 - ClaudeCodeOnly bool // 仅允许 Claude Code 客户端 - FallbackGroupID *int64 // 降级分组 ID + // Codex alpha/search 网页搜索单次价格（USD/次，仅 openai 平台使用）；nil/负数按默认价 0.01 处理 + WebSearchPricePerCall *float64 + ClaudeCodeOnly bool // 仅允许 Claude Code 客户端 + FallbackGroupID *int64 // 降级分组 ID // 无效请求兜底分组 ID（仅 anthropic 平台使用） FallbackGroupIDOnInvalidRequest *int64 // 模型路由配置（仅 anthropic 平台使用） @@ -274,8 +276,10 @@ type UpdateGroupInput struct { VideoPrice480P *float64 VideoPrice720P *float64 VideoPrice1080P *float64 - ClaudeCodeOnly *bool // 仅允许 Claude Code 客户端 - FallbackGroupID *int64 // 降级分组 ID + // Codex alpha/search 网页搜索单次价格（USD/次）；nil 表示不修改，负数表示清除回默认价 0.01 + WebSearchPricePerCall *float64 + ClaudeCodeOnly *bool // 仅允许 Claude Code 客户端 + FallbackGroupID *int64 // 降级分组 ID // 无效请求兜底分组 ID（仅 anthropic 平台使用） FallbackGroupIDOnInvalidRequest *int64 // 模型路由配置（仅 anthropic 平台使用） diff --git a/backend/internal/service/api_key_auth_cache.go b/backend/internal/service/api_key_auth_cache.go index 11b5246a1d..0cd1c6d584 100644 --- a/backend/internal/service/api_key_auth_cache.go +++ b/backend/internal/service/api_key_auth_cache.go @@ -78,6 +78,7 @@ type APIKeyAuthGroupSnapshot struct { VideoPrice480P *float64 `json:"video_price_480p,omitempty"` VideoPrice720P *float64 `json:"video_price_720p,omitempty"` VideoPrice1080P *float64 `json:"video_price_1080p,omitempty"` + WebSearchPricePerCall *float64 `json:"web_search_price_per_call,omitempty"` ClaudeCodeOnly bool `json:"claude_code_only"` FallbackGroupID *int64 `json:"fallback_group_id,omitempty"` FallbackGroupIDOnInvalidRequest *int64 `json:"fallback_group_id_on_invalid_request,omitempty"` diff --git a/backend/internal/service/api_key_auth_cache_impl.go b/backend/internal/service/api_key_auth_cache_impl.go index 539c7375d9..5c45408659 100644 --- a/backend/internal/service/api_key_auth_cache_impl.go +++ b/backend/internal/service/api_key_auth_cache_impl.go @@ -14,7 +14,7 @@ import ( "github.com/dgraph-io/ristretto" ) -const apiKeyAuthSnapshotVersion = 14 // v14: include group video pricing fields +const apiKeyAuthSnapshotVersion = 15 // v15: include group web search per-call pricing type apiKeyAuthCacheConfig struct { l1Size int @@ -270,6 +270,7 @@ func (s *APIKeyService) snapshotFromAPIKey(ctx context.Context, apiKey *APIKey) VideoPrice480P: apiKey.Group.VideoPrice480P, VideoPrice720P: apiKey.Group.VideoPrice720P, VideoPrice1080P: apiKey.Group.VideoPrice1080P, + WebSearchPricePerCall: apiKey.Group.WebSearchPricePerCall, ClaudeCodeOnly: apiKey.Group.ClaudeCodeOnly, FallbackGroupID: apiKey.Group.FallbackGroupID, FallbackGroupIDOnInvalidRequest: apiKey.Group.FallbackGroupIDOnInvalidRequest, @@ -353,6 +354,7 @@ func (s *APIKeyService) snapshotToAPIKey(key string, snapshot *APIKeyAuthSnapsho VideoPrice480P: snapshot.Group.VideoPrice480P, VideoPrice720P: snapshot.Group.VideoPrice720P, VideoPrice1080P: snapshot.Group.VideoPrice1080P, + WebSearchPricePerCall: snapshot.Group.WebSearchPricePerCall, ClaudeCodeOnly: snapshot.Group.ClaudeCodeOnly, FallbackGroupID: snapshot.Group.FallbackGroupID, FallbackGroupIDOnInvalidRequest: snapshot.Group.FallbackGroupIDOnInvalidRequest, diff --git a/backend/internal/service/billing_service.go b/backend/internal/service/billing_service.go index 4fc1d7c214..44dae17db4 100644 --- a/backend/internal/service/billing_service.go +++ b/backend/internal/service/billing_service.go @@ -564,15 +564,19 @@ func (s *BillingService) initFallbackPricing() { s.fallbackPrices["grok-4.3"] = &ModelPricing{ InputPricePerToken: 1.25e-6, OutputPricePerToken: 2.5e-6, - CacheReadPricePerToken: 0, + CacheReadPricePerToken: 0.2e-6, SupportsCacheBreakdown: false, LongContextInputThreshold: 1000000, LongContextInputMultiplier: 1, } - // xAI Grok Build 0.1 (official docs: $1 input / $2 output per MTok) + // xAI Grok Build 0.1 (official docs: $1 input / $0.20 cached input / + // $2 output per MTok). Composer is available only through Grok Build and + // has no standalone public API rate card, so its aliases use this coding + // model rate instead of silently billing at zero. s.fallbackPrices["grok-build-0.1"] = &ModelPricing{ InputPricePerToken: 1e-6, OutputPricePerToken: 2e-6, + CacheReadPricePerToken: 0.2e-6, SupportsCacheBreakdown: false, } } @@ -746,9 +750,14 @@ func (s *BillingService) getFallbackPricing(model string) *ModelPricing { switch modelLower { case "grok", "grok-latest", "grok-4.5", "grok-4.5-latest", "grok-build-latest": return s.fallbackPrices["grok-4.5"] - case "grok-4.3": + case "grok-4.3", + "grok-4.20-0309-reasoning", + "grok-4.20-0309-non-reasoning", + "grok-4.20-multi-agent-0309", + "grok-4.20-reasoning", + "grok-4.20-non-reasoning": return s.fallbackPrices["grok-4.3"] - case "grok-build", "grok-build-0.1": + case "grok-build", "grok-build-0.1", "grok-composer", "grok-composer-2.5-fast", "composer-2.5": return s.fallbackPrices["grok-build-0.1"] } @@ -1361,8 +1370,36 @@ const ( defaultGrokImagineVideo15Price480P = 0.08 defaultGrokImagineVideo15Price720P = 0.14 defaultGrokImagineVideo15Price1080P = 0.25 + + // Codex alpha/search 网页搜索单次默认价：OpenAI 官方 web search 定价 $10/1000 次。 + defaultWebSearchPricePerCall = 0.01 ) +// CalculateWebSearchCost 计算 Codex alpha/search 网页搜索按次费用。 +// callCount: 搜索调用次数（每次请求为 1） +// groupPrice: 分组配置的单次价格（nil 表示使用默认价 0.01；0 表示免费） +// rateMultiplier: 分组费率倍数 +func (s *BillingService) CalculateWebSearchCost(callCount int, groupPrice *float64, rateMultiplier float64) *CostBreakdown { + if callCount <= 0 { + return &CostBreakdown{} + } + unitPrice := defaultWebSearchPricePerCall + if groupPrice != nil && *groupPrice >= 0 { + unitPrice = *groupPrice + } + totalCost := unitPrice * float64(callCount) + + // 应用倍率（保存时强制 > 0；负数按 0 处理避免按 1x 误扣） + if rateMultiplier < 0 { + rateMultiplier = 0 + } + return &CostBreakdown{ + TotalCost: totalCost, + ActualCost: totalCost * rateMultiplier, + BillingMode: string(BillingModePerRequest), + } +} + // CalculateImageCost 计算图片生成费用 // model: 请求的模型名称（用于获取 LiteLLM 默认价格） // imageSize: 图片尺寸 "1K", "2K", "4K" diff --git a/backend/internal/service/billing_service_test.go b/backend/internal/service/billing_service_test.go index 8c87a21616..9f7cc16e8b 100644 --- a/backend/internal/service/billing_service_test.go +++ b/backend/internal/service/billing_service_test.go @@ -1056,6 +1056,58 @@ func TestGetModelPricing_Grok45OfficialFallback(t *testing.T) { } } +func TestGetModelPricing_GrokCatalogFallbacks(t *testing.T) { + svc := newTestBillingService() + + tests := []struct { + name string + models []string + input float64 + cacheRead float64 + output float64 + }{ + { + name: "Grok 4.3 family", + models: []string{ + "grok-4.3", + "grok-4.20-0309-reasoning", + "grok-4.20-0309-non-reasoning", + "grok-4.20-multi-agent-0309", + "grok-4.20-reasoning", + "grok-4.20-non-reasoning", + }, + input: 1.25e-6, + cacheRead: 0.2e-6, + output: 2.5e-6, + }, + { + name: "Grok coding and Composer family", + models: []string{ + "grok-build", + "grok-build-0.1", + "grok-composer", + "grok-composer-2.5-fast", + "composer-2.5", + }, + input: 1e-6, + cacheRead: 0.2e-6, + output: 2e-6, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + for _, model := range tt.models { + pricing, err := svc.GetModelPricing(model) + require.NoError(t, err, "model %s", model) + require.InDelta(t, tt.input, pricing.InputPricePerToken, 1e-12, "model %s input", model) + require.InDelta(t, tt.cacheRead, pricing.CacheReadPricePerToken, 1e-12, "model %s cached input", model) + require.InDelta(t, tt.output, pricing.OutputPricePerToken, 1e-12, "model %s output", model) + } + }) + } +} + func TestCalculateCost_SupportsCacheBreakdown(t *testing.T) { svc := &BillingService{ cfg: &config.Config{}, diff --git a/backend/internal/service/grok_media.go b/backend/internal/service/grok_media.go index 154e3003ef..01101ef6b5 100644 --- a/backend/internal/service/grok_media.go +++ b/backend/internal/service/grok_media.go @@ -333,7 +333,7 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( } upstreamReq.Header.Set("Authorization", "Bearer "+token) upstreamReq.Header.Set("Accept", "application/json") - upstreamReq.Header.Set("User-Agent", "sub2api-grok/1.0") + applyGrokCLIHeaders(upstreamReq.Header) if endpoint.RequiresRequestBody() { contentType = strings.TrimSpace(contentType) if contentType == "" { @@ -357,11 +357,10 @@ func (s *OpenAIGatewayService) ForwardGrokMedia( requestIDHeader := firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")) requestModel := requestInfo.Model if resp.StatusCode >= 400 { - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) return s.handleGrokMediaErrorResponse(ctx, resp, c, account, requestIDHeader, requestModel) } - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError) if err != nil { return nil, err @@ -564,6 +563,9 @@ func (s *OpenAIGatewayService) handleGrokMediaErrorResponse( requestedModel string, ) (*OpenAIForwardResult, error) { body := s.readUpstreamErrorBody(resp) + // Reconcile readiness before configurable passthrough branches can return; + // otherwise a Grok 429 can remain schedulable. + s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body) upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(body))) if upstreamMsg == "" { upstreamMsg = fmt.Sprintf("xAI upstream returned status %d", resp.StatusCode) @@ -609,7 +611,6 @@ func (s *OpenAIGatewayService) handleGrokMediaErrorResponse( return nil, fmt.Errorf("upstream error: %d (not in custom error codes) message=%s", resp.StatusCode, upstreamMsg) } - s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body) kind := "http_error" if s.shouldFailoverUpstreamError(resp.StatusCode) { kind = "failover" diff --git a/backend/internal/service/grok_oauth_service.go b/backend/internal/service/grok_oauth_service.go index acad39985f..0a610ef811 100644 --- a/backend/internal/service/grok_oauth_service.go +++ b/backend/internal/service/grok_oauth_service.go @@ -235,7 +235,7 @@ func (s *GrokOAuthService) BuildAccountCredentials(tokenInfo *GrokTokenInfo) map if tokenInfo.EntitlementStatus != "" { creds["entitlement_status"] = tokenInfo.EntitlementStatus } - creds["base_url"] = xai.DefaultBaseURL + creds["base_url"] = xai.DefaultCLIBaseURL return creds } diff --git a/backend/internal/service/grok_oauth_service_test.go b/backend/internal/service/grok_oauth_service_test.go index d0caa5e527..f042c96337 100644 --- a/backend/internal/service/grok_oauth_service_test.go +++ b/backend/internal/service/grok_oauth_service_test.go @@ -5,6 +5,7 @@ package service import ( "context" "testing" + "time" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/stretchr/testify/require" @@ -66,3 +67,15 @@ func TestGrokOAuthServiceExchangeCodeRequiresStateForCallbackURLAndConsumesSessi require.Contains(t, err.Error(), "GROK_OAUTH_SESSION_NOT_FOUND") require.Zero(t, client.exchangeCalls) } + +func TestGrokOAuthServiceBuildAccountCredentialsDefaultsToSubscriptionProxy(t *testing.T) { + svc := NewGrokOAuthService(nil, &grokOAuthClientStub{}) + defer svc.Stop() + + credentials := svc.BuildAccountCredentials(&GrokTokenInfo{ + AccessToken: "access-token", + ExpiresAt: time.Now().Add(time.Hour).Unix(), + }) + + require.Equal(t, xai.DefaultCLIBaseURL, credentials["base_url"]) +} diff --git a/backend/internal/service/grok_quota_service.go b/backend/internal/service/grok_quota_service.go index 06d90a609e..17e91dac6f 100644 --- a/backend/internal/service/grok_quota_service.go +++ b/backend/internal/service/grok_quota_service.go @@ -82,7 +82,7 @@ func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*Gr req.Header.Set("Authorization", "Bearer "+token) req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json") - req.Header.Set("User-Agent", "sub2api-grok-quota-probe/1.0") + applyGrokCLIHeaders(req.Header) resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 1)) if err != nil { @@ -91,9 +91,16 @@ func (s *GrokQuotaService) ProbeUsage(ctx context.Context, accountID int64) (*Gr defer func() { _ = resp.Body.Close() }() snapshot := xai.ObserveQuotaHeaders(resp.Header, resp.StatusCode, "active_probe") + resetAt, limited := grokRateLimitResetAt(snapshot, time.Now()) + if limited { + normalizeGrokExhaustedWindowResets(snapshot, resetAt, time.Now()) + } _ = s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ grokQuotaSnapshotExtraKey: snapshot, }) + if limited { + persistGrokRateLimit(ctx, s.accountRepo, account, resetAt) + } result := &GrokQuotaProbeResult{ Source: "active_probe", diff --git a/backend/internal/service/grok_quota_service_test.go b/backend/internal/service/grok_quota_service_test.go index fe1a00aa89..2248674899 100644 --- a/backend/internal/service/grok_quota_service_test.go +++ b/backend/internal/service/grok_quota_service_test.go @@ -19,6 +19,10 @@ import ( type grokQuotaAccountRepo struct { *mockAccountRepoForPlatform updates map[int64]map[string]any + updateCalls int + rateLimitedCalls int + lastRateLimitedID int64 + lastRateLimitResetAt time.Time tempUnschedCalls int lastTempUnschedID int64 lastTempUnschedUntil time.Time @@ -26,6 +30,7 @@ type grokQuotaAccountRepo struct { } func (r *grokQuotaAccountRepo) UpdateExtra(_ context.Context, id int64, updates map[string]any) error { + r.updateCalls++ if r.updates == nil { r.updates = make(map[int64]map[string]any) } @@ -33,6 +38,17 @@ func (r *grokQuotaAccountRepo) UpdateExtra(_ context.Context, id int64, updates return nil } +func (r *grokQuotaAccountRepo) SetRateLimited(_ context.Context, id int64, resetAt time.Time) error { + r.rateLimitedCalls++ + r.lastRateLimitedID = id + r.lastRateLimitResetAt = resetAt + return nil +} + +func (r *grokQuotaAccountRepo) SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error { + return r.SetRateLimited(ctx, id, resetAt) +} + func (r *grokQuotaAccountRepo) SetTempUnschedulable(_ context.Context, id int64, until time.Time, reason string) error { r.tempUnschedCalls++ r.lastTempUnschedID = id @@ -96,8 +112,9 @@ func TestGrokQuotaServiceProbeUsageStoresHeaders(t *testing.T) { require.NotNil(t, result.Snapshot.Requests) require.EqualValues(t, 10, *result.Snapshot.Requests.Limit) require.EqualValues(t, 7, *result.Snapshot.Requests.Remaining) - require.Equal(t, "https://api.x.ai/v1/responses", upstream.lastReq.URL.String()) + require.Equal(t, "https://cli-chat-proxy.grok.com/v1/responses", upstream.lastReq.URL.String()) require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization")) + require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version")) require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.lastBody, "model").String()) require.Contains(t, string(upstream.lastBody), `"max_output_tokens":1`) require.Contains(t, string(upstream.lastBody), `"store":false`) @@ -285,6 +302,10 @@ func TestGrokQuotaServiceProbeUsageReturnsRateLimitedSnapshot(t *testing.T) { require.NotNil(t, result.Snapshot) require.NotNil(t, result.Snapshot.RetryAfterSeconds) require.Equal(t, 45, *result.Snapshot.RetryAfterSeconds) + require.Equal(t, 1, repo.rateLimitedCalls) + require.Equal(t, account.ID, repo.lastRateLimitedID) + require.WithinDuration(t, time.Now().Add(45*time.Second), repo.lastRateLimitResetAt, time.Second) + require.Zero(t, repo.tempUnschedCalls) } func TestGrokQuotaServiceResetQuotaUnsupported(t *testing.T) { diff --git a/backend/internal/service/group.go b/backend/internal/service/group.go index a61e356a01..346eef7b95 100644 --- a/backend/internal/service/group.go +++ b/backend/internal/service/group.go @@ -50,6 +50,9 @@ type Group struct { VideoPrice480P *float64 VideoPrice720P *float64 VideoPrice1080P *float64 + // Codex alpha/search 网页搜索单次价格（USD/次，仅 openai 平台使用）； + // nil 表示使用默认价 defaultWebSearchPricePerCall（官方 $10/1000 次）。 + WebSearchPricePerCall *float64 // Claude Code 客户端限制 ClaudeCodeOnly bool diff --git a/backend/internal/service/media_price_config.go b/backend/internal/service/media_price_config.go index ed84998906..0a583b6c72 100644 --- a/backend/internal/service/media_price_config.go +++ b/backend/internal/service/media_price_config.go @@ -29,3 +29,10 @@ func videoPriceConfigFromAPIKey(apiKey *APIKey) *VideoPriceConfig { func apiKeyHasConfiguredVideoPrice(apiKey *APIKey, resolution string) bool { return apiKey != nil && apiKey.Group != nil && apiKey.Group.GetVideoPrice(resolution) != nil } + +func webSearchPricePerCallFromAPIKey(apiKey *APIKey) *float64 { + if apiKey == nil || apiKey.Group == nil { + return nil + } + return apiKey.Group.WebSearchPricePerCall +} diff --git a/backend/internal/service/openai_alpha_search.go b/backend/internal/service/openai_alpha_search.go new file mode 100644 index 0000000000..50d37f95bb --- /dev/null +++ b/backend/internal/service/openai_alpha_search.go @@ -0,0 +1,165 @@ +package service + +import ( + "bytes" + "context" + "fmt" + "io" + "net/http" + "net/url" + "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" +) + +const ( + chatgptCodexAlphaSearchURL = "https://chatgpt.com/backend-api/codex/alpha/search" + openAIPlatformAlphaSearchURL = "https://api.openai.com/v1/alpha/search" +) + +// ForwardAlphaSearch proxies Codex standalone web search without binding the +// evolving alpha request or response schema. +// +// 返回值约定：仅当上游返回 2xx（一次真实成功的搜索）时返回非 nil 的 +// *OpenAIForwardResult（WebSearchCalls=1，供按次计费）；上游错误被原样透传 +// 给客户端时返回 (nil, nil)，不产生计费。 +func (s *OpenAIGatewayService) ForwardAlphaSearch(ctx context.Context, c *gin.Context, account *Account, body []byte) (*OpenAIForwardResult, error) { + if s == nil || c == nil || account == nil { + return nil, fmt.Errorf("service, context, and account are required") + } + modelResult := gjson.GetBytes(body, "model") + requestedModel := strings.TrimSpace(modelResult.String()) + if modelResult.Type != gjson.String || requestedModel == "" { + return nil, fmt.Errorf("model is required") + } + + upstreamModel := normalizeOpenAIModelForUpstream(account, account.GetMappedModel(requestedModel)) + if upstreamModel != "" && upstreamModel != requestedModel { + body = ReplaceModelInBody(body, upstreamModel) + } + + token, _, err := s.GetAccessToken(ctx, account) + if err != nil { + return nil, err + } + + req, err := s.buildOpenAIAlphaSearchRequest(ctx, c, account, body, token) + if err != nil { + return nil, err + } + + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + upstreamStart := time.Now() + resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) + SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds()) + if err != nil { + return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, true) + } + defer func() { _ = resp.Body.Close() }() + + respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError) + if err != nil { + return nil, fmt.Errorf("read alpha search response: %w", err) + } + + if resp.StatusCode >= http.StatusBadRequest { + upstreamMessage := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) + if s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMessage, respBody) { + resp.Body = io.NopCloser(bytes.NewReader(respBody)) + s.handleFailoverSideEffects(ctx, resp, account, respBody, upstreamModel) + return nil, &UpstreamFailoverError{ + StatusCode: resp.StatusCode, + ResponseBody: respBody, + RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + } + } + } + + if !account.IsShadow() { + s.UpdateCodexUsageSnapshotFromHeaders(ctx, account.ID, resp.Header) + } + writeOpenAIPassthroughResponseHeaders(c.Writer.Header(), resp.Header, s.responseHeaderFilter) + contentType := resp.Header.Get("Content-Type") + if contentType == "" { + contentType = "application/json" + } + c.Data(resp.StatusCode, contentType, respBody) + if resp.StatusCode < http.StatusOK || resp.StatusCode >= http.StatusMultipleChoices { + // 非 2xx（错误/重定向）已原样透传给客户端：不是一次成功的搜索，不计费。 + return nil, nil + } + return &OpenAIForwardResult{ + RequestID: strings.TrimSpace(resp.Header.Get("x-request-id")), + Model: requestedModel, + UpstreamModel: upstreamModel, + Duration: time.Since(upstreamStart), + WebSearchCalls: 1, + }, nil +} + +func (s *OpenAIGatewayService) buildOpenAIAlphaSearchRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token string) (*http.Request, error) { + clientBeta := "" + if c != nil { + clientBeta = c.GetHeader("OpenAI-Beta") + } + req, err := s.buildUpstreamRequestOpenAIPassthrough(ctx, c, account, body, token) + if err != nil { + return nil, err + } + + targetURL, err := s.openAIAlphaSearchURL(account) + if err != nil { + return nil, err + } + parsedURL, err := url.Parse(targetURL) + if err != nil { + return nil, fmt.Errorf("parse alpha search URL: %w", err) + } + if c != nil && c.Request != nil && c.Request.URL != nil { + query := parsedURL.Query() + for key, values := range c.Request.URL.Query() { + for _, value := range values { + query.Add(key, value) + } + } + parsedURL.RawQuery = query.Encode() + } + req.URL = parsedURL + req.Header.Set("Accept", "application/json") + if clientBeta == "" { + req.Header.Del("OpenAI-Beta") + } + if version := strings.TrimSpace(c.GetHeader("Version")); version != "" { + req.Header.Set("Version", version) + } else if account.Type == AccountTypeOAuth { + req.Header.Set("Version", codexCLIVersion) + } + return req, nil +} + +func (s *OpenAIGatewayService) openAIAlphaSearchURL(account *Account) (string, error) { + if account == nil { + return "", fmt.Errorf("account is required") + } + switch account.Type { + case AccountTypeOAuth: + return chatgptCodexAlphaSearchURL, nil + case AccountTypeAPIKey: + baseURL := account.GetOpenAIBaseURL() + if baseURL == "" { + return openAIPlatformAlphaSearchURL, nil + } + validatedURL, err := s.validateUpstreamBaseURL(baseURL) + if err != nil { + return "", err + } + return buildOpenAIEndpointURL(validatedURL, "/v1/alpha/search"), nil + default: + return "", fmt.Errorf("unsupported OpenAI account type: %s", account.Type) + } +} diff --git a/backend/internal/service/openai_alpha_search_billing_test.go b/backend/internal/service/openai_alpha_search_billing_test.go new file mode 100644 index 0000000000..7151725763 --- /dev/null +++ b/backend/internal/service/openai_alpha_search_billing_test.go @@ -0,0 +1,101 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/stretchr/testify/require" +) + +func TestCalculateWebSearchCostDefaultAndOverride(t *testing.T) { + t.Parallel() + s := &BillingService{} + + // 默认价：官方 $10/1000 次 = 0.01/次 + cost := s.CalculateWebSearchCost(1, nil, 1.0) + require.InDelta(t, 0.01, cost.TotalCost, 1e-12) + require.InDelta(t, 0.01, cost.ActualCost, 1e-12) + require.Equal(t, string(BillingModePerRequest), cost.BillingMode) + + // 分组覆盖价 + 倍率 + cost = s.CalculateWebSearchCost(1, float64Ptr(0.02), 2.5) + require.InDelta(t, 0.02, cost.TotalCost, 1e-12) + require.InDelta(t, 0.05, cost.ActualCost, 1e-12) + + // 0 = 免费（区别于 nil = 默认价） + cost = s.CalculateWebSearchCost(1, float64Ptr(0), 3.0) + require.Zero(t, cost.TotalCost) + require.Zero(t, cost.ActualCost) + + // 负数倍率按 0 处理，避免按 1x 误扣 + cost = s.CalculateWebSearchCost(1, nil, -1) + require.InDelta(t, 0.01, cost.TotalCost, 1e-12) + require.Zero(t, cost.ActualCost) + + // 次数 <= 0 不产生费用 + cost = s.CalculateWebSearchCost(0, float64Ptr(0.02), 1.0) + require.Zero(t, cost.TotalCost) + require.Empty(t, cost.BillingMode) +} + +func TestCalculateOpenAIRecordUsageCostWebSearchPerCall(t *testing.T) { + t.Parallel() + svc := &OpenAIGatewayService{billingService: &BillingService{}} + groupID := int64(11) + + // 分组未配置单价：默认 0.01。按次搜索使用不含高峰因子的基础倍率（第 4 个倍率参数 2.0）， + // 即使 token 倍率（含高峰，3.0）更高也不采用。 + apiKey := &APIKey{ID: 1, GroupID: &groupID, Group: &Group{ID: groupID, Platform: PlatformOpenAI}} + result := &OpenAIForwardResult{Model: "gpt-5.6-sol", UpstreamModel: "gpt-5.6-sol", WebSearchCalls: 1} + cost, err := svc.calculateOpenAIRecordUsageCost(context.Background(), result, apiKey, []string{"gpt-5.6-sol"}, 3.0, 1.0, 1.0, 2.0, UsageTokens{}, "", false) + require.NoError(t, err) + require.Equal(t, string(BillingModePerRequest), cost.BillingMode) + require.InDelta(t, 0.01, cost.TotalCost, 1e-12) + require.InDelta(t, 0.02, cost.ActualCost, 1e-12) + + // 分组配置单价 0.005 + apiKey.Group.WebSearchPricePerCall = float64Ptr(0.005) + cost, err = svc.calculateOpenAIRecordUsageCost(context.Background(), result, apiKey, []string{"gpt-5.6-sol"}, 1.0, 1.0, 1.0, 1.0, UsageTokens{}, "", false) + require.NoError(t, err) + require.InDelta(t, 0.005, cost.TotalCost, 1e-12) + require.InDelta(t, 0.005, cost.ActualCost, 1e-12) + + // WebSearchCalls = 0 时不得走按次分支（无定价数据会返回 pricing 错误， + // 证明回落到了 token 路径而不是被按次分支吞掉）。 + result.WebSearchCalls = 0 + _, err = svc.calculateOpenAIRecordUsageCost(context.Background(), result, apiKey, []string{"gpt-5.6-sol"}, 1.0, 1.0, 1.0, 1.0, UsageTokens{InputTokens: 10}, "", false) + require.Error(t, err) +} + +func TestAPIKeyService_SnapshotRoundTrip_PreservesWebSearchPricePerCall(t *testing.T) { + svc := NewAPIKeyService(nil, nil, nil, nil, nil, nil, &config.Config{}) + groupID := int64(9) + apiKey := &APIKey{ + ID: 1, + UserID: 2, + GroupID: &groupID, + Key: "k-websearch", + Status: StatusActive, + User: &User{ID: 2, Status: StatusActive, Role: RoleUser}, + Group: &Group{ + ID: groupID, + Name: "openai", + Platform: PlatformOpenAI, + Status: StatusActive, + SubscriptionType: SubscriptionTypeStandard, + RateMultiplier: 1, + WebSearchPricePerCall: float64Ptr(0.008), + }, + } + + snapshot := svc.snapshotFromAPIKey(context.Background(), apiKey) + roundTrip := svc.snapshotToAPIKey(apiKey.Key, snapshot) + + require.NotNil(t, roundTrip) + require.NotNil(t, roundTrip.Group) + require.NotNil(t, roundTrip.Group.WebSearchPricePerCall) + require.InDelta(t, 0.008, *roundTrip.Group.WebSearchPricePerCall, 1e-12) +} diff --git a/backend/internal/service/openai_alpha_search_test.go b/backend/internal/service/openai_alpha_search_test.go new file mode 100644 index 0000000000..458dbe9b5b --- /dev/null +++ b/backend/internal/service/openai_alpha_search_test.go @@ -0,0 +1,145 @@ +package service + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestForwardAlphaSearchOAuthPreservesWire(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{ + "id":"search-session", + "model":"gpt-5.6-sol", + "reasoning":{"effort":"max","context":"all_turns"}, + "input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"latest news"}]}], + "commands":{"search_query":[{"q":"OpenAI news","recency":1}]}, + "settings":{"allowed_callers":["direct"],"external_web_access":true}, + "max_output_tokens":2000, + "future_field":{"keep":true} + }`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search?feature=standalone", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + c.Request.Header.Set("User-Agent", codexCLIUserAgent) + c.Request.Header.Set("Originator", "codex_cli_rs") + c.Request.Header.Set("Version", "0.144.1") + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"encrypted_output":"ciphertext","output":"search result"}`)), + }} + service := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ + ID: 42, + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "oauth-token", + "chatgpt_account_id": "chatgpt-account", + }, + } + + result, err := service.ForwardAlphaSearch(context.Background(), c, account, body) + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 1, result.WebSearchCalls) + require.Equal(t, "gpt-5.6-sol", result.Model) + require.Equal(t, http.StatusOK, recorder.Code) + require.JSONEq(t, `{"encrypted_output":"ciphertext","output":"search result"}`, recorder.Body.String()) + require.Equal(t, chatgptCodexAlphaSearchURL+"?feature=standalone", upstream.lastReq.URL.String()) + require.Equal(t, "chatgpt.com", upstream.lastReq.Host) + require.Equal(t, "Bearer oauth-token", upstream.lastReq.Header.Get("Authorization")) + require.Equal(t, "chatgpt-account", upstream.lastReq.Header.Get("chatgpt-account-id")) + require.Equal(t, "application/json", upstream.lastReq.Header.Get("Accept")) + require.Equal(t, "0.144.1", upstream.lastReq.Header.Get("Version")) + require.Empty(t, upstream.lastReq.Header.Get("OpenAI-Beta")) + require.JSONEq(t, string(body), string(upstream.lastBody)) +} + +func TestForwardAlphaSearchAPIKeyMapsModelAndPassesThroughError(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{"search_query":[{"q":"news"}]}}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/alpha/search", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + upstreamBody := `{"error":{"type":"invalid_request_error","message":"bad search"}}` + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(upstreamBody)), + }} + service := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ + ID: 7, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://compat.example/v4", + "model_mapping": map[string]any{ + "gpt-5.6-sol": "upstream-5.6", + }, + }, + } + + result, err := service.ForwardAlphaSearch(context.Background(), c, account, body) + + require.NoError(t, err) + // 上游错误透传不是一次成功的搜索：不返回 result、不产生按次计费。 + require.Nil(t, result) + require.Equal(t, http.StatusBadRequest, recorder.Code) + require.JSONEq(t, upstreamBody, recorder.Body.String()) + require.Equal(t, "https://compat.example/v4/alpha/search", upstream.lastReq.URL.String()) + require.Equal(t, "Bearer sk-test", upstream.lastReq.Header.Get("Authorization")) + require.Equal(t, "upstream-5.6", gjson.GetBytes(upstream.lastBody, "model").String()) + require.True(t, gjson.GetBytes(upstream.lastBody, "commands.search_query").IsArray()) +} + +func TestForwardAlphaSearchReturnsFailoverBeforeWriting(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"id":"search-session","model":"gpt-5.6-sol","commands":{}}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/alpha/search", bytes.NewReader(body)) + + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)), + }} + service := &OpenAIGatewayService{cfg: &config.Config{}, httpUpstream: upstream} + account := &Account{ + ID: 8, + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Credentials: map[string]any{ + "api_key": "sk-test", + }, + } + + result, err := service.ForwardAlphaSearch(context.Background(), c, account, body) + + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode) + require.Equal(t, openAIPlatformAlphaSearchURL, upstream.lastReq.URL.String()) + require.False(t, c.Writer.Written()) + require.Empty(t, recorder.Body.String()) +} diff --git a/backend/internal/service/openai_codex_function_call_id_test.go b/backend/internal/service/openai_codex_function_call_id_test.go index 2ac59e0520..edfcac1d71 100644 --- a/backend/internal/service/openai_codex_function_call_id_test.go +++ b/backend/internal/service/openai_codex_function_call_id_test.go @@ -114,14 +114,15 @@ func TestFilterCodexInput_OutputTypeKeepsItemID(t *testing.T) { require.Equal(t, "o1", out["id"], "output item id should be preserved") } -// TestFilterCodexInput_NonToolCallItemKeepsID ensures non-tool-call items -// (e.g. message) still keep their id when PreserveReferences is true. +// TestFilterCodexInput_NonToolCallItemKeepsID ensures items subject to neither +// the fc* (call-input) nor the msg* (message) prefix rule still keep their id +// when PreserveReferences is true. +// message is covered separately in openai_codex_message_item_id_test.go (#3981). func TestFilterCodexInput_NonToolCallItemKeepsID(t *testing.T) { input := []any{ map[string]any{ - "type": "message", - "id": "item_msg_001", - "role": "user", + "type": "web_search_call", + "id": "ws_001", }, } @@ -130,7 +131,7 @@ func TestFilterCodexInput_NonToolCallItemKeepsID(t *testing.T) { }) require.Len(t, filtered, 1) - msg, ok := filtered[0].(map[string]any) + item, ok := filtered[0].(map[string]any) require.True(t, ok) - require.Equal(t, "item_msg_001", msg["id"], "non-tool-call items keep their id in preserve mode") + require.Equal(t, "ws_001", item["id"], "unconstrained items keep their id in preserve mode") } diff --git a/backend/internal/service/openai_codex_identity.go b/backend/internal/service/openai_codex_identity.go index 68d4105d84..47bc19ce0c 100644 --- a/backend/internal/service/openai_codex_identity.go +++ b/backend/internal/service/openai_codex_identity.go @@ -11,12 +11,32 @@ import ( // 若请求携带 version 且低于该值，上游直接 404（issue #3901，2026-07 实测）。 const codexUpstreamMinVersion = "0.144.0" +// ensureCodexIdentityHeaders 补齐 OAuth（ChatGPT 内部接口）出站请求所需的 Codex 身份头。 +// 已有 User-Agent 与 version 保持不变，交给紧随其后的 enforceCodexIdentityHeaders +// 做官方身份配对与最低版本校正。 +func ensureCodexIdentityHeaders(h http.Header) { + if h == nil { + return + } + if strings.TrimSpace(h.Get("user-agent")) == "" { + h.Set("user-agent", codexCLIUserAgent) + } + if strings.TrimSpace(h.Get("originator")) == "" { + h.Set("originator", "codex_cli_rs") + } + if strings.TrimSpace(h.Get("version")) == "" { + h.Set("version", codexCLIVersion) + } + h.Set("OpenAI-Beta", "responses=experimental") +} + // enforceCodexIdentityHeaders 收口 OAuth（ChatGPT 内部接口）出站请求的客户端身份头。 // 上游要求 originator 与 User-Agent 首段配套且为官方客户端标识，version 头（若携带） // 不低于 0.144.0，任一不满足即 404（issue #3901）。以最终 User-Agent 为准推导配套 // originator；推导不出官方身份（第三方 UA / UA 缺失）时整体回退为默认 Codex CLI 身份。 // -// 仅对携带 originator 的请求生效——compat messages bridge 故意不带 originator，保持原样。 +// 仅对携带 originator 的请求生效；需要从缺失身份头恢复的调用方应先调用 +// ensureCodexIdentityHeaders。 // 必须在所有 User-Agent 改写（自定义 UA / ForceCodexCLI / 浏览器 UA 兜底）之后调用。 func enforceCodexIdentityHeaders(h http.Header) { if h == nil || h.Get("originator") == "" { diff --git a/backend/internal/service/openai_codex_identity_test.go b/backend/internal/service/openai_codex_identity_test.go index 7d2c6d8520..ecb6eb7f6f 100644 --- a/backend/internal/service/openai_codex_identity_test.go +++ b/backend/internal/service/openai_codex_identity_test.go @@ -7,6 +7,36 @@ import ( "github.com/stretchr/testify/require" ) +func TestEnsureCodexIdentityHeaders(t *testing.T) { + t.Run("补齐缺失身份头", func(t *testing.T) { + h := make(http.Header) + + ensureCodexIdentityHeaders(h) + enforceCodexIdentityHeaders(h) + + require.Equal(t, "codex_cli_rs", h.Get("originator")) + require.Equal(t, codexCLIUserAgent, h.Get("user-agent")) + require.Equal(t, codexCLIVersion, h.Get("version")) + require.Equal(t, "responses=experimental", h.Get("OpenAI-Beta")) + }) + + t.Run("保留已有官方UA和合法version并重新配对", func(t *testing.T) { + const tuiUA = "codex-tui/9.9.9 (Mac OS X 14.0; arm64) iTerm (codex-tui; 9.9.9)" + h := make(http.Header) + h.Set("user-agent", tuiUA) + h.Set("version", "9.9.9") + h.Set("OpenAI-Beta", "assistants=v2") + + ensureCodexIdentityHeaders(h) + enforceCodexIdentityHeaders(h) + + require.Equal(t, "codex-tui", h.Get("originator")) + require.Equal(t, tuiUA, h.Get("user-agent")) + require.Equal(t, "9.9.9", h.Get("version")) + require.Equal(t, "responses=experimental", h.Get("OpenAI-Beta")) + }) +} + func TestEnforceCodexIdentityHeaders(t *testing.T) { const tuiUA = "codex-tui/0.140.2 (Mac OS X 14.0; arm64) iTerm (codex-tui; 0.140.2)" @@ -102,13 +132,14 @@ func TestEnforceCodexIdentityHeaders(t *testing.T) { } } -// compat messages bridge 故意不带 originator：收口必须保持 no-op，不得注入身份头。 +// enforce 本身仍只负责收口：缺少 originator 时必须保持 no-op，由需要恢复身份的 +// 调用方先显式调用 ensureCodexIdentityHeaders。 func TestEnforceCodexIdentityHeaders_NoOriginatorIsNoop(t *testing.T) { h := make(http.Header) - h.Set("user-agent", "luna/1.0.0") + h.Set("user-agent", "third-party-client/1.0.0") enforceCodexIdentityHeaders(h) require.Empty(t, h.Get("originator")) - require.Equal(t, "luna/1.0.0", h.Get("user-agent")) + require.Equal(t, "third-party-client/1.0.0", h.Get("user-agent")) } diff --git a/backend/internal/service/openai_codex_message_item_id_test.go b/backend/internal/service/openai_codex_message_item_id_test.go new file mode 100644 index 0000000000..54c2a9ef22 --- /dev/null +++ b/backend/internal/service/openai_codex_message_item_id_test.go @@ -0,0 +1,160 @@ +//go:build unit + +package service + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +// TestFilterCodexInput_StripsMessageItemID_WhenPreservingReferences +// verifies that message items with a non-msg id (e.g. item_*) have their id +// stripped even when PreserveReferences is true. OpenAI upstream requires +// message ids to begin with "msg" and rejects item_* with 400: +// "Expected an ID that begins with 'msg'." (#3981) +func TestFilterCodexInput_StripsMessageItemID_WhenPreservingReferences(t *testing.T) { + input := []any{ + map[string]any{ + "type": "message", + "id": "item_3bc5a3fa8ccde25f1c0000d4", + "role": "user", + "content": []any{ + map[string]any{"type": "input_text", "text": "hello"}, + }, + }, + } + + filtered := filterCodexInputWithOptions(input, codexInputFilterOptions{ + PreserveReferences: true, + }) + + require.Len(t, filtered, 1) + + msg, ok := filtered[0].(map[string]any) + require.True(t, ok) + require.Equal(t, "message", msg["type"]) + _, hasID := msg["id"] + require.False(t, hasID, "item_* id should be stripped from message") + require.Equal(t, "user", msg["role"], "role must be preserved") + require.NotNil(t, msg["content"], "content must be preserved") +} + +// TestFilterCodexInput_KeepsMsgID_WhenPreservingReferences +// verifies that message items with a valid msg* id are kept when +// PreserveReferences is true, so context references are not lost. +func TestFilterCodexInput_KeepsMsgID_WhenPreservingReferences(t *testing.T) { + input := []any{ + map[string]any{ + "type": "message", + "id": "msg_validID123", + "role": "assistant", + }, + } + + filtered := filterCodexInputWithOptions(input, codexInputFilterOptions{ + PreserveReferences: true, + }) + + require.Len(t, filtered, 1) + msg, ok := filtered[0].(map[string]any) + require.True(t, ok) + require.Equal(t, "msg_validID123", msg["id"], "valid msg* id must be preserved") +} + +// TestFilterCodexInput_StripsMessageIDWhenNotPreservingReferences ensures the +// non-continuation path still drops every message id regardless of prefix. +func TestFilterCodexInput_StripsMessageIDWhenNotPreservingReferences(t *testing.T) { + for _, id := range []string{"item_abc", "msg_validID123"} { + input := []any{ + map[string]any{ + "type": "message", + "id": id, + "role": "user", + }, + } + + filtered := filterCodexInputWithOptions(input, codexInputFilterOptions{ + PreserveReferences: false, + }) + + require.Len(t, filtered, 1) + msg, ok := filtered[0].(map[string]any) + require.True(t, ok) + _, hasID := msg["id"] + require.False(t, hasID, "id %q should be stripped when not preserving references", id) + } +} + +// TestFilterCodexInput_MessageIDStripDoesNotMutateInput ensures the original +// input map is not modified in place when the id is stripped. +func TestFilterCodexInput_MessageIDStripDoesNotMutateInput(t *testing.T) { + original := map[string]any{ + "type": "message", + "id": "item_abc", + "role": "user", + } + + filtered := filterCodexInputWithOptions([]any{original}, codexInputFilterOptions{ + PreserveReferences: true, + }) + + require.Len(t, filtered, 1) + require.Equal(t, "item_abc", original["id"], "original input must not be mutated") +} + +// TestFilterCodexInput_MessageStripKeepsFunctionCallBehavior guards against a +// regression of #3785: message and function_call id rules are independent. +func TestFilterCodexInput_MessageStripKeepsFunctionCallBehavior(t *testing.T) { + input := []any{ + map[string]any{ + "type": "message", + "id": "item_msg_001", + "role": "user", + }, + map[string]any{ + "type": "function_call", + "id": "fc_validID123", + "call_id": "fc_validID123", + "name": "bash", + }, + map[string]any{ + "type": "function_call", + "id": "item_A9v0SNfS3VaLrfX0j3y4xhyK", + "call_id": "fc_abc123", + "name": "bash", + }, + map[string]any{ + "type": "function_call_output", + "id": "o1", + "call_id": "fc_abc123", + "output": "done", + }, + } + + filtered := filterCodexInputWithOptions(input, codexInputFilterOptions{ + PreserveReferences: true, + }) + + require.Len(t, filtered, 4) + + msg, ok := filtered[0].(map[string]any) + require.True(t, ok) + _, hasID := msg["id"] + require.False(t, hasID, "message item_* id should be stripped") + + fcValid, ok := filtered[1].(map[string]any) + require.True(t, ok) + require.Equal(t, "fc_validID123", fcValid["id"], "valid fc* id must be preserved") + + fcBad, ok := filtered[2].(map[string]any) + require.True(t, ok) + _, hasID = fcBad["id"] + require.False(t, hasID, "function_call item_* id should still be stripped") + require.Equal(t, "fc_abc123", fcBad["call_id"], "call_id pairing must survive") + + out, ok := filtered[3].(map[string]any) + require.True(t, ok) + require.Equal(t, "o1", out["id"], "output item id should be preserved") + require.Equal(t, "fc_abc123", out["call_id"], "call_id pairing must survive") +} diff --git a/backend/internal/service/openai_codex_transform.go b/backend/internal/service/openai_codex_transform.go index 99355628f2..3869e97c99 100644 --- a/backend/internal/service/openai_codex_transform.go +++ b/backend/internal/service/openai_codex_transform.go @@ -1405,6 +1405,15 @@ func filterCodexInputWithOptions(input []any, opts codexInputFilterOptions) []an ensureCopy() delete(newItem, "id") } + } else if typ == "message" { + // 同理，message 类 item 的 id 必须以 "msg" 开头（上游校验 + // "Expected an ID that begins with 'msg'"）。item_* 形式的 id + // 来自客户端回放，需要删除。 + // 注意：不改写成 msg_*，改写出的 id 未必对应真实的上游对象。 + if id, ok := m["id"].(string); ok && id != "" && !strings.HasPrefix(id, "msg") { + ensureCopy() + delete(newItem, "id") + } } filtered = append(filtered, newItem) diff --git a/backend/internal/service/openai_compact_body_signal.go b/backend/internal/service/openai_compact_body_signal.go index fce62046c1..ce561b0c5a 100644 --- a/backend/internal/service/openai_compact_body_signal.go +++ b/backend/internal/service/openai_compact_body_signal.go @@ -2,18 +2,10 @@ package service import "github.com/tidwall/gjson" -// HasCompactionTriggerInInput detects the Codex remote compact v2 body signal: -// an input item with type "compaction_trigger". When the client sends this -// inside a normal POST /v1/responses (instead of POST /v1/responses/compact), -// the request must still be treated as a compact request — otherwise the -// upstream path, model mapping, and body normalization are all wrong, causing -// Codex to receive a non-compact response and fail with: -// -// "remote compaction v2 expected exactly one compaction output item, got 0" -// -// The gateway handler promotes such requests by rewriting the URL path to the -// compact form before stream parsing, compact body normalization, and -// compact-capable account scheduling, so both inbound forms share one code path. +// HasCompactionTriggerInInput detects an input item with +// type="compaction_trigger". The handler combines this body signal with the +// request path, stream flag, and Codex beta feature header to distinguish the +// native remote compaction v2 wire from the legacy /responses/compact bridge. func HasCompactionTriggerInInput(body []byte) bool { if len(body) == 0 { return false diff --git a/backend/internal/service/openai_compact_sse_keepalive.go b/backend/internal/service/openai_compact_sse_keepalive.go index 70ef3fc01a..70e6af9454 100644 --- a/backend/internal/service/openai_compact_sse_keepalive.go +++ b/backend/internal/service/openai_compact_sse_keepalive.go @@ -1,6 +1,9 @@ package service import ( + "bufio" + "errors" + "net" "net/http" "sync" "time" @@ -43,12 +46,14 @@ func StartOpenAICompactSSEKeepalive(c *gin.Context, interval time.Duration) func if c == nil || c.Writer == nil || interval <= 0 || !openAICompactClientWantsStream(c) { return func() {} } + originalWriter := c.Writer k := &openAICompactSSEKeepalive{ - writer: c.Writer, + writer: originalWriter, stop: make(chan struct{}), } c.Set(openAICompactSSEKeepaliveKey, k) - c.Writer = &openAICompactKeepaliveWriter{ResponseWriter: c.Writer, k: k} + wrappedWriter := &openAICompactKeepaliveWriter{ResponseWriter: originalWriter, k: k} + c.Writer = wrappedWriter var reqDone <-chan struct{} if c.Request != nil { @@ -71,7 +76,14 @@ func StartOpenAICompactSSEKeepalive(c *gin.Context, interval time.Duration) func timer.Reset(interval) } }() - return k.Stop + return func() { + k.Stop() + // Do not leave a pooled middleware writer reachable through the compact + // wrapper after the request finishes. + if current, ok := c.Writer.(*openAICompactKeepaliveWriter); ok && current == wrappedWriter { + c.Writer = originalWriter + } + } } // beat 在锁内提交（首次）响应头并写出一条 SSE 注释行；返回 false 表示心跳已 @@ -181,52 +193,105 @@ type openAICompactKeepaliveWriter struct { // suspend 停拍心跳；幂等。任何响应构造（含 Header 访问——写响应必先操作 // 响应头）都视为请求侧接管 ResponseWriter。 func (w *openAICompactKeepaliveWriter) suspend() { + if w.k == nil { + return + } w.k.Stop() } func (w *openAICompactKeepaliveWriter) Header() http.Header { w.suspend() + if w.ResponseWriter == nil { + return http.Header{} + } return w.ResponseWriter.Header() } func (w *openAICompactKeepaliveWriter) Write(data []byte) (int, error) { w.suspend() + if w.ResponseWriter == nil { + return 0, nil + } return w.ResponseWriter.Write(data) } func (w *openAICompactKeepaliveWriter) WriteString(s string) (int, error) { w.suspend() + if w.ResponseWriter == nil { + return 0, nil + } return w.ResponseWriter.WriteString(s) } func (w *openAICompactKeepaliveWriter) WriteHeader(code int) { w.suspend() + if w.ResponseWriter == nil { + return + } w.ResponseWriter.WriteHeader(code) } func (w *openAICompactKeepaliveWriter) WriteHeaderNow() { w.suspend() + if w.ResponseWriter == nil { + return + } w.ResponseWriter.WriteHeaderNow() } func (w *openAICompactKeepaliveWriter) Flush() { w.suspend() + if w.ResponseWriter == nil { + return + } w.ResponseWriter.Flush() } +func (w *openAICompactKeepaliveWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + if w.ResponseWriter == nil { + return nil, nil, errors.New("response writer released") + } + return w.ResponseWriter.Hijack() +} + +func (w *openAICompactKeepaliveWriter) CloseNotify() <-chan bool { + if w.ResponseWriter == nil { + ch := make(chan bool) + close(ch) + return ch + } + return w.ResponseWriter.CloseNotify() +} + +func (w *openAICompactKeepaliveWriter) Pusher() http.Pusher { + if w.ResponseWriter == nil { + return nil + } + return w.ResponseWriter.Pusher() +} + func (w *openAICompactKeepaliveWriter) Status() int { + if w.k == nil || w.ResponseWriter == nil { + return 0 + } w.k.mu.Lock() defer w.k.mu.Unlock() return w.ResponseWriter.Status() } func (w *openAICompactKeepaliveWriter) Size() int { + if w.k == nil || w.ResponseWriter == nil { + return 0 + } w.k.mu.Lock() defer w.k.mu.Unlock() return w.ResponseWriter.Size() } func (w *openAICompactKeepaliveWriter) Written() bool { + if w.k == nil || w.ResponseWriter == nil { + return false + } w.k.mu.Lock() defer w.k.mu.Unlock() return w.ResponseWriter.Written() diff --git a/backend/internal/service/openai_compact_sse_keepalive_test.go b/backend/internal/service/openai_compact_sse_keepalive_test.go index 3b217a0718..1efed7e9e4 100644 --- a/backend/internal/service/openai_compact_sse_keepalive_test.go +++ b/backend/internal/service/openai_compact_sse_keepalive_test.go @@ -6,6 +6,7 @@ import ( "testing" "time" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/tidwall/gjson" ) @@ -141,6 +142,110 @@ func TestOpenAICompactKeepaliveWriter_RequestSideWriteSuspendsBeats(t *testing.T require.Contains(t, rec.Body.String(), `{"error":"local reject"}`) } +func TestOpenAICompactKeepaliveWriter_NilInnerWriter_NoPanic(t *testing.T) { + w := &openAICompactKeepaliveWriter{ + k: &openAICompactSSEKeepalive{stop: make(chan struct{})}, + } + w.ResponseWriter = nil + + assert.NotPanics(t, func() { + assert.Equal(t, 0, w.Status()) + }) + assert.NotPanics(t, func() { + assert.Equal(t, 0, w.Size()) + }) + assert.NotPanics(t, func() { + assert.False(t, w.Written()) + }) + assert.NotPanics(t, func() { + assert.NotNil(t, w.Header()) + }) + assert.NotPanics(t, func() { + n, err := w.Write([]byte("test")) + assert.Equal(t, 0, n) + assert.NoError(t, err) + }) + assert.NotPanics(t, func() { + n, err := w.WriteString("test") + assert.Equal(t, 0, n) + assert.NoError(t, err) + }) + assert.NotPanics(t, func() { + w.WriteHeader(http.StatusOK) + }) + assert.NotPanics(t, func() { + w.WriteHeaderNow() + }) + assert.NotPanics(t, func() { + w.Flush() + }) + assert.NotPanics(t, func() { + conn, rw, err := w.Hijack() + assert.Nil(t, conn) + assert.Nil(t, rw) + assert.Error(t, err) + }) + assert.NotPanics(t, func() { + ch := w.CloseNotify() + assert.NotNil(t, ch) + }) + assert.NotPanics(t, func() { + assert.Nil(t, w.Pusher()) + }) +} + +func TestOpenAICompactKeepaliveWriter_NilKeepalive_NoPanic(t *testing.T) { + c, rec := newCompactBridgeTestContext(t, true) + w := &openAICompactKeepaliveWriter{ResponseWriter: c.Writer} + + assert.NotPanics(t, func() { + assert.Equal(t, 0, w.Status()) + }) + assert.NotPanics(t, func() { + assert.Equal(t, 0, w.Size()) + }) + assert.NotPanics(t, func() { + assert.False(t, w.Written()) + }) + assert.NotPanics(t, func() { + w.Header().Set("X-Test", "ok") + }) + assert.NotPanics(t, func() { + w.WriteHeader(http.StatusAccepted) + }) + assert.NotPanics(t, func() { + n, err := w.WriteString("ok") + assert.Equal(t, 2, n) + assert.NoError(t, err) + }) + assert.NotPanics(t, func() { + w.Flush() + }) + require.Equal(t, "ok", rec.Header().Get("X-Test")) + require.Equal(t, "ok", rec.Body.String()) +} + +func TestOpenAICompactKeepaliveWriter_DelegatesWhenReady(t *testing.T) { + c, rec := newCompactBridgeTestContext(t, true) + stop := StartOpenAICompactSSEKeepalive(c, time.Hour) + defer stop() + + w, ok := c.Writer.(*openAICompactKeepaliveWriter) + require.True(t, ok) + + w.Header().Set("X-Test", "ok") + w.WriteHeader(http.StatusAccepted) + n, err := w.WriteString("ready") + require.NoError(t, err) + require.Equal(t, len("ready"), n) + + require.Equal(t, http.StatusAccepted, w.Status()) + require.Equal(t, len("ready"), w.Size()) + require.True(t, w.Written()) + require.Equal(t, "ok", rec.Header().Get("X-Test")) + require.Equal(t, "ready", rec.Body.String()) +} + // fast policy block 在心跳提交后必须降级为 response.failed 终止事件。 func TestWriteOpenAIFastPolicyBlockedResponse_AfterKeepaliveCommit(t *testing.T) { c, rec := newCompactBridgeTestContext(t, true) diff --git a/backend/internal/service/openai_compat_model_test.go b/backend/internal/service/openai_compat_model_test.go index 7c7ac1b94f..69b6ddbca2 100644 --- a/backend/internal/service/openai_compat_model_test.go +++ b/backend/internal/service/openai_compat_model_test.go @@ -837,8 +837,7 @@ func TestForwardAsAnthropic_ReusesOAuthCodexTurnState(t *testing.T) { require.NoError(t, err) require.NotNil(t, firstResult) require.Empty(t, upstream.requests[0].Header.Get("x-codex-turn-state")) - require.Empty(t, upstream.requests[0].Header.Get("OpenAI-Beta")) - require.Empty(t, upstream.requests[0].Header.Get("originator")) + requireOpenAIMessagesCodexIdentity(t, upstream.requests[0], codexCLIUserAgent, "codex_cli_rs") secondBody := []byte(`{"model":"claude-sonnet-4-5","max_tokens":16,"messages":[{"role":"user","content":"first"},{"role":"assistant","content":"ok"},{"role":"user","content":"second"}],"stream":false}`) secondRec := httptest.NewRecorder() @@ -852,12 +851,73 @@ func TestForwardAsAnthropic_ReusesOAuthCodexTurnState(t *testing.T) { require.Equal(t, "turn_state_first", upstream.requests[1].Header.Get("x-codex-turn-state")) require.Equal(t, generateSessionUUID(isolateOpenAISessionID(0, "stable-cache-key")), upstream.requests[1].Header.Get("session_id")) require.Empty(t, upstream.requests[1].Header.Get("conversation_id")) - require.Empty(t, upstream.requests[1].Header.Get("OpenAI-Beta")) - require.Empty(t, upstream.requests[1].Header.Get("originator")) + requireOpenAIMessagesCodexIdentity(t, upstream.requests[1], codexCLIUserAgent, "codex_cli_rs") require.False(t, gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").Exists()) require.False(t, gjson.GetBytes(upstream.bodies[1], "previous_response_id").Exists()) } +func TestForwardAsAnthropic_OAuthRestoresCodexIdentityHeaders(t *testing.T) { + gin.SetMode(gin.TestMode) + + const tuiUA = "codex-tui/9.9.9 (Mac OS X 14.0; arm64) iTerm (codex-tui; 9.9.9)" + tests := []struct { + name string + userAgent string + originator string + wantUserAgent string + wantOriginator string + }{ + { + name: "官方UA逐字保留并重新配对", + userAgent: tuiUA, + originator: "opencode", + wantUserAgent: tuiUA, + wantOriginator: "codex-tui", + }, + { + name: "第三方UA回退为默认Codex身份", + userAgent: "third-party-client/1.0.0", + originator: "opencode", + wantUserAgent: codexCLIUserAgent, + wantOriginator: "codex_cli_rs", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + body := []byte(`{"model":"claude-sonnet-4-5","max_tokens":16,"messages":[{"role":"user","content":"hello"}],"stream":false}`) + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + c.Request.Header.Set("User-Agent", tt.userAgent) + c.Request.Header.Set("originator", tt.originator) + + upstream := &httpUpstreamRecorder{resp: openAICompatSSECompletedResponse("resp_identity", "gpt-5.4")} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + cfg: &config.Config{Security: config.SecurityConfig{URLAllowlist: config.URLAllowlistConfig{Enabled: false}}}, + } + account := &Account{ + ID: 1, + Name: "openai-oauth", + Platform: PlatformOpenAI, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "oauth-token", + "chatgpt_account_id": "chatgpt-acc", + }, + } + + result, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "gpt-5.4") + require.NoError(t, err) + require.NotNil(t, result) + requireOpenAIMessagesCodexIdentity(t, upstream.lastReq, tt.wantUserAgent, tt.wantOriginator) + }) + } +} + func TestForwardAsAnthropic_OAuthDigestFallbackReusesTurnStateWithoutExplicitKey(t *testing.T) { t.Parallel() gin.SetMode(gin.TestMode) @@ -896,6 +956,7 @@ func TestForwardAsAnthropic_OAuthDigestFallbackReusesTurnStateWithoutExplicitKey firstSessionID := upstream.requests[0].Header.Get("session_id") require.NotEmpty(t, firstSessionID) require.Empty(t, upstream.requests[0].Header.Get("x-codex-turn-state")) + requireOpenAIMessagesCodexIdentity(t, upstream.requests[0], codexCLIUserAgent, "codex_cli_rs") require.False(t, gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").Exists()) secondBody := []byte(`{"model":"claude-sonnet-4-5","max_tokens":16,"messages":[{"role":"user","content":"first"},{"role":"assistant","content":"ok"},{"role":"user","content":"second"}],"stream":false}`) @@ -910,6 +971,7 @@ func TestForwardAsAnthropic_OAuthDigestFallbackReusesTurnStateWithoutExplicitKey require.Equal(t, firstSessionID, upstream.requests[1].Header.Get("session_id")) require.Equal(t, "turn_state_digest_first", upstream.requests[1].Header.Get("x-codex-turn-state")) require.Empty(t, upstream.requests[1].Header.Get("conversation_id")) + requireOpenAIMessagesCodexIdentity(t, upstream.requests[1], codexCLIUserAgent, "codex_cli_rs") require.False(t, gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").Exists()) require.False(t, gjson.GetBytes(upstream.bodies[1], "previous_response_id").Exists()) } @@ -1064,8 +1126,7 @@ func TestForwardAsAnthropic_OAuthKeepsSystemAsDeveloperInput(t *testing.T) { instructions := gjson.GetBytes(upstream.lastBody, "instructions") require.True(t, instructions.Exists()) require.Empty(t, instructions.String()) - require.Empty(t, upstream.requests[0].Header.Get("OpenAI-Beta")) - require.Empty(t, upstream.requests[0].Header.Get("originator")) + requireOpenAIMessagesCodexIdentity(t, upstream.requests[0], codexCLIUserAgent, "codex_cli_rs") } func TestForwardAsAnthropic_OAuthAddsClaudeCodeTodoGuardForCompatModel(t *testing.T) { @@ -1202,6 +1263,15 @@ func openAICompatSSECompletedResponse(responseID, model string) *http.Response { } } +func requireOpenAIMessagesCodexIdentity(t *testing.T, req *http.Request, wantUserAgent, wantOriginator string) { + t.Helper() + require.NotNil(t, req) + require.Equal(t, wantUserAgent, req.Header.Get("User-Agent")) + require.Equal(t, wantOriginator, req.Header.Get("originator")) + require.Equal(t, codexCLIVersion, req.Header.Get("version")) + require.Equal(t, "responses=experimental", req.Header.Get("OpenAI-Beta")) +} + func openAICompatSSEResponseWithoutUsage(responseID, model string) *http.Response { body := strings.Join([]string{ `data: {"type":"response.completed","response":{"id":"` + responseID + `","object":"response","model":"` + model + `","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}]}}`, diff --git a/backend/internal/service/openai_gateway_cc_pipeline.go b/backend/internal/service/openai_gateway_cc_pipeline.go index 816f5a26e4..690b9991bb 100644 --- a/backend/internal/service/openai_gateway_cc_pipeline.go +++ b/backend/internal/service/openai_gateway_cc_pipeline.go @@ -88,6 +88,9 @@ func (s *OpenAIGatewayService) failoverOpenAIUpstreamHTTPError( upstreamMsg string, upstreamModel string, ) *UpstreamFailoverError { + if account != nil && account.Platform == PlatformGrok { + s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + } if !s.shouldFailoverOpenAIUpstreamResponse(resp.StatusCode, upstreamMsg, respBody) { return nil } @@ -109,7 +112,9 @@ func (s *OpenAIGatewayService) failoverOpenAIUpstreamHTTPError( Message: upstreamMsg, Detail: upstreamDetail, }) - s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel) + if account.Platform != PlatformGrok { + s.handleOpenAIAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody, upstreamModel) + } return &UpstreamFailoverError{ StatusCode: resp.StatusCode, ResponseBody: respBody, @@ -159,6 +164,7 @@ func (s *OpenAIGatewayService) sendCCUpstreamRequest( stream bool, bearerToken string, userAgent string, + grokCacheIdentity string, ) (*http.Response, error) { upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) upstreamReq, err := http.NewRequestWithContext(upstreamCtx, http.MethodPost, targetURL, bytes.NewReader(body)) @@ -190,6 +196,10 @@ func (s *OpenAIGatewayService) sendCCUpstreamRequest( // 账号级请求头覆写（仅 openai api_key 账号启用时生效） account.ApplyHeaderOverrides(upstreamReq.Header) + if account.Platform == PlatformGrok { + applyGrokCLIHeaders(upstreamReq.Header) + applyGrokCacheHeaders(upstreamReq.Header, grokCacheIdentity) + } proxyURL := "" if account.Proxy != nil { diff --git a/backend/internal/service/openai_gateway_chat_completions.go b/backend/internal/service/openai_gateway_chat_completions.go index 15c4e47635..533b562dc9 100644 --- a/backend/internal/service/openai_gateway_chat_completions.go +++ b/backend/internal/service/openai_gateway_chat_completions.go @@ -72,6 +72,16 @@ func (s *OpenAIGatewayService) ForwardAsChatCompletions( } if account.Platform == PlatformGrok { + if account.IsGrokOAuth() { + if eligible, reason := grokChatResponsesBridgeEligibility(body); eligible { + return s.forwardGrokChatCompletionsViaResponses(ctx, c, account, body, promptCacheKey, defaultMappedModel) + } else { + logger.L().Debug("grok chat_completions: using raw fallback", + zap.Int64("account_id", account.ID), + zap.String("reason", reason), + ) + } + } return s.forwardAsRawChatCompletions(ctx, c, account, body, defaultMappedModel) } diff --git a/backend/internal/service/openai_gateway_chat_completions_raw.go b/backend/internal/service/openai_gateway_chat_completions_raw.go index 6def67c7ec..7693986b0d 100644 --- a/backend/internal/service/openai_gateway_chat_completions_raw.go +++ b/backend/internal/service/openai_gateway_chat_completions_raw.go @@ -76,6 +76,12 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( // 2. Resolve model mapping (same as ForwardAsChatCompletions) billingModel := resolveOpenAIForwardModel(account, originalModel, defaultMappedModel) upstreamModel := normalizeOpenAIModelForUpstream(account, billingModel) + grokCacheIdentity := "" + if account.Platform == PlatformGrok { + // Resolve before image bridging or other body rewrites so the fallback is + // anchored to the client's stable conversation prefix. + grokCacheIdentity = resolveGrokCacheIdentity(c, body, "", upstreamModel) + } reasoningEffort := extractOpenAIReasoningEffortFromBody(body, upstreamModel, billingModel, originalModel) // 国产模型默认 effort 补充：需要 mappedModel 判定，推迟到 billingModel 算出之后。 reasoningEffort = ApplyThinkingEnabledFallback(reasoningEffort, body, billingModel) @@ -134,6 +140,12 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( return nil, fmt.Errorf("enable stream usage: %w", usageErr) } } + if account.Platform == PlatformGrok { + upstreamBody, err = stripGrokChatPromptCacheKey(upstreamBody) + if err != nil { + return nil, fmt.Errorf("remove Responses-only Grok prompt cache key: %w", err) + } + } logger.L().Debug("openai chat_completions raw: forwarding without protocol conversion", zap.Int64("account_id", account.ID), @@ -148,11 +160,12 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( if err != nil { return nil, err } + SetActualOpenAIUpstreamEndpoint(c, grokChatRawEndpoint) customUA := account.GetOpenAIUserAgent() if customUA == "" && account.Platform == PlatformGrok { customUA = "sub2api-grok/1.0" } - resp, err := s.sendCCUpstreamRequest(ctx, c, account, targetURL, upstreamBody, clientStream, token, customUA) + resp, err := s.sendCCUpstreamRequest(ctx, c, account, targetURL, upstreamBody, clientStream, token, customUA, grokCacheIdentity) if err != nil { return nil, err } @@ -162,7 +175,6 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( if resp.StatusCode >= 400 { respBody, upstreamMsg := s.readOpenAIUpstreamError(resp) if account.Platform == PlatformGrok { - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ Platform: account.Platform, AccountID: account.ID, @@ -189,7 +201,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( } if account.Platform == PlatformGrok { - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) } // 8. Forward response @@ -202,6 +214,7 @@ func (s *OpenAIGatewayService) forwardAsRawChatCompletions( } if result != nil { addOpenAIUsage(&result.Usage, bridgeUsage) + result.UpstreamEndpoint = grokChatRawEndpoint } return result, forwardErr } diff --git a/backend/internal/service/openai_gateway_forward.go b/backend/internal/service/openai_gateway_forward.go index 53cc06989e..1c13738020 100644 --- a/backend/internal/service/openai_gateway_forward.go +++ b/backend/internal/service/openai_gateway_forward.go @@ -49,7 +49,6 @@ func (s *OpenAIGatewayService) Forward(ctx context.Context, c *gin.Context, acco originalModel := reqModel if account.Platform == PlatformGrok { - _ = promptCacheKey return s.forwardGrokResponses(ctx, c, account, body, originalModel, reqStream, startTime) } diff --git a/backend/internal/service/openai_gateway_grok.go b/backend/internal/service/openai_gateway_grok.go index ef823fa2fd..379b586136 100644 --- a/backend/internal/service/openai_gateway_grok.go +++ b/backend/internal/service/openai_gateway_grok.go @@ -6,6 +6,7 @@ import ( "encoding/json" "fmt" "io" + "log/slog" "net/http" "strings" "time" @@ -20,6 +21,9 @@ import ( const ( grokComposerImageBridgeVisionModel = "grok-build-0.1" grokComposerImageBridgeMaxOutputTokens = 512 + grokUpstreamUserAgent = "sub2api-grok/1.0" + grokCLIVersion = "0.2.93" + grokRateLimitFallbackCooldown = 2 * time.Minute ) func (s *OpenAIGatewayService) forwardGrokResponses( @@ -31,18 +35,23 @@ func (s *OpenAIGatewayService) forwardGrokResponses( reqStream bool, startTime time.Time, ) (*OpenAIForwardResult, error) { - if account.Type != AccountTypeOAuth { - return nil, fmt.Errorf("grok account type %s is not supported by subscription forwarding", account.Type) + if account.Type != AccountTypeOAuth && account.Type != AccountTypeAPIKey { + return nil, fmt.Errorf("grok account type %s is not supported by Responses forwarding", account.Type) } upstreamModel := account.GetMappedModel(originalModel) if strings.TrimSpace(upstreamModel) == "" { upstreamModel = "grok-4.3" } + cacheIdentity := resolveGrokCacheIdentity(c, body, "", upstreamModel) patchedBody, err := patchGrokResponsesBody(body, upstreamModel) if err != nil { return nil, err } + patchedBody, err = applyGrokResponsesCacheIdentity(patchedBody, body, cacheIdentity, account.IsGrokOAuth()) + if err != nil { + return nil, fmt.Errorf("apply grok prompt cache identity: %w", err) + } token, _, err := s.GetAccessToken(ctx, account) if err != nil { @@ -51,7 +60,7 @@ func (s *OpenAIGatewayService) forwardGrokResponses( upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) defer releaseUpstreamCtx() - upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, patchedBody, token) + upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, patchedBody, token, cacheIdentity) if err != nil { return nil, err } @@ -72,7 +81,6 @@ func (s *OpenAIGatewayService) forwardGrokResponses( if resp.StatusCode >= 400 { respBody := s.readUpstreamErrorBody(resp) resp.Body = io.NopCloser(bytes.NewReader(respBody)) - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) upstreamMsg := sanitizeUpstreamErrorMessage(extractUpstreamErrorMessage(respBody)) if upstreamMsg == "" { upstreamMsg = fmt.Sprintf("xAI upstream returned status %d", resp.StatusCode) @@ -97,7 +105,7 @@ func (s *OpenAIGatewayService) forwardGrokResponses( return s.handleErrorResponse(ctx, resp, c, account, patchedBody, upstreamModel) } - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) var usage *OpenAIUsage var firstTokenMs *int @@ -146,6 +154,10 @@ func patchGrokResponsesBody(body []byte, upstreamModel string) ([]byte, error) { if err != nil { return nil, err } + out, err = sanitizeGrokResponsesModelCapabilities(out, upstreamModel) + if err != nil { + return nil, err + } for _, unsupportedField := range []string{"prompt_cache_retention", "safety_identifier"} { if gjson.GetBytes(out, unsupportedField).Exists() { out, err = sjson.DeleteBytes(out, unsupportedField) @@ -168,6 +180,10 @@ func patchGrokResponsesBody(body []byte, upstreamModel string) ([]byte, error) { if err != nil { return nil, err } + out, err = sanitizeGrokResponsesInput(out) + if err != nil { + return nil, err + } out, err = sanitizeGrokResponsesTools(out) if err != nil { return nil, err @@ -175,6 +191,38 @@ func patchGrokResponsesBody(body []byte, upstreamModel string) ([]byte, error) { return out, nil } +func sanitizeGrokResponsesModelCapabilities(body []byte, upstreamModel string) ([]byte, error) { + if !grokModelRejectsReasoningEffort(upstreamModel) { + return body, nil + } + + out := body + for _, field := range []string{"reasoning", "reasoning_effort", "reasoningEffort"} { + if !gjson.GetBytes(out, field).Exists() { + continue + } + var err error + out, err = sjson.DeleteBytes(out, field) + if err != nil { + return nil, fmt.Errorf("remove unsupported Grok Composer %s: %w", field, err) + } + } + return out, nil +} + +func grokModelRejectsReasoningEffort(model string) bool { + model = strings.TrimSpace(strings.ToLower(model)) + if slash := strings.LastIndex(model, "/"); slash >= 0 { + model = strings.TrimSpace(model[slash+1:]) + } + switch model { + case "grok-composer", "grok-composer-2.5-fast", "composer-2.5": + return true + default: + return false + } +} + var grokResponsesUnsupportedRecursiveFields = map[string]struct{}{ "external_web_access": {}, } @@ -223,6 +271,38 @@ func deleteJSONFields(value any, fields map[string]struct{}) bool { } } +// additional_tools is a Codex/Responses Lite private input carrier. xAI's +// Responses schema accepts ordinary message/function-call input items but +// rejects this carrier before inference with a ModelInput deserialization +// error. Top-level supported tools remain available through the separate +// sanitizeGrokResponsesTools path. +func sanitizeGrokResponsesInput(body []byte) ([]byte, error) { + if !bytes.Contains(body, []byte(`"additional_tools"`)) { + return body, nil + } + input := gjson.GetBytes(body, "input") + if !input.Exists() || !input.IsArray() { + return body, nil + } + + rawItems := input.Array() + filtered := make([]json.RawMessage, 0, len(rawItems)) + for _, item := range rawItems { + if strings.TrimSpace(item.Get("type").String()) == "additional_tools" { + continue + } + filtered = append(filtered, json.RawMessage(item.Raw)) + } + if len(filtered) == len(rawItems) { + return body, nil + } + encoded, err := json.Marshal(filtered) + if err != nil { + return nil, err + } + return sjson.SetRawBytes(body, "input", encoded) +} + var grokResponsesSupportedToolTypes = map[string]struct{}{ "code_execution": {}, "code_interpreter": {}, @@ -457,7 +537,9 @@ func (s *OpenAIGatewayService) describeGrokComposerImage( } upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) - upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, body, token) + // Image-description probes are auxiliary requests, not conversation turns. + // Do not bind them to the caller's Grok prompt-cache identity. + upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, body, token, "") releaseUpstreamCtx() if err != nil { return "", OpenAIUsage{}, fmt.Errorf("build grok composer image bridge request: %w", err) @@ -476,7 +558,6 @@ func (s *OpenAIGatewayService) describeGrokComposerImage( if resp.StatusCode >= 400 { respBody := s.readUpstreamErrorBody(resp) - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) upstreamMsg := sanitizeUpstreamErrorMessage(extractUpstreamErrorMessage(respBody)) if upstreamMsg == "" { upstreamMsg = fmt.Sprintf("xAI image bridge upstream returned status %d", resp.StatusCode) @@ -501,7 +582,7 @@ func (s *OpenAIGatewayService) describeGrokComposerImage( return "", OpenAIUsage{}, fmt.Errorf("grok composer image bridge upstream error: %s", upstreamMsg) } - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, nil) if err != nil { return "", OpenAIUsage{}, fmt.Errorf("read grok composer image bridge response: %w", err) @@ -623,7 +704,7 @@ func addOpenAIUsage(dst *OpenAIUsage, usage OpenAIUsage) { dst.ImageOutputTokens += usage.ImageOutputTokens } -func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token string) (*http.Request, error) { +func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, cacheIdentity string) (*http.Request, error) { targetURL, err := xai.BuildResponsesURL(account.GetGrokBaseURL()) if err != nil { return nil, err @@ -635,7 +716,8 @@ func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Acc req.Header.Set("Authorization", "Bearer "+token) req.Header.Set("Content-Type", "application/json") req.Header.Set("Accept", "application/json, text/event-stream") - req.Header.Set("User-Agent", "sub2api-grok/1.0") + applyGrokCLIHeaders(req.Header) + applyGrokCacheHeaders(req.Header, cacheIdentity) if c != nil { if v := c.GetHeader("OpenAI-Beta"); strings.TrimSpace(v) != "" { req.Header.Set("OpenAI-Beta", v) @@ -644,33 +726,201 @@ func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Acc return req, nil } -func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, accountID int64, snapshot *xai.QuotaSnapshot) { - if s == nil || s.accountRepo == nil || accountID <= 0 || snapshot == nil { +// applyGrokCLIHeaders identifies subscription traffic as a supported Grok CLI +// version. The CLI gateway rejects otherwise valid OAuth requests without it. +func applyGrokCLIHeaders(headers http.Header) { + if headers == nil { return } - if s.codexSnapshotThrottle != nil && !s.codexSnapshotThrottle.Allow(accountID, time.Now()) { + headers.Set("User-Agent", grokUpstreamUserAgent) + headers.Set("X-Grok-Client-Version", grokCLIVersion) +} + +func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, account *Account, snapshot *xai.QuotaSnapshot) { + if s == nil || account == nil || account.ID <= 0 || snapshot == nil { return } - _ = s.accountRepo.UpdateExtra(ctx, accountID, map[string]any{ - grokQuotaSnapshotExtraKey: snapshot, - }) + accountID := account.ID + now := time.Now() + resetAt, hasActiveLimit := grokRateLimitResetAt(snapshot, now) + if hasActiveLimit { + normalizeGrokExhaustedWindowResets(snapshot, resetAt, now) + } + critical := snapshot.StatusCode == http.StatusTooManyRequests || hasActiveLimit + if s.codexSnapshotThrottle != nil { + allowed := s.codexSnapshotThrottle.Allow(accountID, now) + if !critical && !allowed { + return + } + } + + stateCtx := ctx + if hasActiveLimit { + var cancel context.CancelFunc + stateCtx, cancel = openAIAccountStateContext(ctx) + defer cancel() + } + if s.accountRepo != nil { + _ = s.accountRepo.UpdateExtra(stateCtx, accountID, map[string]any{ + grokQuotaSnapshotExtraKey: snapshot, + }) + } + // Error responses are reconciled by handleGrokAccountUpstreamError, which + // also installs the immediate in-memory scheduling block. Successful + // responses can still consume the last available request/token, so persist + // that exhausted window here as a real rate limit rather than relying only + // on the passive snapshot scheduler check. + if hasActiveLimit { + s.rateLimitGrok(stateCtx, account, resetAt) + } +} + +func parseGrokQuotaSnapshot(headers http.Header, statusCode int, now time.Time) *xai.QuotaSnapshot { + snapshot := xai.ParseQuotaHeaders(headers, statusCode) + if snapshot == nil && statusCode == http.StatusTooManyRequests { + return &xai.QuotaSnapshot{ + StatusCode: statusCode, + UpdatedAt: now.UTC().Format(time.RFC3339), + } + } + return snapshot +} + +func normalizeGrokExhaustedWindowResets(snapshot *xai.QuotaSnapshot, resetAt, now time.Time) { + if snapshot == nil || !resetAt.After(now) { + return + } + for _, window := range []*xai.QuotaWindow{snapshot.Requests, snapshot.Tokens} { + if window == nil || window.Remaining == nil || *window.Remaining > 0 { + continue + } + candidate := time.Time{} + if window.ResetUnix != nil && *window.ResetUnix > 0 { + candidate = time.Unix(*window.ResetUnix, 0) + } else if parsed, err := time.Parse(time.RFC3339, strings.TrimSpace(window.ResetAt)); err == nil { + candidate = parsed + } + if !candidate.After(now) { + candidate = resetAt + } + resetUnix := candidate.Unix() + window.ResetUnix = &resetUnix + window.ResetAt = candidate.UTC().Format(time.RFC3339) + } +} + +func grokRateLimitResetAt(snapshot *xai.QuotaSnapshot, now time.Time) (time.Time, bool) { + if snapshot == nil { + return time.Time{}, false + } + + // Retry-After is xAI's explicit retry boundary. Use the observation time so + // a persisted snapshot does not start a fresh cooldown every time it is read. + retryAfterExpired := false + var resetAt time.Time + if snapshot.RetryAfterSeconds != nil && *snapshot.RetryAfterSeconds > 0 { + observedAt := now + if parsed, err := time.Parse(time.RFC3339, strings.TrimSpace(snapshot.UpdatedAt)); err == nil { + observedAt = parsed + } + retryAfterResetAt := observedAt.Add(time.Duration(*snapshot.RetryAfterSeconds) * time.Second) + if retryAfterResetAt.After(now) { + resetAt = retryAfterResetAt + } else { + retryAfterExpired = true + } + } + + exhausted := false + for _, window := range []*xai.QuotaWindow{snapshot.Requests, snapshot.Tokens} { + if window == nil || window.Remaining == nil || *window.Remaining > 0 { + continue + } + exhausted = true + candidate := time.Time{} + if window.ResetUnix != nil && *window.ResetUnix > 0 { + candidate = time.Unix(*window.ResetUnix, 0) + } else if parsed, err := time.Parse(time.RFC3339, strings.TrimSpace(window.ResetAt)); err == nil { + candidate = parsed + } + if candidate.After(now) && candidate.After(resetAt) { + resetAt = candidate + } + } + if !resetAt.IsZero() { + return resetAt, true + } + // An observed Retry-After is an absolute boundary once combined with the + // snapshot timestamp. Do not turn an expired persisted snapshot into a new + // rolling fallback cooldown, but still allow a later explicit window reset. + if retryAfterExpired { + return time.Time{}, false + } + if exhausted || snapshot.StatusCode == http.StatusTooManyRequests { + return now.Add(grokRateLimitFallbackCooldown), true + } + return time.Time{}, false +} + +func normalizeGrokRateLimitResetAt(account *Account, resetAt, now time.Time) time.Time { + if !resetAt.After(now) { + resetAt = now.Add(grokRateLimitFallbackCooldown) + } + if account != nil && account.RateLimitResetAt != nil && account.RateLimitResetAt.After(resetAt) { + resetAt = *account.RateLimitResetAt + } + return resetAt +} + +type grokRateLimitExtendingRepository interface { + SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error +} + +func persistGrokRateLimit(ctx context.Context, repo AccountRepository, account *Account, resetAt time.Time) { + if repo == nil || account == nil || account.ID <= 0 { + return + } + resetAt = normalizeGrokRateLimitResetAt(account, resetAt, time.Now()) + stateCtx, cancel := openAIAccountStateContext(ctx) + defer cancel() + var err error + if extendingRepo, ok := repo.(grokRateLimitExtendingRepository); ok { + err = extendingRepo.SetRateLimitedIfLater(stateCtx, account.ID, resetAt) + } else { + err = repo.SetRateLimited(stateCtx, account.ID, resetAt) + } + if err != nil { + slog.Warn("persist_grok_rate_limit_failed", "account_id", account.ID, "reset_at", resetAt.UTC(), "error", err) + } +} + +func (s *OpenAIGatewayService) rateLimitGrok(ctx context.Context, account *Account, resetAt time.Time) { + if s == nil || account == nil { + return + } + resetAt = normalizeGrokRateLimitResetAt(account, resetAt, time.Now()) + + runtimeUntil := resetAt + if account.TempUnschedulableUntil != nil && account.TempUnschedulableUntil.After(runtimeUntil) { + runtimeUntil = *account.TempUnschedulableUntil + } + s.BlockAccountScheduling(account, runtimeUntil, "429") + persistGrokRateLimit(ctx, s.accountRepo, account, resetAt) } func (s *OpenAIGatewayService) handleGrokAccountUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) { if s == nil || account == nil { return } + now := time.Now() + s.updateGrokUsageSnapshot(ctx, account, parseGrokQuotaSnapshot(headers, statusCode, now)) switch statusCode { case http.StatusUnauthorized: - s.tempUnscheduleGrok(ctx, account, 10*time.Minute, "grok oauth token unauthorized") + s.tempUnscheduleGrok(ctx, account, 10*time.Minute, "grok credentials unauthorized") case http.StatusForbidden: - s.tempUnscheduleGrok(ctx, account, 30*time.Minute, "grok entitlement or subscription tier denied") + s.tempUnscheduleGrok(ctx, account, 30*time.Minute, "grok access or entitlement denied") case http.StatusTooManyRequests: - cooldown := 2 * time.Minute - if snapshot := xai.ParseQuotaHeaders(headers, statusCode); snapshot != nil && snapshot.RetryAfterSeconds != nil && *snapshot.RetryAfterSeconds > 0 { - cooldown = time.Duration(*snapshot.RetryAfterSeconds) * time.Second - } - s.tempUnscheduleGrok(ctx, account, cooldown, "grok rate limited") + // updateGrokUsageSnapshot installs both runtime and durable rate-limit state. default: if statusCode >= 500 { s.tempUnscheduleGrok(ctx, account, 2*time.Minute, "grok upstream temporary error") diff --git a/backend/internal/service/openai_gateway_grok_cache.go b/backend/internal/service/openai_gateway_grok_cache.go new file mode 100644 index 0000000000..20934b94c3 --- /dev/null +++ b/backend/internal/service/openai_gateway_grok_cache.go @@ -0,0 +1,149 @@ +package service + +import ( + "fmt" + "net/http" + "strings" + + "github.com/gin-gonic/gin" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +const ( + grokConversationIDHeader = "X-Grok-Conv-Id" + grokFreeCacheNativeToolsJSON = `[{"type":"web_search"},{"type":"x_search"}]` + grokFreeCacheDisabledToolChoice = "none" +) + +// resolveGrokCacheIdentity derives one stable, tenant-isolated routing identity +// for xAI's server-side prompt cache. The returned value is safe to expose to +// the upstream: it never contains the client's raw session identifier. +// +// A valid downstream API key is required. This intentionally fails closed on +// internal probes and incomplete request contexts instead of creating a cache +// identity that could be shared by unrelated tenants. +func resolveGrokCacheIdentity(c *gin.Context, body []byte, explicitKey, upstreamModel string) string { + apiKeyID := getAPIKeyIDFromContext(c) + if apiKeyID <= 0 { + return "" + } + // /responses/compact rejects tool_choice and does not represent a normal + // conversation turn. Keep both cache identity and Free-tier routing + // augmentation out of this path. + if isOpenAIResponsesCompactPath(c) { + return "" + } + + model := strings.ToLower(strings.TrimSpace(upstreamModel)) + if model == "" { + return "" + } + + seed := explicitGrokCacheSeed(c, body, explicitKey) + if seed == "" { + seed = deriveOpenAIContentSessionSeed(body) + } + if seed == "" { + return "" + } + + // generateSessionUUID hashes the whole seed before formatting it as a UUID. + // Include a versioned namespace so this identity cannot collide with other + // upstream session identifiers derived by sub2api. + isolatedSeed := fmt.Sprintf("grok-prompt-cache:v1:%d:%s:%s", apiKeyID, model, seed) + return generateSessionUUID(isolatedSeed) +} + +func explicitGrokCacheSeed(c *gin.Context, body []byte, explicitKey string) string { + seed := "" + if c != nil { + seed = strings.TrimSpace(c.GetHeader("session_id")) + if seed == "" { + seed = strings.TrimSpace(c.GetHeader("conversation_id")) + } + if seed == "" { + seed = strings.TrimSpace(c.GetHeader(grokConversationIDHeader)) + } + } + if seed == "" && len(body) > 0 { + seed = strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String()) + } + if seed == "" { + seed = strings.TrimSpace(explicitKey) + } + return seed +} + +func isGrokRequestContext(c *gin.Context) bool { + if c == nil { + return false + } + v, exists := c.Get("api_key") + if !exists { + return false + } + apiKey, ok := v.(*APIKey) + return ok && apiKey != nil && apiKey.Group != nil && apiKey.Group.Platform == PlatformGrok +} + +// applyGrokResponsesCacheIdentity writes the cache routing identity into an +// xAI Responses request. Existing client values are deliberately replaced by +// the tenant-isolated value to prevent collisions on shared OAuth accounts. +// +// Free OAuth requests without native search tools are routed by xAI to the +// non-cacheable build-free model. For otherwise tool-free requests, add the +// native tools with tool_choice=none: this selects the cache-capable tier +// without allowing an actual search. Any explicit client tools or tool_choice +// disable this augmentation so client function-calling semantics stay intact. +func applyGrokResponsesCacheIdentity(body, intentSourceBody []byte, identity string, injectFreeTierTools bool) ([]byte, error) { + identity = strings.TrimSpace(identity) + if identity == "" { + if gjson.GetBytes(body, "prompt_cache_key").Exists() { + return sjson.DeleteBytes(body, "prompt_cache_key") + } + return body, nil + } + out, err := sjson.SetBytes(body, "prompt_cache_key", identity) + if err != nil { + return nil, err + } + if !injectFreeTierTools { + return out, nil + } + // Inspect the pre-sanitization source. patchGrokResponsesBody may remove an + // unsupported client tool and its tool_choice; that must not turn an + // explicit client tool intent into an eligible native-tool request. + if gjson.GetBytes(intentSourceBody, "tools").Exists() || gjson.GetBytes(intentSourceBody, "tool_choice").Exists() { + return out, nil + } + out, err = sjson.SetRawBytes(out, "tools", []byte(grokFreeCacheNativeToolsJSON)) + if err != nil { + return nil, err + } + return sjson.SetBytes(out, "tool_choice", grokFreeCacheDisabledToolChoice) +} + +// applyGrokCacheHeaders applies the documented Chat Completions conversation +// routing header. The request is built from a fresh header map, so client +// supplied x-grok headers cannot override this server-derived value. +func applyGrokCacheHeaders(headers http.Header, identity string) { + if headers == nil { + return + } + identity = strings.TrimSpace(identity) + if identity == "" { + headers.Del(grokConversationIDHeader) + return + } + headers.Set(grokConversationIDHeader, identity) +} + +// stripGrokChatPromptCacheKey removes the Responses-only body field after it +// has been used as an identity seed. Chat Completions routes cache by header. +func stripGrokChatPromptCacheKey(body []byte) ([]byte, error) { + if !gjson.GetBytes(body, "prompt_cache_key").Exists() { + return body, nil + } + return sjson.DeleteBytes(body, "prompt_cache_key") +} diff --git a/backend/internal/service/openai_gateway_grok_cache_test.go b/backend/internal/service/openai_gateway_grok_cache_test.go new file mode 100644 index 0000000000..556f19304f --- /dev/null +++ b/backend/internal/service/openai_gateway_grok_cache_test.go @@ -0,0 +1,273 @@ +//go:build unit + +package service + +import ( + "net/http" + "net/http/httptest" + "sync" + "testing" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func newGrokCacheTestContext(apiKeyID int64) *gin.Context { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + if apiKeyID > 0 { + c.Set("api_key", &APIKey{ID: apiKeyID, Group: &Group{Platform: PlatformGrok}}) + } + return c +} + +func TestResolveGrokCacheIdentityStableAcrossAppendOnlyTurns(t *testing.T) { + gin.SetMode(gin.TestMode) + c := newGrokCacheTestContext(101) + round1 := []byte(`{"model":"grok","instructions":"be concise","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}],"input":[{"role":"user","content":"first question"}]}`) + round2 := []byte(`{"model":"grok","instructions":"be concise","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}],"input":[{"role":"user","content":"first question"},{"role":"assistant","content":"first answer"},{"role":"user","content":"second question"}]}`) + + first := resolveGrokCacheIdentity(c, round1, "", "grok-4.5") + second := resolveGrokCacheIdentity(c, round2, "", "grok-4.5") + + require.NotEmpty(t, first) + require.Len(t, first, 36) + require.Equal(t, first, second) +} + +func TestResolveGrokCacheIdentityIsolatesAPIKeyAndMappedModel(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"model":"grok","input":"same prompt"}`) + + base := resolveGrokCacheIdentity(newGrokCacheTestContext(201), body, "", "grok-4.5") + otherTenant := resolveGrokCacheIdentity(newGrokCacheTestContext(202), body, "", "grok-4.5") + otherModel := resolveGrokCacheIdentity(newGrokCacheTestContext(201), body, "", "grok-4.3") + + require.NotEmpty(t, base) + require.NotEqual(t, base, otherTenant) + require.NotEqual(t, base, otherModel) +} + +func TestResolveGrokCacheIdentityUsesAndIsolatesNativeConversationHeader(t *testing.T) { + gin.SetMode(gin.TestMode) + c := newGrokCacheTestContext(301) + c.Request.Header.Set(grokConversationIDHeader, "raw-native-conversation") + body1 := []byte(`{"model":"grok","input":"one"}`) + body2 := []byte(`{"model":"grok","input":"different body that must not replace the explicit session"}`) + + first := resolveGrokCacheIdentity(c, body1, "body-cache-key", "grok-4.5") + second := resolveGrokCacheIdentity(c, body2, "another-body-cache-key", "grok-4.5") + + require.Equal(t, "raw-native-conversation", (&OpenAIGatewayService{}).ExtractSessionID(c, body1)) + require.Equal(t, first, second) + require.NotEqual(t, "raw-native-conversation", first) + require.NotContains(t, first, "raw-native-conversation") +} + +func TestResolveGrokCacheIdentityExplicitHeaderPriority(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"model":"grok","prompt_cache_key":"body-key","input":"hi"}`) + c := newGrokCacheTestContext(401) + c.Request.Header.Set(grokConversationIDHeader, "grok-key") + c.Request.Header.Set("conversation_id", "conversation-key") + c.Request.Header.Set("session_id", "session-key") + + got := resolveGrokCacheIdentity(c, body, "explicit-argument", "grok-4.5") + onlySession := newGrokCacheTestContext(401) + onlySession.Request.Header.Set("session_id", "session-key") + want := resolveGrokCacheIdentity(onlySession, []byte(`{"model":"grok","input":"unrelated"}`), "", "grok-4.5") + + require.Equal(t, want, got) +} + +func TestResolveGrokCacheIdentityFailsClosedWithoutAPIKeyContext(t *testing.T) { + gin.SetMode(gin.TestMode) + c := newGrokCacheTestContext(0) + c.Request.Header.Set(grokConversationIDHeader, "native-session") + + require.Empty(t, resolveGrokCacheIdentity(c, []byte(`{"model":"grok","input":"hi"}`), "", "grok-4.5")) + require.Empty(t, resolveGrokCacheIdentity(nil, []byte(`{"model":"grok","prompt_cache_key":"key"}`), "key", "grok-4.5")) +} + +func TestGrokConversationHeaderIsScopedToGrokRequestScheduling(t *testing.T) { + gin.SetMode(gin.TestMode) + body := []byte(`{"model":"grok","prompt_cache_key":"body-session","input":"hi"}`) + + grokContext := newGrokCacheTestContext(601) + grokContext.Request.Header.Set(grokConversationIDHeader, "native-grok-session") + require.Equal(t, "native-grok-session", (&OpenAIGatewayService{}).ExtractSessionID(grokContext, body)) + + openAIContext := newGrokCacheTestContext(601) + openAIContext.Set("api_key", &APIKey{ID: 601, Group: &Group{Platform: PlatformOpenAI}}) + openAIContext.Request.Header.Set(grokConversationIDHeader, "must-be-ignored") + require.Equal(t, "body-session", (&OpenAIGatewayService{}).ExtractSessionID(openAIContext, body)) + + withoutGrokHeader := newGrokCacheTestContext(601) + withoutGrokHeader.Set("api_key", &APIKey{ID: 601, Group: &Group{Platform: PlatformOpenAI}}) + require.Equal(t, + (&OpenAIGatewayService{}).GenerateSessionHash(withoutGrokHeader, body), + (&OpenAIGatewayService{}).GenerateSessionHash(openAIContext, body), + ) +} + +func TestApplyGrokCacheIdentityWritesResponsesBodyAndHeader(t *testing.T) { + sourceBody := []byte(`{"model":"grok-4.5","prompt_cache_key":"raw-client-key"}`) + body, err := applyGrokResponsesCacheIdentity(sourceBody, sourceBody, "isolated-id", true) + require.NoError(t, err) + require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String()) + require.Equal(t, "web_search", gjson.GetBytes(body, "tools.0.type").String()) + require.Equal(t, "x_search", gjson.GetBytes(body, "tools.1.type").String()) + require.Equal(t, grokFreeCacheDisabledToolChoice, gjson.GetBytes(body, "tool_choice").String()) + + headers := make(http.Header) + headers.Set(grokConversationIDHeader, "spoofed-client-value") + applyGrokCacheHeaders(headers, "isolated-id") + require.Equal(t, "isolated-id", headers.Get(grokConversationIDHeader)) + applyGrokCacheHeaders(headers, "") + require.Empty(t, headers.Get(grokConversationIDHeader)) + + chatBody, err := stripGrokChatPromptCacheKey(body) + require.NoError(t, err) + require.False(t, gjson.GetBytes(chatBody, "prompt_cache_key").Exists()) + + unscopedSourceBody := []byte(`{"model":"grok","prompt_cache_key":"raw-client-key"}`) + unscopedBody, err := applyGrokResponsesCacheIdentity(unscopedSourceBody, unscopedSourceBody, "", true) + require.NoError(t, err) + require.False(t, gjson.GetBytes(unscopedBody, "prompt_cache_key").Exists()) + require.False(t, gjson.GetBytes(unscopedBody, "tools").Exists()) + require.False(t, gjson.GetBytes(unscopedBody, "tool_choice").Exists()) +} + +func TestApplyGrokCacheIdentityPreservesExplicitClientToolFields(t *testing.T) { + tests := []struct { + name string + body string + }{ + { + name: "tools only", + body: `{"model":"grok","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}]}`, + }, + { + name: "empty tools array", + body: `{"model":"grok","tools":[]}`, + }, + { + name: "null tools", + body: `{"model":"grok","tools":null}`, + }, + { + name: "tool choice only", + body: `{"model":"grok","tool_choice":{"type":"function","name":"lookup"}}`, + }, + { + name: "null tool choice", + body: `{"model":"grok","tool_choice":null}`, + }, + { + name: "both fields", + body: `{"model":"grok","tools":[{"type":"web_search"}],"tool_choice":"auto"}`, + }, + { + name: "unsupported tool", + body: `{"model":"grok","tools":[{"type":"namespace","name":"client_tools"}]}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + beforeTools := gjson.Get(tt.body, "tools") + beforeChoice := gjson.Get(tt.body, "tool_choice") + body, err := applyGrokResponsesCacheIdentity([]byte(tt.body), []byte(tt.body), "isolated-id", true) + require.NoError(t, err) + require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String()) + require.Equal(t, beforeTools.Exists(), gjson.GetBytes(body, "tools").Exists()) + require.Equal(t, beforeTools.Raw, gjson.GetBytes(body, "tools").Raw) + require.Equal(t, beforeChoice.Exists(), gjson.GetBytes(body, "tool_choice").Exists()) + require.Equal(t, beforeChoice.Raw, gjson.GetBytes(body, "tool_choice").Raw) + }) + } +} + +func TestApplyGrokCacheIdentityUsesPreSanitizationToolIntent(t *testing.T) { + tests := []struct { + name string + intentBody string + }{ + { + name: "unsupported tools removed by sanitizer", + intentBody: `{"model":"grok","tools":[{"type":"namespace","name":"client_tools"}]}`, + }, + { + name: "tool choice removed with unsupported tool", + intentBody: `{"model":"grok","tools":[{"type":"namespace","name":"client_tools"}],"tool_choice":{"type":"namespace","name":"client_tools"}}`, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // This is the shape apply receives after patchGrokResponsesBody has + // removed unsupported tools and their associated tool_choice. + patchedBody := []byte(`{"model":"grok-4.5","input":"hello"}`) + body, err := applyGrokResponsesCacheIdentity(patchedBody, []byte(tt.intentBody), "isolated-id", true) + + require.NoError(t, err) + require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String()) + require.False(t, gjson.GetBytes(body, "tools").Exists()) + require.False(t, gjson.GetBytes(body, "tool_choice").Exists()) + }) + } +} + +func TestApplyGrokCacheIdentityWithoutFreeTierRoutingOnlyWritesIdentity(t *testing.T) { + sourceBody := []byte(`{"model":"grok-4.5","input":"hello"}`) + body, err := applyGrokResponsesCacheIdentity(sourceBody, sourceBody, "isolated-id", false) + + require.NoError(t, err) + require.Equal(t, "isolated-id", gjson.GetBytes(body, "prompt_cache_key").String()) + require.False(t, gjson.GetBytes(body, "tools").Exists()) + require.False(t, gjson.GetBytes(body, "tool_choice").Exists()) +} + +func TestGrokCompactRequestSkipsCacheIdentityAndNativeTools(t *testing.T) { + gin.SetMode(gin.TestMode) + c := newGrokCacheTestContext(701) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", nil) + body := []byte(`{"model":"grok","input":"compact this","prompt_cache_key":"raw-client-key"}`) + + identity := resolveGrokCacheIdentity(c, body, "", "grok-4.5") + patched, err := applyGrokResponsesCacheIdentity(body, body, identity, true) + + require.NoError(t, err) + require.Empty(t, identity) + require.False(t, gjson.GetBytes(patched, "prompt_cache_key").Exists()) + require.False(t, gjson.GetBytes(patched, "tools").Exists()) + require.False(t, gjson.GetBytes(patched, "tool_choice").Exists()) +} + +func TestResolveGrokCacheIdentityConcurrentDeterminism(t *testing.T) { + gin.SetMode(gin.TestMode) + const workers = 50 + body := []byte(`{"model":"grok","messages":[{"role":"system","content":"stable"},{"role":"user","content":"hello"}]}`) + identities := make(chan string, workers) + var wg sync.WaitGroup + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + identities <- resolveGrokCacheIdentity(newGrokCacheTestContext(501), body, "", "grok-4.5") + }() + } + wg.Wait() + close(identities) + + var first string + for identity := range identities { + if first == "" { + first = identity + } + require.Equal(t, first, identity) + } + require.NotEmpty(t, first) +} diff --git a/backend/internal/service/openai_gateway_grok_chat_bridge.go b/backend/internal/service/openai_gateway_grok_chat_bridge.go new file mode 100644 index 0000000000..fd6b426536 --- /dev/null +++ b/backend/internal/service/openai_gateway_grok_chat_bridge.go @@ -0,0 +1,320 @@ +package service + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "net/http" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/gin-gonic/gin" +) + +const ( + grokChatResponsesEndpoint = "/v1/responses" + grokChatRawEndpoint = "/v1/chat/completions" +) + +var grokChatResponsesBridgeTopLevelFields = map[string]struct{}{ + "model": {}, + "messages": {}, + "stream": {}, + "stream_options": {}, + "max_tokens": {}, + "max_completion_tokens": {}, + "temperature": {}, + "top_p": {}, + "prompt_cache_key": {}, + "tools": {}, + "tool_choice": {}, + "functions": {}, + "function_call": {}, +} + +// grokChatResponsesBridgeEligibility deliberately accepts only request shapes +// whose Chat Completions semantics are preserved by the Responses bridge. +// Everything else stays on raw Chat Completions rather than being silently +// dropped or rewritten. +func grokChatResponsesBridgeEligibility(body []byte) (bool, string) { + var root map[string]json.RawMessage + if err := json.Unmarshal(body, &root); err != nil || root == nil { + return false, "invalid_json" + } + + for _, field := range []string{"stop", "reasoning_effort"} { + if _, exists := root[field]; exists { + return false, "unsupported_" + field + } + } + for _, field := range []string{"tools", "functions"} { + if raw, exists := root[field]; exists && !grokChatNullOrEmptyArray(raw) { + return false, "unsupported_" + field + } + } + if raw, exists := root["tool_choice"]; exists && !grokChatNullOrNone(raw) { + return false, "unsupported_tool_choice" + } + if raw, exists := root["function_call"]; exists && !grokChatNullOrNone(raw) { + return false, "unsupported_function_call" + } + for field := range root { + if _, supported := grokChatResponsesBridgeTopLevelFields[field]; !supported { + return false, "unknown_field_" + field + } + } + + var model string + if raw, ok := root["model"]; !ok || json.Unmarshal(raw, &model) != nil || strings.TrimSpace(model) == "" { + return false, "invalid_model" + } + + if raw, ok := root["stream"]; ok { + var stream *bool + if json.Unmarshal(raw, &stream) != nil || stream == nil { + return false, "invalid_stream" + } + } + if raw, ok := root["stream_options"]; ok { + var options map[string]json.RawMessage + if json.Unmarshal(raw, &options) != nil || options == nil { + return false, "invalid_stream_options" + } + for field, value := range options { + if field != "include_usage" { + return false, "unknown_stream_option_" + field + } + var includeUsage *bool + if json.Unmarshal(value, &includeUsage) != nil || includeUsage == nil { + return false, "invalid_stream_include_usage" + } + } + } + + for _, field := range []string{"max_tokens", "max_completion_tokens"} { + if raw, ok := root[field]; ok { + var value *int + if json.Unmarshal(raw, &value) != nil || value == nil || *value < 128 { + return false, "unsafe_" + field + } + } + } + if _, hasMaxTokens := root["max_tokens"]; hasMaxTokens { + if _, hasMaxCompletionTokens := root["max_completion_tokens"]; hasMaxCompletionTokens { + return false, "conflicting_max_tokens" + } + } + for _, field := range []string{"temperature", "top_p"} { + if raw, ok := root[field]; ok { + var value *float64 + if json.Unmarshal(raw, &value) != nil || value == nil { + return false, "invalid_" + field + } + } + } + if raw, ok := root["prompt_cache_key"]; ok { + var key string + if json.Unmarshal(raw, &key) != nil { + return false, "invalid_prompt_cache_key" + } + } + + var messages []map[string]json.RawMessage + rawMessages, ok := root["messages"] + if !ok || json.Unmarshal(rawMessages, &messages) != nil || len(messages) == 0 { + return false, "invalid_messages" + } + for _, message := range messages { + for field := range message { + if field != "role" && field != "content" { + return false, "unsafe_message_field_" + field + } + } + var role string + if raw, exists := message["role"]; !exists || json.Unmarshal(raw, &role) != nil { + return false, "invalid_message_role" + } + switch role { + case "system", "user", "assistant": + default: + return false, "unsupported_message_role_" + role + } + var content string + if raw, exists := message["content"]; !exists || json.Unmarshal(raw, &content) != nil { + // Structured content includes image_url and other parts whose exact + // behavior is not guaranteed by this bridge. + return false, "non_text_message_content" + } + if strings.TrimSpace(content) == "" { + return false, "empty_message_content" + } + } + + return true, "" +} + +func grokChatNullOrEmptyArray(raw json.RawMessage) bool { + if strings.TrimSpace(string(raw)) == "null" { + return true + } + var values []json.RawMessage + return json.Unmarshal(raw, &values) == nil && len(values) == 0 +} + +func grokChatNullOrNone(raw json.RawMessage) bool { + if strings.TrimSpace(string(raw)) == "null" { + return true + } + var value string + return json.Unmarshal(raw, &value) == nil && strings.EqualFold(strings.TrimSpace(value), "none") +} + +func grokChatCacheIntentBody(body []byte) ([]byte, error) { + var root map[string]json.RawMessage + if err := json.Unmarshal(body, &root); err != nil { + return nil, err + } + for _, field := range []string{"tools", "tool_choice", "functions", "function_call"} { + delete(root, field) + } + return json.Marshal(root) +} + +func grokChatResponsesRuntimeEligible(upstreamModel, cacheIdentity string) bool { + return strings.TrimSpace(upstreamModel) == "grok-4.5" && strings.TrimSpace(cacheIdentity) != "" +} + +// forwardGrokChatCompletionsViaResponses converts a strictly compatible Chat +// request into xAI Responses format and reuses the established Responses-to- +// Chat response translators. It intentionally does not run the Codex OAuth +// transform because Grok CLI is a separate upstream protocol. +func (s *OpenAIGatewayService) forwardGrokChatCompletionsViaResponses( + ctx context.Context, + c *gin.Context, + account *Account, + body []byte, + promptCacheKey string, + defaultMappedModel string, +) (*OpenAIForwardResult, error) { + startTime := time.Now() + + var chatReq apicompat.ChatCompletionsRequest + if err := json.Unmarshal(body, &chatReq); err != nil { + return nil, fmt.Errorf("parse grok chat completions request: %w", err) + } + originalModel := chatReq.Model + clientStream := chatReq.Stream + billingModel := resolveOpenAIForwardModel(account, originalModel, defaultMappedModel) + upstreamModel := normalizeOpenAIModelForUpstream(account, billingModel) + cacheIdentity := resolveGrokCacheIdentity(c, body, promptCacheKey, upstreamModel) + if !grokChatResponsesRuntimeEligible(upstreamModel, cacheIdentity) { + return s.forwardAsRawChatCompletions(ctx, c, account, body, defaultMappedModel) + } + + responsesReq, err := apicompat.ChatCompletionsToResponses(&chatReq) + if err != nil { + return nil, fmt.Errorf("convert grok chat completions to responses: %w", err) + } + responsesReq.Model = upstreamModel + responsesReq.Stream = true + // These fields are useful to Codex but are not needed by the Grok CLI + // protocol. Keep the bridge request as close as possible to native Grok. + responsesReq.Include = nil + responsesReq.Store = nil + + responsesBody, err := json.Marshal(responsesReq) + if err != nil { + return nil, fmt.Errorf("marshal grok responses bridge request: %w", err) + } + responsesBody, err = patchGrokResponsesBody(responsesBody, upstreamModel) + if err != nil { + return nil, fmt.Errorf("patch grok responses bridge request: %w", err) + } + intentBody, err := grokChatCacheIntentBody(body) + if err != nil { + return nil, fmt.Errorf("normalize grok responses bridge tool intent: %w", err) + } + responsesBody, err = applyGrokResponsesCacheIdentity(responsesBody, intentBody, cacheIdentity, true) + if err != nil { + return nil, fmt.Errorf("apply grok responses bridge cache identity: %w", err) + } + + updatedBody, policyErr := s.applyOpenAIFastPolicyToBody(ctx, account, upstreamModel, responsesBody) + if policyErr != nil { + var blocked *OpenAIFastBlockedError + if errors.As(policyErr, &blocked) { + MarkOpsClientBusinessLimited(c, OpsClientBusinessLimitedReasonLocalPolicyDenied) + writeChatCompletionsError(c, http.StatusForbidden, "permission_error", blocked.Message) + } + return nil, policyErr + } + responsesBody = updatedBody + + token, _, err := s.GetAccessToken(ctx, account) + if err != nil { + return nil, fmt.Errorf("get grok access token: %w", err) + } + upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) + upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, cacheIdentity) + releaseUpstreamCtx() + if err != nil { + return nil, fmt.Errorf("build grok responses bridge request: %w", err) + } + SetActualOpenAIUpstreamEndpoint(c, grokChatResponsesEndpoint) + + proxyURL := "" + if account.ProxyID != nil && account.Proxy != nil { + proxyURL = account.Proxy.URL() + } + resp, err := s.httpUpstream.Do(upstreamReq, proxyURL, account.ID, account.Concurrency) + if err != nil { + return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false) + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode >= http.StatusBadRequest { + respBody, upstreamMsg := s.readOpenAIUpstreamError(resp) + if upstreamMsg == "" { + upstreamMsg = fmt.Sprintf("xAI upstream returned status %d", resp.StatusCode) + } + appendOpsUpstreamError(c, OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + UpstreamStatusCode: resp.StatusCode, + UpstreamRequestID: firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")), + Kind: "failover", + Message: upstreamMsg, + }) + s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + if s.shouldFailoverUpstreamError(resp.StatusCode) { + return nil, &UpstreamFailoverError{ + StatusCode: resp.StatusCode, + ResponseBody: respBody, + RetryableOnSameAccount: account.IsPoolMode() && account.IsPoolModeRetryableStatus(resp.StatusCode), + } + } + return s.handleChatCompletionsErrorResponse(resp, c, account, billingModel) + } + + s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + + var result *OpenAIForwardResult + if clientStream { + result, err = s.handleChatStreamingResponse(resp, c, account, originalModel, billingModel, upstreamModel, startTime, len(body)) + } else { + result, err = s.handleChatBufferedStreamingResponse(resp, c, account, originalModel, billingModel, upstreamModel, startTime) + } + if result != nil { + result.UpstreamEndpoint = grokChatResponsesEndpoint + result.ResponseHeaders = resp.Header.Clone() + if result.RequestID == "" { + result.RequestID = firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")) + } + result.ReasoningEffort = extractOpenAIReasoningEffortFromBody(body, upstreamModel, billingModel, originalModel) + } + return result, err +} diff --git a/backend/internal/service/openai_gateway_grok_chat_bridge_test.go b/backend/internal/service/openai_gateway_grok_chat_bridge_test.go new file mode 100644 index 0000000000..0843f1ea42 --- /dev/null +++ b/backend/internal/service/openai_gateway_grok_chat_bridge_test.go @@ -0,0 +1,358 @@ +//go:build unit + +package service + +import ( + "bytes" + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "github.com/tidwall/gjson" +) + +func TestGrokChatResponsesBridgeEligibility(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + body string + want bool + reason string + }{ + { + name: "plain text chat", + body: `{"model":"grok","messages":[{"role":"system","content":"concise"},{"role":"user","content":"hi"}],"stream":false}`, + want: true, + }, + { + name: "safe generation options", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":true,"stream_options":{"include_usage":true},"max_completion_tokens":256,"temperature":0.2,"top_p":0.9,"prompt_cache_key":"session","tools":[],"functions":null,"tool_choice":"none"}`, + want: true, + }, + { + name: "stop falls back", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"stop":"done"}`, + reason: "unsupported_stop", + }, + { + name: "developer role falls back", + body: `{"model":"grok","messages":[{"role":"developer","content":"rules"},{"role":"user","content":"hi"}]}`, + reason: "unsupported_message_role_developer", + }, + { + name: "image content falls back", + body: `{"model":"grok","messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":"data:image/png;base64,QQ=="}}]}]}`, + reason: "non_text_message_content", + }, + { + name: "function tools fall back", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[{"type":"function","function":{"name":"lookup"}}]}`, + reason: "unsupported_tools", + }, + { + name: "automatic tool choice falls back", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"tools":[],"tool_choice":"auto"}`, + reason: "unsupported_tool_choice", + }, + { + name: "reasoning effort falls back because conversion adds summary", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"reasoning_effort":"high"}`, + reason: "unsupported_reasoning_effort", + }, + { + name: "both token limits fall back", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"max_tokens":256,"max_completion_tokens":256}`, + reason: "conflicting_max_tokens", + }, + { + name: "empty message falls back", + body: `{"model":"grok","messages":[{"role":"assistant","content":""},{"role":"user","content":"hi"}]}`, + reason: "empty_message_content", + }, + { + name: "tool history falls back", + body: `{"model":"grok","messages":[{"role":"assistant","content":"","tool_calls":[]}]}`, + reason: "unsafe_message_field_tool_calls", + }, + { + name: "unknown field falls back", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"seed":7}`, + reason: "unknown_field_seed", + }, + { + name: "small max tokens falls back because conversion clamps it", + body: `{"model":"grok","messages":[{"role":"user","content":"hi"}],"max_tokens":32}`, + reason: "unsafe_max_tokens", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, reason := grokChatResponsesBridgeEligibility([]byte(tt.body)) + require.Equal(t, tt.want, got) + require.Equal(t, tt.reason, reason) + }) + } +} + +func TestGrokChatResponsesRuntimeEligibility(t *testing.T) { + t.Parallel() + require.True(t, grokChatResponsesRuntimeEligible("grok-4.5", "isolated-id")) + require.False(t, grokChatResponsesRuntimeEligible("grok-4.3", "isolated-id")) + require.False(t, grokChatResponsesRuntimeEligible("grok-4.5-build-free", "isolated-id")) + require.False(t, grokChatResponsesRuntimeEligible("grok-4.5", "")) +} + +func TestForwardGrokChatViaResponsesNonStreamingCachesAndReturnsChat(t *testing.T) { + gin.SetMode(gin.TestMode) + + body := []byte(`{"model":"grok","messages":[{"role":"system","content":"be concise"},{"role":"user","content":"hi"}],"stream":false,"prompt_cache_key":"stable-session","tools":[],"functions":null,"tool_choice":"none"}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 7101}) + + account := grokChatBridgeTestAccount(71) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &httpUpstreamRecorder{resp: grokChatBridgeCompletedResponse("resp_grok_chat_cache", 9856)} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "") + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) + require.Equal(t, grokChatResponsesEndpoint, result.UpstreamEndpoint) + require.Equal(t, "grok-4.5", result.UpstreamModel) + require.Equal(t, 9908, result.Usage.InputTokens) + require.Equal(t, 12, result.Usage.OutputTokens) + require.Equal(t, 9856, result.Usage.CacheReadInputTokens) + + identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String() + require.NotEmpty(t, identity) + require.NotEqual(t, "stable-session", identity) + require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader)) + require.Equal(t, "web_search", gjson.GetBytes(upstream.lastBody, "tools.0.type").String()) + require.Equal(t, "x_search", gjson.GetBytes(upstream.lastBody, "tools.1.type").String()) + require.Equal(t, grokFreeCacheDisabledToolChoice, gjson.GetBytes(upstream.lastBody, "tool_choice").String()) + require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool()) + require.Equal(t, "system", gjson.GetBytes(upstream.lastBody, "input.0.role").String()) + require.Equal(t, "user", gjson.GetBytes(upstream.lastBody, "input.1.role").String()) + require.False(t, gjson.GetBytes(upstream.lastBody, "instructions").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, "include").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, "store").Exists()) + + require.Equal(t, http.StatusOK, recorder.Code) + require.Equal(t, "cached ok", gjson.Get(recorder.Body.String(), "choices.0.message.content").String()) + require.Equal(t, int64(9856), gjson.Get(recorder.Body.String(), "usage.prompt_tokens_details.cached_tokens").Int()) + require.NotNil(t, repo.updates[account.ID][grokQuotaSnapshotExtraKey]) +} + +func TestForwardGrokChatViaResponsesStreamingPropagatesCachedUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + + body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":true}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 7201}) + + account := grokChatBridgeTestAccount(72) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &httpUpstreamRecorder{resp: grokChatBridgeCompletedResponse("resp_grok_chat_stream", 4096)} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "") + require.NoError(t, err) + require.NotNil(t, result) + require.True(t, result.Stream) + require.Equal(t, grokChatResponsesEndpoint, result.UpstreamEndpoint) + require.Equal(t, 4096, result.Usage.CacheReadInputTokens) + require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream") + require.Contains(t, recorder.Body.String(), `"content":"cached ok"`) + require.Contains(t, recorder.Body.String(), `"cached_tokens":4096`) + require.Contains(t, recorder.Body.String(), "data: [DONE]") +} + +func TestForwardGrokChatRuntimeGateFallsBackToRaw(t *testing.T) { + gin.SetMode(gin.TestMode) + + tests := []struct { + name string + setAPIKey bool + mappedModel string + wantUpstream string + }{ + {name: "missing cache identity", wantUpstream: "grok-4.5"}, + {name: "non cache capable mapped model", setAPIKey: true, mappedModel: "grok-4.3", wantUpstream: "grok-4.3"}, + } + + for index, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body)) + if tt.setAPIKey { + c.Set("api_key", &APIKey{ID: int64(7301 + index)}) + } + + account := grokChatBridgeTestAccount(int64(73 + index)) + if tt.mappedModel != "" { + account.Credentials["model_mapping"] = map[string]any{"grok": tt.mappedModel} + } + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader( + `{"id":"chat_raw","object":"chat.completion","model":"` + tt.wantUpstream + `","choices":[{"index":0,"message":{"role":"assistant","content":"raw ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":2,"completion_tokens":1,"total_tokens":3}}`, + )), + }} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "") + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String()) + require.Equal(t, grokChatRawEndpoint, result.UpstreamEndpoint) + require.Equal(t, tt.wantUpstream, result.UpstreamModel) + require.False(t, gjson.GetBytes(upstream.lastBody, "tools").Exists()) + require.Equal(t, "raw ok", gjson.Get(recorder.Body.String(), "choices.0.message.content").String()) + }) + } +} + +func TestForwardGrokChatViaResponses429UsesGrokRateLimitPolicy(t *testing.T) { + gin.SetMode(gin.TestMode) + + body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 7501}) + + account := grokChatBridgeTestAccount(75) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Retry-After": []string{"45"}, + }, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)), + }} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + before := time.Now() + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "") + require.Error(t, err) + require.Nil(t, result) + var failoverErr *UpstreamFailoverError + require.True(t, errors.As(err, &failoverErr)) + require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode) + require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) + require.Equal(t, grokChatResponsesEndpoint, GetActualOpenAIUpstreamEndpoint(c)) + require.Equal(t, 1, repo.rateLimitedCalls) + require.Zero(t, repo.tempUnschedCalls) + require.WithinDuration(t, before.Add(45*time.Second), repo.lastRateLimitResetAt, time.Second) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestForwardGrokRawChatErrorRecordsActualEndpoint(t *testing.T) { + gin.SetMode(gin.TestMode) + + body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false,"stop":"done"}`) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, grokChatRawEndpoint, bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 7601}) + + account := grokChatBridgeTestAccount(76) + repo := &grokQuotaAccountRepo{mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{account.ID: account}, + }} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusBadRequest, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"bad request"}}`)), + }} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "") + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String()) + require.Equal(t, grokChatRawEndpoint, GetActualOpenAIUpstreamEndpoint(c)) +} + +func grokChatBridgeTestAccount(id int64) *Account { + return &Account{ + ID: id, + Name: "grok-cache-bridge", + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + "base_url": xai.DefaultCLIBaseURL, + }, + } +} + +func grokChatBridgeCompletedResponse(responseID string, cachedTokens int) *http.Response { + body := strings.Join([]string{ + `data: {"type":"response.output_text.delta","sequence_number":0,"delta":"cached ok"}`, + "", + `data: {"type":"response.completed","sequence_number":1,"response":{"id":"` + responseID + `","object":"response","model":"grok-4.5","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"cached ok"}]}],"usage":{"input_tokens":9908,"output_tokens":12,"total_tokens":9920,"input_tokens_details":{"cached_tokens":` + strconv.Itoa(cachedTokens) + `}}}}`, + "", + }, "\n") + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "Xai-Request-Id": []string{responseID + "-request"}, + "X-Ratelimit-Limit-Requests": []string{"10"}, + "X-Ratelimit-Remaining-Requests": []string{"9"}, + }, + Body: io.NopCloser(strings.NewReader(body)), + } +} diff --git a/backend/internal/service/openai_gateway_grok_test.go b/backend/internal/service/openai_gateway_grok_test.go index 8dbd9ddad0..71e427c790 100644 --- a/backend/internal/service/openai_gateway_grok_test.go +++ b/backend/internal/service/openai_gateway_grok_test.go @@ -6,6 +6,7 @@ import ( "bytes" "context" "encoding/json" + "fmt" "io" "mime/multipart" "net/http" @@ -41,6 +42,50 @@ func TestPatchGrokResponsesBodySetsMappedModelAndDropsUnsupportedFields(t *testi require.Equal(t, "high", gjson.GetBytes(patched, "reasoning.effort").String()) } +func TestPatchGrokResponsesBodySanitizesComposerReasoningParameters(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + upstreamModel string + wantReasoning bool + }{ + {name: "composer fast", upstreamModel: "grok-composer-2.5-fast"}, + {name: "composer shorthand", upstreamModel: "grok-composer"}, + {name: "composer legacy alias", upstreamModel: "composer-2.5"}, + {name: "provider-prefixed composer", upstreamModel: "xai/grok-composer-2.5-fast"}, + {name: "grok 4.5", upstreamModel: "grok-4.5", wantReasoning: true}, + } + + body := []byte(`{ + "model": "grok", + "input": "hello", + "reasoning": {"effort": "medium", "summary": "auto"}, + "reasoning_effort": "medium", + "reasoningEffort": "medium" + }`) + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + patched, err := patchGrokResponsesBody(body, tt.upstreamModel) + require.NoError(t, err) + require.True(t, json.Valid(patched)) + require.Equal(t, tt.upstreamModel, gjson.GetBytes(patched, "model").String()) + + if tt.wantReasoning { + require.Equal(t, "medium", gjson.GetBytes(patched, "reasoning.effort").String()) + require.Equal(t, "medium", gjson.GetBytes(patched, "reasoning_effort").String()) + require.Equal(t, "medium", gjson.GetBytes(patched, "reasoningEffort").String()) + return + } + + require.False(t, gjson.GetBytes(patched, "reasoning").Exists()) + require.False(t, gjson.GetBytes(patched, "reasoning_effort").Exists()) + require.False(t, gjson.GetBytes(patched, "reasoningEffort").Exists()) + }) + } +} + func TestExtractGrokResponsesReasoningEffortSupportsOpenAICompatibleField(t *testing.T) { t.Parallel() @@ -161,6 +206,45 @@ func TestPatchGrokResponsesBodyDropsToolChoiceWhenNoSupportedToolsRemain(t *test require.False(t, gjson.GetBytes(patched, "tool_choice").Exists()) } +func TestPatchGrokResponsesBodyDropsCodexAdditionalToolsInputItems(t *testing.T) { + t.Parallel() + + body := []byte(`{ + "model": "grok", + "input": [ + { + "type": "additional_tools", + "role": "developer", + "tools": [ + {"type": "namespace", "name": "image_gen"}, + {"type": "function", "name": "wait"} + ] + }, + { + "type": "message", + "role": "developer", + "content": [{"type": "input_text", "text": "system prompt"}] + }, + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "hello"}] + } + ] + }`) + + patched, err := patchGrokResponsesBody(body, "grok-4.5") + require.NoError(t, err) + require.True(t, json.Valid(patched)) + require.Equal(t, "grok-4.5", gjson.GetBytes(patched, "model").String()) + require.Equal(t, 2, len(gjson.GetBytes(patched, "input").Array())) + require.False(t, gjson.GetBytes(patched, `input.#(type=="additional_tools")`).Exists()) + require.Equal(t, "developer", gjson.GetBytes(patched, "input.0.role").String()) + require.Equal(t, "system prompt", gjson.GetBytes(patched, "input.0.content.0.text").String()) + require.Equal(t, "user", gjson.GetBytes(patched, "input.1.role").String()) + require.Equal(t, "hello", gjson.GetBytes(patched, "input.1.content.0.text").String()) +} + func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T) { t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") @@ -172,13 +256,15 @@ func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T) }, } - req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token") + req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "isolated-cache-id") require.NoError(t, err) require.Equal(t, http.MethodPost, req.Method) require.Equal(t, "https://xai.test/v1/responses", req.URL.String()) require.Equal(t, "Bearer access-token", req.Header.Get("Authorization")) require.Equal(t, "application/json", req.Header.Get("Content-Type")) require.Contains(t, req.Header.Get("Accept"), "text/event-stream") + require.Equal(t, grokCLIVersion, req.Header.Get("X-Grok-Client-Version")) + require.Equal(t, "isolated-cache-id", req.Header.Get(grokConversationIDHeader)) data, err := io.ReadAll(req.Body) require.NoError(t, err) @@ -196,7 +282,7 @@ func TestBuildGrokResponsesRequestRejectsUnsafeAccountBaseURL(t *testing.T) { }, } - _, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token") + _, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "") require.Error(t, err) require.Contains(t, err.Error(), "invalid base url") } @@ -324,6 +410,7 @@ func TestForwardGrokMediaImagesGenerationNormalizesImagineAlias(t *testing.T) { require.Equal(t, http.MethodPost, upstream.lastReq.Method) require.Equal(t, "Bearer api-key", upstream.lastReq.Header.Get("Authorization")) require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type")) + require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version")) require.JSONEq(t, `{"model":"grok-imagine-image-quality","prompt":"draw a cat"}`, string(upstream.lastBody)) require.Equal(t, http.StatusOK, recorder.Code) require.JSONEq(t, `{"data":[]}`, recorder.Body.String()) @@ -565,7 +652,7 @@ func TestBindGrokMediaVideoRequestAccountUsesRequestIDStickyHash(t *testing.T) { require.Equal(t, int64(63), accountID) } -func TestForwardGrokMediaErrorHonorsCustomErrorCodes(t *testing.T) { +func TestForwardGrokMedia429ReconcilesRateLimitBeforeCustomErrorBypass(t *testing.T) { t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true") gin.SetMode(gin.TestMode) @@ -585,18 +672,20 @@ func TestForwardGrokMediaErrorHonorsCustomErrorCodes(t *testing.T) { "api_key": "api-key", "base_url": "https://xai.test/v1", "custom_error_codes_enabled": true, - "custom_error_codes": []any{float64(http.StatusTooManyRequests)}, + "custom_error_codes": []any{float64(http.StatusBadRequest)}, }, } upstream := &httpUpstreamRecorder{resp: &http.Response{ - StatusCode: http.StatusBadRequest, + StatusCode: http.StatusTooManyRequests, Header: http.Header{ "Content-Type": []string{"application/json"}, "Xai-Request-Id": []string{"xai-error-req"}, + "Retry-After": []string{"45"}, }, Body: io.NopCloser(strings.NewReader(`{"error":{"message":"do not expose this upstream detail"}}`)), }} - svc := &OpenAIGatewayService{httpUpstream: upstream} + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{httpUpstream: upstream, accountRepo: repo} result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json") require.Error(t, err) @@ -604,15 +693,19 @@ func TestForwardGrokMediaErrorHonorsCustomErrorCodes(t *testing.T) { require.Equal(t, http.StatusInternalServerError, recorder.Code) require.Contains(t, recorder.Body.String(), "Upstream gateway error") require.NotContains(t, recorder.Body.String(), "do not expose") + require.Equal(t, 1, repo.rateLimitedCalls) + require.Zero(t, repo.tempUnschedCalls) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) } -func TestForwardAsChatCompletionsForGrokUsesXAIChatCompletionsAndSnapshots(t *testing.T) { +func TestForwardAsChatCompletionsForGrokStopFallsBackToXAIChatCompletions(t *testing.T) { gin.SetMode(gin.TestMode) recorder := httptest.NewRecorder() c, _ := gin.CreateTestContext(recorder) - body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false}`) + body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false,"stop":"done","prompt_cache_key":"raw-client-cache-key"}`) c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 5101}) account := &Account{ ID: 51, @@ -641,7 +734,7 @@ func TestForwardAsChatCompletionsForGrokUsesXAIChatCompletionsAndSnapshots(t *te "X-Ratelimit-Limit-Tokens": []string{"1000"}, "X-Ratelimit-Remaining-Tokens": []string{"990"}, }, - Body: io.NopCloser(strings.NewReader(`{"id":"chatcmpl","object":"chat.completion","model":"grok-4.3","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2}}`)), + Body: io.NopCloser(strings.NewReader(`{"id":"chatcmpl","object":"chat.completion","model":"grok-4.3","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2,"prompt_tokens_details":{"cached_tokens":1}}}`)), }} svc := &OpenAIGatewayService{ httpUpstream: upstream, @@ -653,11 +746,15 @@ func TestForwardAsChatCompletionsForGrokUsesXAIChatCompletionsAndSnapshots(t *te require.NoError(t, err) require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String()) require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization")) + require.NotEmpty(t, upstream.lastReq.Header.Get(grokConversationIDHeader)) + require.NotEqual(t, "raw-client-cache-key", upstream.lastReq.Header.Get(grokConversationIDHeader)) require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String()) + require.False(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").Exists()) require.Equal(t, "grok", result.Model) require.Equal(t, "grok-4.5", result.UpstreamModel) require.Equal(t, 1, result.Usage.InputTokens) require.Equal(t, 2, result.Usage.OutputTokens) + require.Equal(t, 1, result.Usage.CacheReadInputTokens) require.NotNil(t, repo.updates[51][grokQuotaSnapshotExtraKey]) require.Equal(t, http.StatusOK, recorder.Code) } @@ -671,6 +768,7 @@ func TestForwardGrokResponsesStreamingUsesXAIResponsesAndSnapshots(t *testing.T) c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) c.Request.Header.Set("Content-Type", "application/json") c.Request.Header.Set("OpenAI-Beta", "responses=experimental") + c.Set("api_key", &APIKey{ID: 5201}) account := &Account{ ID: 52, @@ -719,6 +817,11 @@ func TestForwardGrokResponsesStreamingUsesXAIResponsesAndSnapshots(t *testing.T) require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization")) require.Equal(t, "responses=experimental", upstream.lastReq.Header.Get("OpenAI-Beta")) require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String()) + require.NotEmpty(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()) + require.Equal(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String(), upstream.lastReq.Header.Get(grokConversationIDHeader)) + require.Equal(t, "web_search", gjson.GetBytes(upstream.lastBody, "tools.0.type").String()) + require.Equal(t, "x_search", gjson.GetBytes(upstream.lastBody, "tools.1.type").String()) + require.Equal(t, "none", gjson.GetBytes(upstream.lastBody, "tool_choice").String()) require.Equal(t, "high", gjson.GetBytes(upstream.lastBody, "reasoning_effort").String()) require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool()) require.True(t, result.Stream) @@ -734,6 +837,83 @@ func TestForwardGrokResponsesStreamingUsesXAIResponsesAndSnapshots(t *testing.T) require.NotNil(t, repo.updates[52][grokQuotaSnapshotExtraKey]) } +func TestForwardGrokResponsesAPIKeyUsesXAIResponses(t *testing.T) { + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok","input":"hi","stream":true}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + + account := &Account{ + ID: 53, + Name: "grok-api-key", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 2, + Credentials: map[string]any{ + "api_key": "xai-test-key", + "base_url": "https://api.x.ai/v1", + }, + } + upstreamBody := strings.Join([]string{ + `data: {"type":"response.output_text.delta","sequence_number":0,"delta":"ok"}`, + "", + `data: {"type":"response.completed","sequence_number":1,"response":{"id":"resp_grok_api_key","model":"grok-4.5","usage":{"input_tokens":2,"output_tokens":1}}}`, + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(upstreamBody)), + }} + svc := &OpenAIGatewayService{httpUpstream: upstream} + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", true, time.Now()) + require.NoError(t, err) + require.Equal(t, "https://api.x.ai/v1/responses", upstream.lastReq.URL.String()) + require.Equal(t, "Bearer xai-test-key", upstream.lastReq.Header.Get("Authorization")) + require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String()) + require.Equal(t, "resp_grok_api_key", result.ResponseID) + require.Equal(t, 2, result.Usage.InputTokens) + require.Equal(t, 1, result.Usage.OutputTokens) +} + +func TestAccountTestServiceGrokAPIKeyUsesXAIResponses(t *testing.T) { + gin.SetMode(gin.TestMode) + + account := &Account{ + ID: 54, + Name: "grok-api-key", + Platform: PlatformGrok, + Type: AccountTypeAPIKey, + Concurrency: 2, + Credentials: map[string]any{ + "api_key": "xai-test-key", + "base_url": "https://api.x.ai/v1", + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader( + "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" + + "data: {\"type\":\"response.completed\"}\n\n", + )), + }} + svc := &AccountTestService{httpUpstream: upstream} + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/54/test", nil) + + err := svc.testGrokAccountConnection(c, account, "grok") + require.NoError(t, err) + require.Equal(t, "https://api.x.ai/v1/responses", upstream.lastReq.URL.String()) + require.Equal(t, "Bearer xai-test-key", upstream.lastReq.Header.Get("Authorization")) + require.Contains(t, recorder.Body.String(), `"type":"test_complete"`) +} + func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *testing.T) { gin.SetMode(gin.TestMode) @@ -801,6 +981,205 @@ func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *te require.NotNil(t, repo.updates[53][grokQuotaSnapshotExtraKey]) } +func TestForwardGrokResponsesNonStreamingUsesCacheIdentityAndCachedUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok","input":"hi","stream":false,"tools":[{"type":"namespace","name":"client_tools"}],"tool_choice":{"type":"namespace","name":"client_tools"}}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + c.Set("api_key", &APIKey{ID: 5202}) + + account := &Account{ + ID: 56, + Name: "grok", + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + "base_url": xai.DefaultCLIBaseURL, + }, + } + repo := &grokQuotaAccountRepo{ + mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{56: account}, + }, + } + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"application/json"}, + "Xai-Request-Id": []string{"xai-non-stream-req"}, + }, + Body: io.NopCloser(strings.NewReader(`{"id":"resp_grok_non_stream","object":"response","model":"grok-4.3","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":7,"output_tokens":2,"total_tokens":9,"input_tokens_details":{"cached_tokens":4}}}`)), + }} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now()) + require.NoError(t, err) + require.NotNil(t, result) + require.False(t, result.Stream) + require.Equal(t, "resp_grok_non_stream", result.ResponseID) + require.Equal(t, 7, result.Usage.InputTokens) + require.Equal(t, 2, result.Usage.OutputTokens) + require.Equal(t, 4, result.Usage.CacheReadInputTokens) + require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String()) + identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String() + require.NotEmpty(t, identity) + require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader)) + // The sanitizer drops this unsupported client tool, but its explicit intent + // must still prevent native cache-routing tools from being injected. + require.False(t, gjson.GetBytes(upstream.lastBody, "tools").Exists()) + require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists()) + require.Equal(t, "resp_grok_non_stream", gjson.Get(recorder.Body.String(), "id").String()) +} + +func TestForwardGrokResponsesFailoverKeepsCacheIdentityAcrossAccounts(t *testing.T) { + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok","input":[{"role":"user","content":"stable prefix"}],"stream":false}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 5203}) + + newAccount := func(id int64, token string) *Account { + return &Account{ + ID: id, + Name: fmt.Sprintf("grok-%d", id), + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": token, + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + "base_url": xai.DefaultCLIBaseURL, + }, + } + } + firstAccount := newAccount(58, "access-token-a") + secondAccount := newAccount(59, "access-token-b") + repo := &grokQuotaAccountRepo{ + mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{58: firstAccount, 59: secondAccount}, + }, + } + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + { + StatusCode: http.StatusServiceUnavailable, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"temporary"}}`)), + }, + { + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"application/json"}}, + Body: io.NopCloser(strings.NewReader(`{"id":"resp_after_failover","object":"response","model":"grok-4.3","status":"completed","output":[],"usage":{"input_tokens":5,"output_tokens":1}}`)), + }, + }} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + _, err := svc.forwardGrokResponses(context.Background(), c, firstAccount, body, "grok", false, time.Now()) + var failoverErr *UpstreamFailoverError + require.ErrorAs(t, err, &failoverErr) + + result, err := svc.forwardGrokResponses(context.Background(), c, secondAccount, body, "grok", false, time.Now()) + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, upstream.requests, 2) + require.Len(t, upstream.bodies, 2) + firstIdentity := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String() + secondIdentity := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String() + require.NotEmpty(t, firstIdentity) + require.Equal(t, firstIdentity, secondIdentity) + require.Equal(t, firstIdentity, upstream.requests[0].Header.Get(grokConversationIDHeader)) + require.Equal(t, secondIdentity, upstream.requests[1].Header.Get(grokConversationIDHeader)) + require.Equal(t, "Bearer access-token-a", upstream.requests[0].Header.Get("Authorization")) + require.Equal(t, "Bearer access-token-b", upstream.requests[1].Header.Get("Authorization")) +} + +func TestForwardAsChatCompletionsForGrokStreamingStopFallsBackToRawXAIChatCompletions(t *testing.T) { + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":true,"stop":"done"}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + c.Request.Header.Set(grokConversationIDHeader, "native-client-conversation") + c.Set("api_key", &APIKey{ID: 5301}) + + account := &Account{ + ID: 53, + Name: "grok", + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + "base_url": xai.DefaultCLIBaseURL, + }, + } + repo := &grokQuotaAccountRepo{ + mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{53: account}, + }, + } + upstreamBody := strings.Join([]string{ + `data: {"id":"chatcmpl_grok","object":"chat.completion.chunk","model":"grok-4.3","choices":[{"index":0,"delta":{"content":"ok"}}]}`, + "", + `data: {"id":"chatcmpl_grok","object":"chat.completion.chunk","model":"grok-4.3","choices":[],"usage":{"prompt_tokens":6,"completion_tokens":4,"total_tokens":10,"prompt_tokens_details":{"cached_tokens":1}}}`, + "", + "data: [DONE]", + "", + }, "\n") + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "X-Request-Id": []string{"chat-stream-req"}, + "X-Ratelimit-Limit-Requests": []string{"10"}, + "X-Ratelimit-Remaining-Requests": []string{"7"}, + }, + Body: io.NopCloser(strings.NewReader(upstreamBody)), + }} + svc := &OpenAIGatewayService{ + cfg: rawChatCompletionsTestConfig(), + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "") + require.NoError(t, err) + require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String()) + require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization")) + require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept")) + require.Equal(t, "sub2api-grok/1.0", upstream.lastReq.Header.Get("User-Agent")) + require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version")) + require.NotEmpty(t, upstream.lastReq.Header.Get(grokConversationIDHeader)) + require.NotEqual(t, "native-client-conversation", upstream.lastReq.Header.Get(grokConversationIDHeader)) + require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String()) + require.True(t, gjson.GetBytes(upstream.lastBody, "stream_options.include_usage").Bool()) + require.True(t, result.Stream) + require.Equal(t, 6, result.Usage.InputTokens) + require.Equal(t, 4, result.Usage.OutputTokens) + require.Equal(t, 1, result.Usage.CacheReadInputTokens) + require.Contains(t, recorder.Body.String(), "data: [DONE]") + require.NotNil(t, repo.updates[53][grokQuotaSnapshotExtraKey]) +} + func TestForwardAsChatCompletionsForGrokComposerBridgesImageInput(t *testing.T) { gin.SetMode(gin.TestMode) @@ -809,6 +1188,7 @@ func TestForwardAsChatCompletionsForGrokComposerBridgesImageInput(t *testing.T) body := []byte(`{"model":"grok-composer-2.5-fast","messages":[{"role":"system","content":"You are concise."},{"role":"user","content":[{"type":"text","text":"What is shown?"},{"type":"image_url","image_url":{"url":"data:image/png;base64,QUJD"}}]}],"stream":false}`) c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) c.Request.Header.Set("Content-Type", "application/json") + c.Set("api_key", &APIKey{ID: 5501}) account := &Account{ ID: 55, @@ -858,9 +1238,11 @@ func TestForwardAsChatCompletionsForGrokComposerBridgesImageInput(t *testing.T) require.NotNil(t, result) require.Len(t, upstream.requests, 2) require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.requests[0].URL.String()) + require.Empty(t, upstream.requests[0].Header.Get(grokConversationIDHeader)) require.Equal(t, "grok-build-0.1", gjson.GetBytes(upstream.bodies[0], "model").String()) require.Equal(t, "input_image", gjson.GetBytes(upstream.bodies[0], "input.0.content.1.type").String()) require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.requests[1].URL.String()) + require.NotEmpty(t, upstream.requests[1].Header.Get(grokConversationIDHeader)) require.Equal(t, "grok-composer-2.5-fast", gjson.GetBytes(upstream.bodies[1], "model").String()) require.False(t, strings.Contains(string(upstream.bodies[1]), "image_url")) require.Contains(t, gjson.GetBytes(upstream.bodies[1], "messages.1.content").String(), "Image 1 description") @@ -878,6 +1260,9 @@ func TestForwardAsAnthropicForGrokUsesXAIResponses(t *testing.T) { c, _ := gin.CreateTestContext(recorder) body := []byte(`{"model":"grok","max_tokens":32,"stream":false,"messages":[{"role":"user","content":"hi"}]}`) c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 5401}) + c.Request.Header.Set("OpenAI-Beta", "grok-experimental") + c.Request.Header.Set("originator", "opencode") account := &Account{ ID: 54, @@ -896,7 +1281,7 @@ func TestForwardAsAnthropicForGrokUsesXAIResponses(t *testing.T) { accountsByID: map[int64]*Account{54: account}, }, } - upstream := &httpUpstreamRecorder{resp: openAICompatSSECompletedResponse("resp_grok_messages", "grok-4.3")} + upstream := &httpUpstreamRecorder{resp: grokMessagesSSECompletedResponse("resp_grok_messages", 3)} svc := &OpenAIGatewayService{ httpUpstream: upstream, grokTokenProvider: NewGrokTokenProvider(repo, nil), @@ -908,18 +1293,88 @@ func TestForwardAsAnthropicForGrokUsesXAIResponses(t *testing.T) { require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization")) require.Equal(t, "sub2api-grok/1.0", upstream.lastReq.Header.Get("User-Agent")) + require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version")) + require.Equal(t, "grok-experimental", upstream.lastReq.Header.Get("OpenAI-Beta")) + require.Empty(t, upstream.lastReq.Header.Get("originator")) + require.Empty(t, upstream.lastReq.Header.Get("version")) require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String()) + require.NotEmpty(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()) + require.Equal(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String(), upstream.lastReq.Header.Get(grokConversationIDHeader)) + require.Equal(t, "web_search", gjson.GetBytes(upstream.lastBody, "tools.0.type").String()) + require.Equal(t, "x_search", gjson.GetBytes(upstream.lastBody, "tools.1.type").String()) + require.Equal(t, "none", gjson.GetBytes(upstream.lastBody, "tool_choice").String()) + require.Empty(t, upstream.lastReq.Header.Get("session_id")) require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool()) require.NotContains(t, string(upstream.lastBody), "chatgpt.com") require.Equal(t, "grok", result.Model) require.Equal(t, "grok-4.5", result.UpstreamModel) require.Equal(t, 5, result.Usage.InputTokens) require.Equal(t, 2, result.Usage.OutputTokens) + require.Equal(t, 3, result.Usage.CacheReadInputTokens) require.Contains(t, recorder.Body.String(), `"type":"message"`) + require.Equal(t, int64(3), gjson.Get(recorder.Body.String(), "usage.cache_read_input_tokens").Int()) require.Contains(t, recorder.Body.String(), "ok") } -func TestHandleGrokAccountUpstreamErrorTempUnschedulesReadinessStates(t *testing.T) { +func TestForwardAsAnthropicForGrokStreamingPreservesCacheUsage(t *testing.T) { + gin.SetMode(gin.TestMode) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + body := []byte(`{"model":"grok","max_tokens":32,"stream":true,"messages":[{"role":"user","content":"hi"}]}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body)) + c.Set("api_key", &APIKey{ID: 5402}) + + account := &Account{ + ID: 57, + Name: "grok", + Platform: PlatformGrok, + Type: AccountTypeOAuth, + Concurrency: 1, + Credentials: map[string]any{ + "access_token": "access-token", + "expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339), + "base_url": xai.DefaultCLIBaseURL, + }, + } + repo := &grokQuotaAccountRepo{ + mockAccountRepoForPlatform: &mockAccountRepoForPlatform{ + accountsByID: map[int64]*Account{57: account}, + }, + } + upstream := &httpUpstreamRecorder{resp: grokMessagesSSECompletedResponse("resp_grok_messages_stream", 2)} + svc := &OpenAIGatewayService{ + httpUpstream: upstream, + grokTokenProvider: NewGrokTokenProvider(repo, nil), + accountRepo: repo, + } + + result, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "") + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 2, result.Usage.CacheReadInputTokens) + identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String() + require.NotEmpty(t, identity) + require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader)) + require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream") + require.Contains(t, recorder.Body.String(), `"cache_read_input_tokens":2`) +} + +func grokMessagesSSECompletedResponse(responseID string, cachedTokens int) *http.Response { + body := strings.Join([]string{ + fmt.Sprintf(`data: {"type":"response.completed","response":{"id":%q,"object":"response","model":"grok-4.3","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":5,"output_tokens":2,"total_tokens":7,"input_tokens_details":{"cached_tokens":%d}}}}`, responseID, cachedTokens), + "", + "data: [DONE]", + "", + }, "\n") + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader(body)), + } +} + +func TestHandleGrokAccountUpstreamErrorTempUnschedulesNonRateLimitStates(t *testing.T) { tests := []struct { name string status int @@ -931,24 +1386,23 @@ func TestHandleGrokAccountUpstreamErrorTempUnschedulesReadinessStates(t *testing { name: "unauthorized reauth", status: http.StatusUnauthorized, - wantReason: "grok oauth token unauthorized", + wantReason: "grok credentials unauthorized", wantMinCooldown: 10*time.Minute - time.Second, wantMaxCooldown: 10*time.Minute + time.Second, }, { name: "forbidden entitlement", status: http.StatusForbidden, - wantReason: "grok entitlement or subscription tier denied", + wantReason: "grok access or entitlement denied", wantMinCooldown: 30*time.Minute - time.Second, wantMaxCooldown: 30*time.Minute + time.Second, }, { - name: "rate limited retry after", - status: http.StatusTooManyRequests, - headers: http.Header{"Retry-After": []string{"45"}}, - wantReason: "grok rate limited", - wantMinCooldown: 44 * time.Second, - wantMaxCooldown: 46 * time.Second, + name: "upstream temporary error", + status: http.StatusInternalServerError, + wantReason: "grok upstream temporary error", + wantMinCooldown: 2*time.Minute - time.Second, + wantMaxCooldown: 2*time.Minute + time.Second, }, } @@ -963,6 +1417,7 @@ func TestHandleGrokAccountUpstreamErrorTempUnschedulesReadinessStates(t *testing require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) require.Equal(t, 1, repo.tempUnschedCalls) + require.Zero(t, repo.rateLimitedCalls) require.Equal(t, account.ID, repo.lastTempUnschedID) require.Equal(t, tt.wantReason, repo.lastTempUnschedReason) require.True(t, repo.lastTempUnschedUntil.After(before.Add(tt.wantMinCooldown))) @@ -971,10 +1426,83 @@ func TestHandleGrokAccountUpstreamErrorTempUnschedulesReadinessStates(t *testing } } -func TestHandleGrokAccountUpstreamErrorDoesNotShortenExistingPause(t *testing.T) { +func TestHandleGrokAccountUpstreamError429SetsRateLimitedFromRetryAfter(t *testing.T) { + account := &Account{ID: 61, Platform: PlatformGrok, Type: AccountTypeOAuth} + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + before := time.Now() + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{"Retry-After": []string{"45"}}, nil) + + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) + require.Equal(t, 1, repo.rateLimitedCalls) + require.Equal(t, account.ID, repo.lastRateLimitedID) + require.WithinDuration(t, before.Add(45*time.Second), repo.lastRateLimitResetAt, time.Second) + require.Zero(t, repo.tempUnschedCalls) +} + +func TestHandleGrokAccountUpstreamError429UsesLatestExhaustedWindowReset(t *testing.T) { + now := time.Now() + requestReset := now.Add(10 * time.Minute).Truncate(time.Second) + tokenReset := now.Add(20 * time.Minute).Truncate(time.Second) + headers := http.Header{ + "X-Ratelimit-Limit-Requests": []string{"10"}, + "X-Ratelimit-Remaining-Requests": []string{"0"}, + "X-Ratelimit-Reset-Requests": []string{fmt.Sprintf("%d", requestReset.Unix())}, + "X-Ratelimit-Limit-Tokens": []string{"1000"}, + "X-Ratelimit-Remaining-Tokens": []string{"0"}, + "X-Ratelimit-Reset-Tokens": []string{fmt.Sprintf("%d", tokenReset.Unix())}, + } + account := &Account{ID: 62, Platform: PlatformGrok, Type: AccountTypeOAuth} + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, headers, nil) + + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, tokenReset, repo.lastRateLimitResetAt, time.Second) + require.Zero(t, repo.tempUnschedCalls) +} + +func TestHandleGrokAccountUpstreamError429UsesFallbackReset(t *testing.T) { + account := &Account{ID: 63, Platform: PlatformGrok, Type: AccountTypeOAuth} + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + before := time.Now() + + svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, nil, nil) + + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, before.Add(grokRateLimitFallbackCooldown), repo.lastRateLimitResetAt, time.Second) + require.Zero(t, repo.tempUnschedCalls) +} + +func TestGrokRateLimitResetAtUsesFutureWindowAfterRetryAfterExpires(t *testing.T) { + now := time.Now().UTC().Truncate(time.Second) + observedAt := now.Add(-2 * time.Minute) + windowReset := now.Add(15 * time.Minute) + retryAfter := 30 + snapshot := &xai.QuotaSnapshot{ + StatusCode: http.StatusTooManyRequests, + UpdatedAt: observedAt.Format(time.RFC3339), + RetryAfterSeconds: &retryAfter, + Requests: &xai.QuotaWindow{ + Limit: grokInt64PtrForTest(10), + Remaining: grokInt64PtrForTest(0), + ResetUnix: grokInt64PtrForTest(windowReset.Unix()), + }, + } + + resetAt, limited := grokRateLimitResetAt(snapshot, now) + + require.True(t, limited) + require.WithinDuration(t, windowReset, resetAt, time.Second) +} + +func TestHandleGrokAccountUpstreamError429DoesNotShortenExistingPause(t *testing.T) { existingUntil := time.Now().Add(15 * time.Minute) account := &Account{ - ID: 62, + ID: 64, Platform: PlatformGrok, Type: AccountTypeOAuth, TempUnschedulableUntil: &existingUntil, @@ -985,11 +1513,167 @@ func TestHandleGrokAccountUpstreamErrorDoesNotShortenExistingPause(t *testing.T) svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{"Retry-After": []string{"45"}}, nil) - require.Equal(t, 1, repo.tempUnschedCalls) - require.WithinDuration(t, existingUntil, repo.lastTempUnschedUntil, time.Second) + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, time.Now().Add(45*time.Second), repo.lastRateLimitResetAt, time.Second) + require.Zero(t, repo.tempUnschedCalls) value, ok := svc.openaiAccountRuntimeBlockUntil.Load(account.ID) require.True(t, ok) runtimeUntil, ok := value.(time.Time) require.True(t, ok) require.WithinDuration(t, existingUntil, runtimeUntil, time.Second) } + +func TestUpdateGrokUsageSnapshotExhaustedSuccessBypassesThrottleAndSetsRateLimited(t *testing.T) { + account := &Account{ID: 65, Platform: PlatformGrok, Type: AccountTypeOAuth} + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{ + accountRepo: repo, + codexSnapshotThrottle: newAccountWriteThrottle(time.Hour), + } + now := time.Now() + + // Consume the normal snapshot write allowance first. + svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{ + StatusCode: http.StatusOK, + Requests: &xai.QuotaWindow{ + Limit: grokInt64PtrForTest(10), + Remaining: grokInt64PtrForTest(9), + }, + UpdatedAt: now.UTC().Format(time.RFC3339), + }) + resetAt := now.Add(30 * time.Minute).Truncate(time.Second) + svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{ + StatusCode: http.StatusOK, + Requests: &xai.QuotaWindow{ + Limit: grokInt64PtrForTest(10), + Remaining: grokInt64PtrForTest(0), + ResetUnix: grokInt64PtrForTest(resetAt.Unix()), + ResetAt: resetAt.UTC().Format(time.RFC3339), + }, + UpdatedAt: now.UTC().Format(time.RFC3339), + }) + + require.Equal(t, 2, repo.updateCalls) + require.Equal(t, 1, repo.rateLimitedCalls) + require.Equal(t, account.ID, repo.lastRateLimitedID) + require.WithinDuration(t, resetAt, repo.lastRateLimitResetAt, time.Second) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestUpdateGrokUsageSnapshotAvailableSuccessDoesNotSetRateLimited(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 66, Platform: PlatformGrok, Type: AccountTypeOAuth} + + svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{ + StatusCode: http.StatusOK, + Requests: &xai.QuotaWindow{ + Limit: grokInt64PtrForTest(10), + Remaining: grokInt64PtrForTest(1), + }, + UpdatedAt: time.Now().UTC().Format(time.RFC3339), + }) + + require.Equal(t, 1, repo.updateCalls) + require.Zero(t, repo.rateLimitedCalls) +} + +func TestUpdateGrokUsageSnapshotExhaustedSuccessWithoutResetUsesFallback(t *testing.T) { + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 67, Platform: PlatformGrok, Type: AccountTypeOAuth} + before := time.Now() + + svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{ + StatusCode: http.StatusOK, + Tokens: &xai.QuotaWindow{ + Limit: grokInt64PtrForTest(2_000_000), + Remaining: grokInt64PtrForTest(0), + }, + UpdatedAt: before.UTC().Format(time.RFC3339), + }) + + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, before.Add(grokRateLimitFallbackCooldown), repo.lastRateLimitResetAt, time.Second) + stored, ok := repo.updates[account.ID][grokQuotaSnapshotExtraKey].(*xai.QuotaSnapshot) + require.True(t, ok) + require.NotNil(t, stored.Tokens.ResetUnix) + paused, _ := shouldAutoPauseGrokQuotaWindow("tokens", stored.Tokens, before.Add(time.Second)) + require.True(t, paused) + paused, _ = shouldAutoPauseGrokQuotaWindow("tokens", stored.Tokens, repo.lastRateLimitResetAt.Add(time.Second)) + require.False(t, paused) +} + +func TestOpenAIWSHTTPBridgeGrok429PersistsRateLimit(t *testing.T) { + repo := &grokQuotaAccountRepo{} + upstream := &httpUpstreamRecorder{resp: &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Retry-After": []string{"45"}}, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)), + }} + svc := &OpenAIGatewayService{accountRepo: repo, httpUpstream: upstream} + account := &Account{ID: 68, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1} + before := time.Now() + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), nil, account, "token", + []byte(`{"type":"response.create","model":"grok-4.3","input":"hi"}`), + 64, "grok-4.3", "", "", "", "cache-id", 1, + func([]byte) error { return nil }, + ) + + require.Error(t, err) + require.Nil(t, result) + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, before.Add(45*time.Second), repo.lastRateLimitResetAt, time.Second) + require.Zero(t, repo.tempUnschedCalls) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestOpenAIWSHTTPBridgeGrokExhaustedSuccessPersistsRateLimit(t *testing.T) { + repo := &grokQuotaAccountRepo{} + resetAt := time.Now().Add(20 * time.Minute).UTC().Truncate(time.Second) + resp := grokMessagesSSECompletedResponse("resp_ws_limited", 0) + resp.Header.Set("X-Ratelimit-Limit-Requests", "10") + resp.Header.Set("X-Ratelimit-Remaining-Requests", "0") + resp.Header.Set("X-Ratelimit-Reset-Requests", fmt.Sprintf("%d", resetAt.Unix())) + upstream := &httpUpstreamRecorder{resp: resp} + svc := &OpenAIGatewayService{accountRepo: repo, httpUpstream: upstream} + account := &Account{ID: 69, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1} + + result, err := svc.proxyOpenAIWSHTTPBridgeTurn( + context.Background(), nil, account, "token", + []byte(`{"type":"response.create","model":"grok-4.3","input":"hi"}`), + 64, "grok-4.3", "", "", "", "cache-id", 1, + func([]byte) error { return nil }, + ) + + require.NoError(t, err) + require.NotNil(t, result) + require.Equal(t, 1, repo.rateLimitedCalls) + require.WithinDuration(t, resetAt, repo.lastRateLimitResetAt, time.Second) + require.True(t, svc.isOpenAIAccountRuntimeBlocked(account)) +} + +func TestFailoverOpenAIUpstreamHTTPErrorUsesOnlyGrokRateLimitPolicy(t *testing.T) { + gin.SetMode(gin.TestMode) + repo := &grokQuotaAccountRepo{} + svc := &OpenAIGatewayService{accountRepo: repo} + account := &Account{ID: 70, Platform: PlatformGrok, Type: AccountTypeOAuth} + resp := &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{"Retry-After": []string{"45"}}, + } + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil) + + failoverErr := svc.failoverOpenAIUpstreamHTTPError( + context.Background(), c, account, resp, + []byte(`{"error":{"message":"rate limited"}}`), "rate limited", "grok-4.3", + ) + + require.NotNil(t, failoverErr) + require.Equal(t, 1, repo.rateLimitedCalls) + require.Zero(t, repo.tempUnschedCalls) +} diff --git a/backend/internal/service/openai_gateway_messages.go b/backend/internal/service/openai_gateway_messages.go index 219b5e4be4..8cf62bdc65 100644 --- a/backend/internal/service/openai_gateway_messages.go +++ b/backend/internal/service/openai_gateway_messages.go @@ -244,12 +244,17 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( return nil, policyErr } responsesBody = updatedBody + grokCacheIdentity := "" if account.Platform == PlatformGrok { + grokCacheIdentity = resolveGrokCacheIdentity(c, responsesBody, promptCacheKey, upstreamModel) patchedBody, patchErr := patchGrokResponsesBody(responsesBody, upstreamModel) if patchErr != nil { return nil, patchErr } - responsesBody = patchedBody + responsesBody, patchErr = applyGrokResponsesCacheIdentity(patchedBody, responsesBody, grokCacheIdentity, account.IsGrokOAuth()) + if patchErr != nil { + return nil, fmt.Errorf("apply grok prompt cache identity: %w", patchErr) + } } // 5. Get access token @@ -261,15 +266,14 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( // 6. Build upstream request if account.Type == AccountTypeOAuth && account.Platform != PlatformGrok { // Messages 兼容桥即使 body 未带 todo-guard/prompt_cache_key 标记（如映射到非 - // gpt-5/codex 模型），也必须让 buildUpstreamRequest 走 bridge 分支：不带 - // originator、User-Agent 逐字透传，避免身份收口（issue #3901）误改本路径 - // 刻意最小化的请求形态（下方的 Del(OpenAI-Beta/originator) 兜底保持不变）。 + // gpt-5/codex 模型），也必须让 buildUpstreamRequest 走 bridge 分支，以保留 + // 既有 body/session/conversation 行为。身份头在 post-build 阶段统一恢复。 setOpenAICompatMessagesBridgeContext(c, true) } upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) var upstreamReq *http.Request if account.Platform == PlatformGrok { - upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token) + upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, responsesBody, token, grokCacheIdentity) } else { upstreamReq, err = s.buildUpstreamRequest(upstreamCtx, c, account, responsesBody, token, isStream, promptCacheKey, false) } @@ -280,7 +284,7 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( // Override session_id with a deterministic UUID derived from the isolated // session key, ensuring different API keys produce different upstream sessions. - if promptCacheKey != "" { + if account.Platform != PlatformGrok && promptCacheKey != "" { isolatedSessionID := generateSessionUUID(isolateOpenAISessionID(apiKeyID, promptCacheKey)) upstreamReq.Header.Set("session_id", isolatedSessionID) if upstreamReq.Header.Get("conversation_id") != "" { @@ -288,12 +292,16 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( } } if account.Type == AccountTypeOAuth && account.Platform != PlatformGrok { - // Anthropic Messages compatibility uses the ChatGPT Codex SSE endpoint. - // Match airgate-openai's request shape: the SSE endpoint does not need - // the Responses experimental beta header, and forcing originator can make - // ChatGPT select a different internal continuation path. - upstreamReq.Header.Del("OpenAI-Beta") - upstreamReq.Header.Del("originator") + // buildUpstreamRequest 保留 Messages bridge 的 body/session 兼容行为，并会先 + // 清除身份头。真正发送前恢复完整 Codex 身份，避免 ChatGPT Codex 上游因缺失 + // originator/OpenAI-Beta 返回 404（issue #3901）。 + ensureCodexIdentityHeaders(upstreamReq.Header) + enforceCodexIdentityHeaders(upstreamReq.Header) + logger.L().Debug("openai messages: upstream identity restored", + zap.Int64("account_id", account.ID), + zap.String("upstream_model", upstreamModel), + zap.Bool("compat_identity_restored", true), + ) } if account.Type == AccountTypeOAuth && promptCacheKey != "" && strings.TrimSpace(c.GetHeader("conversation_id")) == "" { upstreamReq.Header.Del("conversation_id") @@ -316,11 +324,6 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( // 8. Handle error response with failover if resp.StatusCode >= 400 { respBody, upstreamMsg := s.readOpenAIUpstreamError(resp) - if account.Platform == PlatformGrok { - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) - s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) - } - if previousResponseID != "" && (isOpenAICompatPreviousResponseNotFound(resp.StatusCode, upstreamMsg, respBody) || isOpenAICompatPreviousResponseUnsupported(resp.StatusCode, upstreamMsg, respBody)) { if isOpenAICompatPreviousResponseUnsupported(resp.StatusCode, upstreamMsg, respBody) { s.disableOpenAICompatSessionContinuation(ctx, c, account, promptCacheKey) @@ -340,6 +343,9 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( // Non-failover error: return Anthropic-formatted error to client return s.handleAnthropicErrorResponse(resp, c, account, billingModel) } + if account.Platform == PlatformGrok && account.Type == AccountTypeOAuth && !account.IsShadow() { + s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + } if account.Type == AccountTypeOAuth && promptCacheKey != "" { if turnState := strings.TrimSpace(resp.Header.Get("x-codex-turn-state")); turnState != "" { @@ -387,10 +393,8 @@ func (s *OpenAIGatewayService) ForwardAsAnthropic( // Extract and save Codex usage snapshot from response headers (for OAuth accounts). // 排除 spark 影子:其 codex_* 仅由 QueryUsage(/wham/usage bengalfox)更新(外审第7轮 P1)。 - if handleErr == nil && account.Type == AccountTypeOAuth && !account.IsShadow() { - if account.Platform == PlatformGrok { - s.updateGrokUsageSnapshot(ctx, account.ID, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) - } else if snapshot := ParseCodexRateLimitHeaders(resp.Header); snapshot != nil { + if handleErr == nil && account.Type == AccountTypeOAuth && !account.IsShadow() && account.Platform != PlatformGrok { + if snapshot := ParseCodexRateLimitHeaders(resp.Header); snapshot != nil { s.updateCodexUsageSnapshot(ctx, account.ID, snapshot) } } diff --git a/backend/internal/service/openai_gateway_messages_chat_fallback.go b/backend/internal/service/openai_gateway_messages_chat_fallback.go index 3861ae9945..a1a7023be4 100644 --- a/backend/internal/service/openai_gateway_messages_chat_fallback.go +++ b/backend/internal/service/openai_gateway_messages_chat_fallback.go @@ -101,7 +101,7 @@ func (s *OpenAIGatewayService) forwardAnthropicViaRawChatCompletions( if err != nil { return nil, err } - resp, err := s.sendCCUpstreamRequest(ctx, c, account, targetURL, chatBody, clientStream, apiKey, account.GetOpenAIUserAgent()) + resp, err := s.sendCCUpstreamRequest(ctx, c, account, targetURL, chatBody, clientStream, apiKey, account.GetOpenAIUserAgent(), "") if err != nil { return nil, err } diff --git a/backend/internal/service/openai_gateway_messages_chat_fallback_test.go b/backend/internal/service/openai_gateway_messages_chat_fallback_test.go index 8bb15c81aa..b36b76ffa8 100644 --- a/backend/internal/service/openai_gateway_messages_chat_fallback_test.go +++ b/backend/internal/service/openai_gateway_messages_chat_fallback_test.go @@ -356,6 +356,8 @@ func TestForwardAsAnthropic_ResponsesSupportedAccountStillUsesResponsesEndpoint( c, _ := gin.CreateTestContext(rec) c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body)) c.Request.Header.Set("Content-Type", "application/json") + c.Request.Header.Set("User-Agent", "third-party-client/1.0.0") + c.Request.Header.Set("originator", "opencode") upstreamBody := strings.Join([]string{ `data: {"type":"response.completed","response":{"id":"resp_native","object":"response","model":"gpt-5.4","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":5,"output_tokens":2,"total_tokens":7}}}`, @@ -385,5 +387,9 @@ func TestForwardAsAnthropic_ResponsesSupportedAccountStillUsesResponsesEndpoint( "responses-capable account must stay on /v1/responses, got %s", upstream.lastReq.URL.String()) require.True(t, gjson.GetBytes(upstream.lastBody, "input").Exists()) require.False(t, gjson.GetBytes(upstream.lastBody, "messages").Exists()) + require.Equal(t, "third-party-client/1.0.0", upstream.lastReq.Header.Get("User-Agent")) + require.Equal(t, "opencode", upstream.lastReq.Header.Get("originator")) + require.Empty(t, upstream.lastReq.Header.Get("version")) + require.Empty(t, upstream.lastReq.Header.Get("OpenAI-Beta")) require.Equal(t, "ok", gjson.Get(rec.Body.String(), "content.0.text").String()) } diff --git a/backend/internal/service/openai_gateway_responses_chat_fallback.go b/backend/internal/service/openai_gateway_responses_chat_fallback.go index 9d3391e4e6..f494429226 100644 --- a/backend/internal/service/openai_gateway_responses_chat_fallback.go +++ b/backend/internal/service/openai_gateway_responses_chat_fallback.go @@ -92,7 +92,7 @@ func (s *OpenAIGatewayService) forwardResponsesViaRawChatCompletions( if err != nil { return nil, err } - resp, err := s.sendCCUpstreamRequest(ctx, c, account, targetURL, chatBody, clientStream, apiKey, account.GetOpenAIUserAgent()) + resp, err := s.sendCCUpstreamRequest(ctx, c, account, targetURL, chatBody, clientStream, apiKey, account.GetOpenAIUserAgent(), "") if err != nil { return nil, err } diff --git a/backend/internal/service/openai_gateway_scheduling.go b/backend/internal/service/openai_gateway_scheduling.go index 6ac0de47b4..f318adc604 100644 --- a/backend/internal/service/openai_gateway_scheduling.go +++ b/backend/internal/service/openai_gateway_scheduling.go @@ -23,17 +23,7 @@ import ( // ExtractSessionID extracts the raw session ID from headers or body without hashing. // Used by ForwardAsAnthropic to pass as prompt_cache_key for upstream cache. func (s *OpenAIGatewayService) ExtractSessionID(c *gin.Context, body []byte) string { - if c == nil { - return "" - } - sessionID := strings.TrimSpace(c.GetHeader("session_id")) - if sessionID == "" { - sessionID = strings.TrimSpace(c.GetHeader("conversation_id")) - } - if sessionID == "" && len(body) > 0 { - sessionID = strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String()) - } - return sessionID + return explicitOpenAIRequestSessionID(c, body) } func explicitOpenAISessionID(c *gin.Context, body []byte) string { @@ -51,11 +41,33 @@ func explicitOpenAISessionID(c *gin.Context, body []byte) string { return sessionID } +// explicitOpenAIRequestSessionID extends the common OpenAI session signals +// with Grok's native conversation header only for requests authenticated to a +// Grok group. This keeps an unrelated x-grok-conv-id header from changing +// scheduling or upstream session behavior for non-Grok groups. +func explicitOpenAIRequestSessionID(c *gin.Context, body []byte) string { + if c == nil { + return "" + } + + sessionID := strings.TrimSpace(c.GetHeader("session_id")) + if sessionID == "" { + sessionID = strings.TrimSpace(c.GetHeader("conversation_id")) + } + if sessionID == "" && isGrokRequestContext(c) { + sessionID = strings.TrimSpace(c.GetHeader(grokConversationIDHeader)) + } + if sessionID == "" && len(body) > 0 { + sessionID = strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String()) + } + return sessionID +} + // GenerateExplicitSessionHash generates a sticky-session hash only from explicit // client session signals. It intentionally skips content-derived fallback and is // used by stateless endpoints such as /v1/images. func (s *OpenAIGatewayService) GenerateExplicitSessionHash(c *gin.Context, body []byte) string { - sessionID := explicitOpenAISessionID(c, body) + sessionID := explicitOpenAIRequestSessionID(c, body) if sessionID == "" { return "" } @@ -70,14 +82,15 @@ func (s *OpenAIGatewayService) GenerateExplicitSessionHash(c *gin.Context, body // Priority: // 1. Header: session_id // 2. Header: conversation_id -// 3. Body: prompt_cache_key (opencode) -// 4. Body: content-based fallback (model + system + tools + first user message) +// 3. Header: x-grok-conv-id (Grok groups only) +// 4. Body: prompt_cache_key (opencode) +// 5. Body: content-based fallback (model + system + tools + first user message) func (s *OpenAIGatewayService) GenerateSessionHash(c *gin.Context, body []byte) string { if c == nil { return "" } - sessionID := explicitOpenAISessionID(c, body) + sessionID := explicitOpenAIRequestSessionID(c, body) if sessionID == "" && len(body) > 0 { sessionID = deriveOpenAIContentSessionSeed(body) } diff --git a/backend/internal/service/openai_gateway_service.go b/backend/internal/service/openai_gateway_service.go index 0d03d83dcf..29c7d968a2 100644 --- a/backend/internal/service/openai_gateway_service.go +++ b/backend/internal/service/openai_gateway_service.go @@ -49,6 +49,7 @@ const ( openAIWSRetryBackoffMaxDefault = 2 * time.Second openAIWSRetryJitterRatioDefault = 0.2 openAICompactSessionSeedKey = "openai_compact_session_seed" + openAIUpstreamEndpointContextKey = "openai_actual_upstream_endpoint" codexCLIVersion = "0.144.1" // Codex 限额快照仅用于后台展示/诊断，不需要每个成功请求都立即落库。 openAICodexSnapshotPersistMinInterval = 30 * time.Second @@ -66,6 +67,7 @@ var openaiAllowedHeaders = map[string]bool{ "user-agent": true, "originator": true, "session_id": true, + "x-codex-beta-features": true, "x-codex-turn-state": true, "x-codex-turn-metadata": true, } @@ -81,6 +83,7 @@ var openaiPassthroughAllowedHeaders = map[string]bool{ "user-agent": true, "originator": true, "session_id": true, + "x-codex-beta-features": true, "x-codex-turn-state": true, "x-codex-turn-metadata": true, } @@ -223,6 +226,9 @@ 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 + // 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 // ServiceTier records the OpenAI Responses API service tier, e.g. "priority" / "flex". // Nil means the request did not specify a recognized tier. ServiceTier *string @@ -246,11 +252,40 @@ type OpenAIForwardResult struct { VideoResolution string // VideoDurationSeconds 是提交时请求的生成时长（xAI 按输出秒数计费），已归一化到 1-15 秒。 VideoDurationSeconds int + // WebSearchCalls 是 Codex alpha/search 网页搜索调用次数（每次成功请求为 1）。 + // 上游不返回 usage 字段，>0 时走按次计费（分组单价 × 次数 × 倍率）。 + WebSearchCalls int wsReplayInput []json.RawMessage wsReplayInputExists bool } +// SetActualOpenAIUpstreamEndpoint records the endpoint selected by the current +// forwarding attempt. It covers error paths where no OpenAIForwardResult is +// available for usage and operations logging. +func SetActualOpenAIUpstreamEndpoint(c *gin.Context, endpoint string) { + if c == nil { + return + } + if endpoint = strings.TrimSpace(endpoint); endpoint != "" { + c.Set(openAIUpstreamEndpointContextKey, endpoint) + } +} + +// GetActualOpenAIUpstreamEndpoint returns the endpoint recorded by the latest +// forwarding attempt in this request. +func GetActualOpenAIUpstreamEndpoint(c *gin.Context) string { + if c == nil { + return "" + } + value, exists := c.Get(openAIUpstreamEndpointContextKey) + if !exists { + return "" + } + endpoint, _ := value.(string) + return strings.TrimSpace(endpoint) +} + type OpenAIWSRetryMetricsSnapshot struct { RetryAttemptsTotal int64 `json:"retry_attempts_total"` RetryBackoffMsTotal int64 `json:"retry_backoff_ms_total"` diff --git a/backend/internal/service/openai_gateway_usage.go b/backend/internal/service/openai_gateway_usage.go index 6e9760bf77..f030be6b46 100644 --- a/backend/internal/service/openai_gateway_usage.go +++ b/backend/internal/service/openai_gateway_usage.go @@ -187,6 +187,7 @@ func (s *OpenAIGatewayService) RecordUsage(ctx context.Context, input *OpenAIRec multiplier, imageMultiplier, videoMultiplier, + baseMultiplier, tokens, serviceTier, longContextBillingEnabled, @@ -376,11 +377,19 @@ func (s *OpenAIGatewayService) calculateOpenAIRecordUsageCost( multiplier float64, imageMultiplier float64, videoMultiplier float64, + webSearchMultiplier float64, tokens UsageTokens, serviceTier string, longContextBillingEnabled bool, ) (*CostBreakdown, error) { billingModel := firstUsageBillingModel(billingModels) + if result != nil && result.WebSearchCalls > 0 { + // Codex alpha/search 网页搜索按次计费：上游不返回 usage/token 字段，单价只取 + // 分组覆盖价（nil 时默认 0.01 = 官方 $10/1000 次），不参与渠道级模型定价。 + // 倍率与 image/video 按次口径一致：使用不含高峰因子的基础倍率 + //（用户专属 > 分组 rate_multiplier > 系统默认），与分组表单的价格预览承诺一致。 + return s.billingService.CalculateWebSearchCost(result.WebSearchCalls, webSearchPricePerCallFromAPIKey(apiKey), webSearchMultiplier), nil + } if isGrokVideoUsageResult(result, billingModels) { if resolved := s.resolveOpenAIChannelPricing(ctx, billingModel, apiKey); resolved == nil || resolved.Mode != BillingModeToken { return s.calculateOpenAIVideoCost(ctx, billingModel, apiKey, result, videoMultiplier), nil diff --git a/backend/internal/service/openai_gpt56_max_test.go b/backend/internal/service/openai_gpt56_max_test.go index 272eb16ff0..cbca2ff3ee 100644 --- a/backend/internal/service/openai_gpt56_max_test.go +++ b/backend/internal/service/openai_gpt56_max_test.go @@ -223,13 +223,17 @@ func TestOpenAIGatewayServiceForwardOAuthCompactDowngradesMaxEffort(t *testing.T require.Equal(t, "xhigh", *result.ReasoningEffort) } -func TestOpenAIGatewayServiceForwardOAuthResponsesPreservesMaxEffort(t *testing.T) { +func TestOpenAIGatewayServiceForwardOAuthRemoteCompactV2PreservesResponsesWire(t *testing.T) { gin.SetMode(gin.TestMode) upstream := &httpUpstreamRecorder{ resp: &http.Response{ StatusCode: http.StatusOK, - Header: http.Header{"Content-Type": []string{"application/json"}}, - Body: io.NopCloser(strings.NewReader(`{"usage":{"input_tokens":1,"output_tokens":2}}`)), + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader( + "data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"compaction\",\"encrypted_content\":\"summary\"}}\n\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":1,\"output_tokens\":2}}}\n\n" + + "data: [DONE]\n\n", + )), }, } cfg := &config.Config{} @@ -244,6 +248,9 @@ func TestOpenAIGatewayServiceForwardOAuthResponsesPreservesMaxEffort(t *testing. Credentials: map[string]any{ "access_token": "oauth-token", "chatgpt_account_id": "chatgpt-acc", + "compact_model_mapping": map[string]any{ + "gpt-5.6-sol": "gpt-5.6-sol-openai-compact", + }, }, Status: StatusActive, Schedulable: true, @@ -251,16 +258,82 @@ func TestOpenAIGatewayServiceForwardOAuthResponsesPreservesMaxEffort(t *testing. rec := httptest.NewRecorder() c, _ := gin.CreateTestContext(rec) c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + c.Request.Header.Set("x-codex-beta-features", "remote_compaction_v2") SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) - body := []byte(`{"model":"gpt-5.6-sol","instructions":"response-test","input":"hello","reasoning":{"effort":"max"}}`) + body := []byte(`{"model":"gpt-5.6-sol","stream":true,"instructions":"response-test","input":[{"type":"message","role":"user","content":"hello"},{"type":"compaction_trigger"}],"reasoning":{"effort":"max","context":"all_turns"}}`) result, err := svc.Forward(context.Background(), c, account, body) require.NoError(t, err) require.NotNil(t, result) require.NotNil(t, upstream.lastReq) require.Equal(t, chatgptCodexURL, upstream.lastReq.URL.String()) + require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(upstream.lastBody, "model").String()) + require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool()) + require.Equal(t, "compaction_trigger", gjson.GetBytes(upstream.lastBody, "input.#(type==\"compaction_trigger\").type").String()) require.Equal(t, "max", gjson.GetBytes(upstream.lastBody, "reasoning.effort").String()) + require.Equal(t, "all_turns", gjson.GetBytes(upstream.lastBody, "reasoning.context").String()) + require.Equal(t, "remote_compaction_v2", upstream.lastReq.Header.Get("x-codex-beta-features")) + require.Contains(t, rec.Body.String(), `"type":"compaction"`) + require.Contains(t, rec.Body.String(), `"encrypted_content":"summary"`) + require.NotNil(t, result.ReasoningEffort) + require.Equal(t, "max", *result.ReasoningEffort) +} + +func TestOpenAIGatewayServiceForwardAPIKeyRemoteCompactV2PreservesResponsesWire(t *testing.T) { + gin.SetMode(gin.TestMode) + upstream := &httpUpstreamRecorder{ + resp: &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": []string{"text/event-stream"}}, + Body: io.NopCloser(strings.NewReader( + "data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"compaction\",\"encrypted_content\":\"summary\"}}\n\n" + + "data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":1,\"output_tokens\":2}}}\n\n" + + "data: [DONE]\n\n", + )), + }, + } + cfg := &config.Config{} + cfg.Security.URLAllowlist.Enabled = false + svc := &OpenAIGatewayService{cfg: cfg, httpUpstream: upstream} + account := &Account{ + ID: 11, + Name: "openai-apikey-responses", + Platform: PlatformOpenAI, + Type: AccountTypeAPIKey, + Concurrency: 1, + Credentials: map[string]any{ + "api_key": "sk-test", + "base_url": "https://example.com/v1", + "compact_model_mapping": map[string]any{ + "gpt-5.6-sol": "gpt-5.6-sol-openai-compact", + }, + }, + Extra: map[string]any{"use_responses_api": true}, + Status: StatusActive, + Schedulable: true, + } + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil) + c.Request.Header.Set("x-codex-beta-features", "remote_compaction_v2") + SetOpenAIClientTransport(c, OpenAIClientTransportHTTP) + + body := []byte(`{"model":"gpt-5.6-sol","stream":true,"instructions":"response-test","input":[{"type":"message","role":"user","content":"hello"},{"type":"compaction_trigger"}],"reasoning":{"effort":"max","context":"all_turns"}}`) + result, err := svc.Forward(context.Background(), c, account, body) + + require.NoError(t, err) + require.NotNil(t, result) + require.NotNil(t, upstream.lastReq) + require.Equal(t, "https://example.com/v1/responses", upstream.lastReq.URL.String()) + require.Equal(t, "gpt-5.6-sol", gjson.GetBytes(upstream.lastBody, "model").String()) + require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool()) + require.Equal(t, "compaction_trigger", gjson.GetBytes(upstream.lastBody, "input.#(type==\"compaction_trigger\").type").String()) + require.Equal(t, "max", gjson.GetBytes(upstream.lastBody, "reasoning.effort").String()) + require.Equal(t, "all_turns", gjson.GetBytes(upstream.lastBody, "reasoning.context").String()) + require.Equal(t, "remote_compaction_v2", upstream.lastReq.Header.Get("x-codex-beta-features")) + require.Contains(t, rec.Body.String(), `"type":"compaction"`) + require.Contains(t, rec.Body.String(), `"encrypted_content":"summary"`) require.NotNil(t, result.ReasoningEffort) require.Equal(t, "max", *result.ReasoningEffort) } diff --git a/backend/internal/service/openai_oauth_passthrough_test.go b/backend/internal/service/openai_oauth_passthrough_test.go index de8ecf030b..60790b9e4c 100644 --- a/backend/internal/service/openai_oauth_passthrough_test.go +++ b/backend/internal/service/openai_oauth_passthrough_test.go @@ -347,6 +347,7 @@ func TestOpenAIGatewayService_OAuthPassthrough_StreamKeepsToolNameAndBodyNormali c.Request.Header.Set("Accept-Encoding", "gzip") c.Request.Header.Set("Proxy-Authorization", "Basic abc") c.Request.Header.Set("X-Test", "keep") + c.Request.Header.Set("x-codex-beta-features", "remote_compaction_v2") originalBody := []byte(`{"model":"gpt-5.2","stream":true,"store":true,"instructions":"local-test-instructions","input":[{"type":"text","text":"hi"}]}`) @@ -409,6 +410,7 @@ func TestOpenAIGatewayService_OAuthPassthrough_StreamKeepsToolNameAndBodyNormali require.Empty(t, upstream.lastReq.Header.Get("Accept-Encoding")) require.Empty(t, upstream.lastReq.Header.Get("Proxy-Authorization")) require.Empty(t, upstream.lastReq.Header.Get("X-Test")) + require.Equal(t, "remote_compaction_v2", upstream.lastReq.Header.Get("x-codex-beta-features")) // 3) required OAuth headers are present require.Equal(t, "chatgpt.com", upstream.lastReq.Host) @@ -1373,6 +1375,7 @@ func TestOpenAIGatewayService_APIKeyPassthrough_PreservesBodyAndUsesResponsesEnd c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(nil)) c.Request.Header.Set("User-Agent", "curl/8.0") c.Request.Header.Set("X-Test", "keep") + c.Request.Header.Set("x-codex-beta-features", "remote_compaction_v2") originalBody := []byte(`{"model":"gpt-5.2","stream":false,"service_tier":"flex","max_output_tokens":128,"input":[{"type":"text","text":"hi"}]}`) resp := &http.Response{ @@ -1410,6 +1413,7 @@ func TestOpenAIGatewayService_APIKeyPassthrough_PreservesBodyAndUsesResponsesEnd require.Equal(t, "https://api.openai.com/v1/responses", upstream.lastReq.URL.String()) require.Equal(t, "Bearer sk-api-key", upstream.lastReq.Header.Get("Authorization")) require.Equal(t, "curl/8.0", upstream.lastReq.Header.Get("User-Agent")) + require.Equal(t, "remote_compaction_v2", upstream.lastReq.Header.Get("x-codex-beta-features")) require.Empty(t, upstream.lastReq.Header.Get("X-Test")) } diff --git a/backend/internal/service/openai_ws_forwarder_ingress.go b/backend/internal/service/openai_ws_forwarder_ingress.go index 45e004b6d0..a2af4b760f 100644 --- a/backend/internal/service/openai_ws_forwarder_ingress.go +++ b/backend/internal/service/openai_ws_forwarder_ingress.go @@ -421,6 +421,10 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( storeDisabled, ) currentBridgePayload := firstPayload + // Keep the first turn as the stable conversation seed. The mapped model + // is resolved again for each turn below so an in-connection model switch + // cannot reuse another model's upstream cache identity. + grokCacheSeedPayload := firstPayload.payloadRaw var bridgeReplayInput []json.RawMessage bridgeReplayInputExists := false for turn := 1; ; turn++ { @@ -469,6 +473,13 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( openAIWSRawPayloadHasToolCallOutput(currentBridgePayload.payloadRaw), ) } + grokCacheIdentity := "" + if account.Platform == PlatformGrok { + grokCacheIdentity, err = resolveGrokWSCacheIdentity(c, account, grokCacheSeedPayload, currentBridgePayload.originalModel) + if err != nil { + return fmt.Errorf("resolve Grok websocket cache identity: %w", err) + } + } result, bridgeErr := s.proxyOpenAIWSHTTPBridgeTurn( ctx, c, @@ -480,6 +491,7 @@ func (s *OpenAIGatewayService) ProxyResponsesWebSocketFromClient( currentBridgePayload.imageBillingModel, currentBridgePayload.imageSizeTier, currentBridgePayload.imageInputSize, + grokCacheIdentity, turn, writeClientMessage, ) diff --git a/backend/internal/service/openai_ws_forwarder_payload.go b/backend/internal/service/openai_ws_forwarder_payload.go index a4d47218e7..5830444815 100644 --- a/backend/internal/service/openai_ws_forwarder_payload.go +++ b/backend/internal/service/openai_ws_forwarder_payload.go @@ -74,6 +74,11 @@ func (s *OpenAIGatewayService) buildOpenAIWSHeaders( if v := strings.TrimSpace(c.Request.Header.Get("accept-language")); v != "" { headers.Set("accept-language", v) } + for _, value := range c.Request.Header.Values("x-codex-beta-features") { + if value = strings.TrimSpace(value); value != "" { + headers.Add("x-codex-beta-features", value) + } + } } // OAuth 账号：将 apiKeyID 混入 session 标识符，防止跨用户会话碰撞。 if account != nil && account.Type == AccountTypeOAuth { diff --git a/backend/internal/service/openai_ws_forwarder_success_test.go b/backend/internal/service/openai_ws_forwarder_success_test.go index adae109e09..bb4ac2242c 100644 --- a/backend/internal/service/openai_ws_forwarder_success_test.go +++ b/backend/internal/service/openai_ws_forwarder_success_test.go @@ -602,6 +602,7 @@ func TestOpenAIGatewayService_Forward_WSv2_OAuthStoreFalseByDefault(t *testing.T c.Request.Header.Set("User-Agent", "codex_cli_rs/0.98.0") c.Request.Header.Set("session_id", "sess-oauth-1") c.Request.Header.Set("conversation_id", "conv-oauth-1") + c.Request.Header.Set("x-codex-beta-features", "remote_compaction_v2") cfg := &config.Config{} cfg.Security.URLAllowlist.Enabled = false @@ -661,6 +662,7 @@ func TestOpenAIGatewayService_Forward_WSv2_OAuthStoreFalseByDefault(t *testing.T require.True(t, gjson.Get(requestJSON, "stream").Exists(), "WSv2 payload 应保留 stream 字段") require.True(t, gjson.Get(requestJSON, "stream").Bool(), "OAuth Codex 规范化后应强制 stream=true") require.Equal(t, openAIWSBetaV2Value, captureDialer.lastHeaders.Get("OpenAI-Beta")) + require.Equal(t, "remote_compaction_v2", captureDialer.lastHeaders.Get("x-codex-beta-features")) // OAuth 账号的 session_id/conversation_id 应被 isolateOpenAISessionID 隔离， // 测试中未设置 api_key 到 context，apiKeyID=0。 require.Equal(t, isolateOpenAISessionID(0, "sess-oauth-1"), captureDialer.lastHeaders.Get("session_id")) diff --git a/backend/internal/service/openai_ws_http_bridge.go b/backend/internal/service/openai_ws_http_bridge.go index dce0c7b6db..0afb9181a1 100644 --- a/backend/internal/service/openai_ws_http_bridge.go +++ b/backend/internal/service/openai_ws_http_bridge.go @@ -12,6 +12,7 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/pkg/xai" "github.com/gin-gonic/gin" "github.com/tidwall/gjson" ) @@ -155,6 +156,7 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( imageBillingModel string, imageSizeTier string, imageInputSize string, + grokCacheIdentity string, turn int, writeClientMessage func([]byte) error, ) (*OpenAIForwardResult, error) { @@ -179,21 +181,19 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx) var upstreamReq *http.Request if account.Platform == PlatformGrok { - upstreamModel := strings.TrimSpace(gjson.GetBytes(body, "model").String()) - if originalModel != "" { - if mappedModel := normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel)); mappedModel != "" { - upstreamModel = mappedModel - } - } - if upstreamModel == "" { - upstreamModel = "grok-4.3" - } + upstreamModel := resolveGrokWSUpstreamModel(account, body, originalModel) + grokIntentSourceBody := body body, err = patchGrokResponsesBody(body, upstreamModel) if err != nil { releaseUpstreamCtx() return nil, err } - upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, body, token) + body, err = applyGrokResponsesCacheIdentity(body, grokIntentSourceBody, grokCacheIdentity, account.IsGrokOAuth()) + if err != nil { + releaseUpstreamCtx() + return nil, fmt.Errorf("apply grok prompt cache identity: %w", err) + } + upstreamReq, err = buildGrokResponsesRequest(upstreamCtx, c, account, body, token, grokCacheIdentity) } else { upstreamReq, err = s.buildUpstreamRequestOpenAIPassthrough(upstreamCtx, c, account, body, token) } @@ -222,6 +222,9 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( if resp.StatusCode >= 400 { respBody, _ := io.ReadAll(io.LimitReader(resp.Body, openAIWSHTTPBridgeErrorBodyLimitBytes)) + if account.Platform == PlatformGrok { + s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, respBody) + } upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(respBody))) if upstreamMsg == "" { upstreamMsg = http.StatusText(resp.StatusCode) @@ -229,6 +232,9 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( _ = writeClientMessage(buildOpenAIWSHTTPBridgeErrorEvent(resp.StatusCode, upstreamMsg)) return nil, fmt.Errorf("upstream http bridge error: status=%d message=%s", resp.StatusCode, upstreamMsg) } + if account.Platform == PlatformGrok { + s.updateGrokUsageSnapshot(ctx, account, xai.ParseQuotaHeaders(resp.Header, resp.StatusCode)) + } responseID := "" usage := OpenAIUsage{} @@ -407,3 +413,25 @@ func (s *OpenAIGatewayService) proxyOpenAIWSHTTPBridgeTurn( } return resultWithUsage(), errors.New("upstream http bridge stream ended before terminal event") } + +func resolveGrokWSCacheIdentity(c *gin.Context, account *Account, payload []byte, originalModel string) (string, error) { + body, err := prepareOpenAIWSHTTPBridgeBody(payload) + if err != nil { + return "", err + } + upstreamModel := resolveGrokWSUpstreamModel(account, body, originalModel) + return resolveGrokCacheIdentity(c, body, "", upstreamModel), nil +} + +func resolveGrokWSUpstreamModel(account *Account, body []byte, originalModel string) string { + upstreamModel := strings.TrimSpace(gjson.GetBytes(body, "model").String()) + if account != nil && originalModel != "" { + if mappedModel := normalizeOpenAIModelForUpstream(account, account.GetMappedModel(originalModel)); mappedModel != "" { + upstreamModel = mappedModel + } + } + if upstreamModel == "" { + upstreamModel = "grok-4.3" + } + return upstreamModel +} diff --git a/backend/internal/service/openai_ws_http_bridge_test.go b/backend/internal/service/openai_ws_http_bridge_test.go index 4d1e4a0374..da2ae77917 100644 --- a/backend/internal/service/openai_ws_http_bridge_test.go +++ b/backend/internal/service/openai_ws_http_bridge_test.go @@ -3,6 +3,7 @@ package service import ( "context" "errors" + "fmt" "io" "net/http" "net/http/httptest" @@ -127,6 +128,7 @@ func TestOpenAIWSHTTPBridgeRelaysSSEFramesAsWebSocketMessages(t *testing.T) { "", "", "", + "", 1, writeClient, ) @@ -179,21 +181,28 @@ func TestOpenAIWSHTTPBridgeRelaysSSEFramesAsWebSocketMessages(t *testing.T) { func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridge(t *testing.T) { gin.SetMode(gin.TestMode) - sseBody := strings.Join([]string{ - `data: {"type":"response.created","response":{"id":"resp_grok_ws","model":"grok-4.3"}}`, - "", - `data: {"type":"response.output_text.delta","response":{"id":"resp_grok_ws"},"delta":"ok"}`, - "", - `data: {"type":"response.completed","response":{"id":"resp_grok_ws","model":"grok-4.3","usage":{"input_tokens":4,"output_tokens":2}}}`, - "", - }, "\n") - upstream := &httpUpstreamRecorder{resp: &http.Response{ - StatusCode: http.StatusOK, - Header: http.Header{ - "Content-Type": []string{"text/event-stream"}, - "Xai-Request-Id": []string{"xai-ws-req"}, - }, - Body: io.NopCloser(strings.NewReader(sseBody)), + bridgeResponse := func(responseID, requestID string, cachedTokens int) *http.Response { + sseBody := strings.Join([]string{ + `data: {"type":"response.created","response":{"id":"` + responseID + `","model":"grok-4.3"}}`, + "", + `data: {"type":"response.output_text.delta","response":{"id":"` + responseID + `"},"delta":"ok"}`, + "", + `data: {"type":"response.completed","response":{"id":"` + responseID + `","model":"grok-4.3","usage":{"input_tokens":4,"output_tokens":2,"input_tokens_details":{"cached_tokens":` + fmt.Sprintf("%d", cachedTokens) + `}}}}`, + "", + }, "\n") + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Type": []string{"text/event-stream"}, + "Xai-Request-Id": []string{requestID}, + }, + Body: io.NopCloser(strings.NewReader(sseBody)), + } + } + upstream := &httpUpstreamRecorder{responses: []*http.Response{ + bridgeResponse("resp_grok_ws_1", "xai-ws-req-1", 0), + bridgeResponse("resp_grok_ws_2", "xai-ws-req-2", 3), + bridgeResponse("resp_grok_ws_3", "xai-ws-req-3", 0), }} svc := &OpenAIGatewayService{ cfg: &config.Config{ @@ -241,6 +250,7 @@ func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridge(t *testing.T) req := r.Clone(r.Context()) req.Header = req.Header.Clone() ginCtx.Request = req + ginCtx.Set("api_key", &APIKey{ID: 7101}) errCh <- svc.ProxyResponsesWebSocketFromClient(r.Context(), ginCtx, conn, account, "access-token", firstMessage, nil) })) @@ -271,6 +281,33 @@ func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridge(t *testing.T) require.Equal(t, "response.created", gjson.GetBytes(created, "type").String()) require.Equal(t, "response.output_text.delta", gjson.GetBytes(delta, "type").String()) require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String()) + require.Equal(t, "resp_grok_ws_1", gjson.GetBytes(completed, "response.id").String()) + + writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","generate":true,"model":"grok","stream":true,"previous_response_id":"resp_grok_ws_1","input":"second turn"}`)) + cancelWrite() + require.NoError(t, err) + + created = readEvent() + delta = readEvent() + completed = readEvent() + require.Equal(t, "response.created", gjson.GetBytes(created, "type").String()) + require.Equal(t, "response.output_text.delta", gjson.GetBytes(delta, "type").String()) + require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String()) + require.Equal(t, "resp_grok_ws_2", gjson.GetBytes(completed, "response.id").String()) + + writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second) + err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","generate":true,"model":"grok-4.3","stream":true,"previous_response_id":"resp_grok_ws_2","input":"third turn with a different model"}`)) + cancelWrite() + require.NoError(t, err) + + created = readEvent() + delta = readEvent() + completed = readEvent() + require.Equal(t, "response.created", gjson.GetBytes(created, "type").String()) + require.Equal(t, "response.output_text.delta", gjson.GetBytes(delta, "type").String()) + require.Equal(t, "response.completed", gjson.GetBytes(completed, "type").String()) + require.Equal(t, "resp_grok_ws_3", gjson.GetBytes(completed, "response.id").String()) _ = clientConn.Close(coderws.StatusNormalClosure, "done") select { @@ -280,10 +317,30 @@ func TestProxyResponsesWebSocketFromClientForGrokUsesXAIHTTPBridge(t *testing.T) require.Fail(t, "proxy did not finish after client close") } + require.Len(t, upstream.requests, 3) + require.Len(t, upstream.bodies, 3) require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String()) require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization")) require.Equal(t, "sub2api-grok/1.0", upstream.lastReq.Header.Get("User-Agent")) - require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String()) + require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version")) + require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.bodies[0], "model").String()) + require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.bodies[1], "model").String()) + require.Equal(t, "grok-4.3", gjson.GetBytes(upstream.bodies[2], "model").String()) + require.NotEmpty(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()) + require.Equal(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String(), upstream.lastReq.Header.Get(grokConversationIDHeader)) + require.Equal(t, "web_search", gjson.GetBytes(upstream.lastBody, "tools.0.type").String()) + require.Equal(t, "x_search", gjson.GetBytes(upstream.lastBody, "tools.1.type").String()) + require.Equal(t, "none", gjson.GetBytes(upstream.lastBody, "tool_choice").String()) + firstIdentity := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String() + secondIdentity := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String() + thirdIdentity := gjson.GetBytes(upstream.bodies[2], "prompt_cache_key").String() + require.NotEmpty(t, firstIdentity) + require.Equal(t, firstIdentity, secondIdentity) + require.NotEmpty(t, thirdIdentity) + require.NotEqual(t, firstIdentity, thirdIdentity) + require.Equal(t, firstIdentity, upstream.requests[0].Header.Get(grokConversationIDHeader)) + require.Equal(t, secondIdentity, upstream.requests[1].Header.Get(grokConversationIDHeader)) + require.Equal(t, thirdIdentity, upstream.requests[2].Header.Get(grokConversationIDHeader)) require.False(t, gjson.GetBytes(upstream.lastBody, "type").Exists()) require.False(t, gjson.GetBytes(upstream.lastBody, "generate").Exists()) require.False(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_retention").Exists()) diff --git a/backend/internal/service/openai_ws_pool.go b/backend/internal/service/openai_ws_pool.go index 5950e02841..329908e762 100644 --- a/backend/internal/service/openai_ws_pool.go +++ b/backend/internal/service/openai_ws_pool.go @@ -218,6 +218,9 @@ func (l *openAIWSConnLease) Release() { return } l.conn.release() + if l.pool != nil { + l.pool.notifyAccountPoolChanged(l.accountID) + } } type openAIWSConn struct { @@ -225,6 +228,7 @@ type openAIWSConn struct { ws openAIWSClientConn handshakeHeaders http.Header + betaFeatures string leaseCh chan struct{} closedCh chan struct{} @@ -498,6 +502,10 @@ 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) isPrewarmed() bool { if c == nil { return false @@ -516,6 +524,7 @@ type openAIWSAccountPool struct { mu sync.Mutex conns map[string]*openAIWSConn pinnedConns map[string]int + changedCh chan struct{} creating int lastCleanupAt time.Time lastAcquire *openAIWSAcquireRequest @@ -525,6 +534,23 @@ type openAIWSAccountPool struct { prewarmFailAt time.Time } +func (ap *openAIWSAccountPool) changeChannelLocked() chan struct{} { + if ap.changedCh == nil { + ap.changedCh = make(chan struct{}) + } + return ap.changedCh +} + +func (ap *openAIWSAccountPool) signalChangedLocked() { + if ap == nil { + return + } + if ap.changedCh != nil { + close(ap.changedCh) + } + ap.changedCh = make(chan struct{}) +} + type OpenAIWSPoolMetricsSnapshot struct { AcquireTotal int64 AcquireReuseTotal int64 @@ -786,7 +812,9 @@ func (p *openAIWSConnPool) acquire(ctx context.Context, req openAIWSAcquireReque return nil, errors.New("ws url is empty") } +retryAcquire: accountID := req.Account.ID + betaFeatures := normalizeOpenAIWSBetaFeatures(req.Headers) effectiveMaxConns := p.effectiveMaxConnsByAccount(req.Account) if effectiveMaxConns <= 0 { return nil, errOpenAIWSConnQueueFull @@ -814,7 +842,7 @@ func (p *openAIWSConnPool) acquire(ctx context.Context, req openAIWSAcquireReque return nil, errOpenAIWSPreferredConnUnavailable } preferredConn, ok := ap.conns[preferredConnID] - if !ok || preferredConn == nil { + if !ok || !preferredConn.matchesBetaFeatures(betaFeatures) { p.recordConnPickDuration(time.Since(pickStartedAt)) ap.mu.Unlock() closeOpenAIWSConns(evicted) @@ -895,7 +923,7 @@ func (p *openAIWSConnPool) acquire(ctx context.Context, req openAIWSAcquireReque } if preferredConnID != "" { - if conn, ok := ap.conns[preferredConnID]; ok && conn.tryAcquire() { + if conn, ok := ap.conns[preferredConnID]; ok && conn.matchesBetaFeatures(betaFeatures) && conn.tryAcquire() { connPick := time.Since(pickStartedAt) p.recordConnPickDuration(connPick) ap.mu.Unlock() @@ -917,7 +945,7 @@ func (p *openAIWSConnPool) acquire(ctx context.Context, req openAIWSAcquireReque } } - best := p.pickLeastBusyConnLocked(ap, "") + best := p.pickLeastBusyConnLocked(ap, "", betaFeatures) if best != nil && best.tryAcquire() { connPick := time.Since(pickStartedAt) p.recordConnPickDuration(connPick) @@ -939,7 +967,7 @@ func (p *openAIWSConnPool) acquire(ctx context.Context, req openAIWSAcquireReque return lease, nil } for _, conn := range ap.conns { - if conn == nil || conn == best { + if conn == nil || conn == best || !conn.matchesBetaFeatures(betaFeatures) { continue } if conn.tryAcquire() { @@ -965,6 +993,37 @@ func (p *openAIWSConnPool) acquire(ctx context.Context, req openAIWSAcquireReque } } + if !req.ForceNewConn && len(ap.conns)+ap.creating >= effectiveMaxConns { + compatible := p.pickLeastBusyConnLocked(ap, "", betaFeatures) + if idle := p.pickOldestIdleConnWithDifferentBetaFeaturesLocked(ap, betaFeatures); idle != nil { + delete(ap.conns, idle.id) + evicted = append(evicted, idle) + p.metrics.scaleDownTotal.Add(1) + } else if compatible == nil { + hasConnection := false + for _, conn := range ap.conns { + if conn != nil { + hasConnection = true + break + } + } + if !hasConnection && ap.creating == 0 { + ap.mu.Unlock() + closeOpenAIWSConns(evicted) + return nil, errOpenAIWSConnClosed + } + changedCh := ap.changeChannelLocked() + ap.mu.Unlock() + closeOpenAIWSConns(evicted) + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-changedCh: + goto retryAcquire + } + } + } + if req.ForceNewConn && len(ap.conns)+ap.creating >= effectiveMaxConns { if idle := p.pickOldestIdleConnLocked(ap); idle != nil { delete(ap.conns, idle.id) @@ -988,6 +1047,7 @@ func (p *openAIWSConnPool) acquire(ctx context.Context, req openAIWSAcquireReque if dialErr != nil { ap.prewarmFails++ ap.prewarmFailAt = time.Now() + ap.signalChangedLocked() ap.mu.Unlock() return nil, dialErr } @@ -1016,7 +1076,7 @@ func (p *openAIWSConnPool) acquire(ctx context.Context, req openAIWSAcquireReque return nil, errOpenAIWSConnQueueFull } - target := p.pickLeastBusyConnLocked(ap, req.PreferredConnID) + target := p.pickLeastBusyConnLocked(ap, req.PreferredConnID, betaFeatures) connPick := time.Since(pickStartedAt) p.recordConnPickDuration(connPick) if target == nil { @@ -1089,6 +1149,22 @@ func (p *openAIWSConnPool) pickOldestIdleConnLocked(ap *openAIWSAccountPool) *op return oldest } +func (p *openAIWSConnPool) pickOldestIdleConnWithDifferentBetaFeaturesLocked(ap *openAIWSAccountPool, betaFeatures 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) { + continue + } + if oldest == nil || conn.lastUsedAt().Before(oldest.lastUsedAt()) { + oldest = conn + } + } + return oldest +} + func (p *openAIWSConnPool) getOrCreateAccountPool(accountID int64) *openAIWSAccountPool { if p == nil || accountID <= 0 { return nil @@ -1101,6 +1177,7 @@ func (p *openAIWSConnPool) getOrCreateAccountPool(accountID int64) *openAIWSAcco ap := &openAIWSAccountPool{ conns: make(map[string]*openAIWSConn), pinnedConns: make(map[string]int), + changedCh: make(chan struct{}), } actual, _ := p.accounts.LoadOrStore(accountID, ap) if typed, ok := actual.(*openAIWSAccountPool); ok && typed != nil { @@ -1126,6 +1203,16 @@ func (p *openAIWSConnPool) getAccountPool(accountID int64) (*openAIWSAccountPool return ap, typed && ap != nil } +func (p *openAIWSConnPool) notifyAccountPoolChanged(accountID int64) { + ap, ok := p.getAccountPool(accountID) + if !ok || ap == nil { + return + } + ap.mu.Lock() + ap.signalChangedLocked() + ap.mu.Unlock() +} + func (p *openAIWSConnPool) isConnPinnedLocked(ap *openAIWSAccountPool, connID string) bool { if ap == nil || connID == "" || len(ap.pinnedConns) == 0 { return false @@ -1212,17 +1299,20 @@ func (p *openAIWSConnPool) cleanupAccountLocked(ap *openAIWSAccountPool, now tim p.metrics.scaleDownTotal.Add(int64(redundant)) } } + if len(evicted) > 0 { + ap.signalChangedLocked() + } return evicted } -func (p *openAIWSConnPool) pickLeastBusyConnLocked(ap *openAIWSAccountPool, preferredConnID string) *openAIWSConn { +func (p *openAIWSConnPool) pickLeastBusyConnLocked(ap *openAIWSAccountPool, preferredConnID, betaFeatures string) *openAIWSConn { if ap == nil || len(ap.conns) == 0 { return nil } preferredConnID = stringsTrim(preferredConnID) if preferredConnID != "" { - if conn, ok := ap.conns[preferredConnID]; ok { + if conn, ok := ap.conns[preferredConnID]; ok && conn.matchesBetaFeatures(betaFeatures) { return conn } } @@ -1230,7 +1320,7 @@ func (p *openAIWSConnPool) pickLeastBusyConnLocked(ap *openAIWSAccountPool, pref var bestWaiters int32 var bestLastUsed time.Time for _, conn := range ap.conns { - if conn == nil { + if conn == nil || !conn.matchesBetaFeatures(betaFeatures) { continue } waiters := conn.waiters.Load() @@ -1395,10 +1485,12 @@ func (p *openAIWSConnPool) prewarmConns(accountID int64, req openAIWSAcquireRequ if err != nil { ap.prewarmFails++ ap.prewarmFailAt = time.Now() + ap.signalChangedLocked() ap.mu.Unlock() continue } if len(ap.conns) >= p.effectiveMaxConnsByAccount(req.Account) { + ap.signalChangedLocked() ap.mu.Unlock() conn.close() continue @@ -1406,6 +1498,7 @@ func (p *openAIWSConnPool) prewarmConns(accountID int64, req openAIWSAcquireRequ ap.conns[conn.id] = conn ap.prewarmFails = 0 ap.prewarmFailAt = time.Time{} + ap.signalChangedLocked() ap.mu.Unlock() } } @@ -1424,6 +1517,7 @@ func (p *openAIWSConnPool) evictConn(accountID int64, connID string) { if len(ap.pinnedConns) > 0 { delete(ap.pinnedConns, connID) } + ap.signalChangedLocked() } ap.mu.Unlock() } @@ -1476,9 +1570,11 @@ func (p *openAIWSConnPool) UnpinConn(accountID int64, connID string) { count := ap.pinnedConns[connID] if count <= 1 { delete(ap.pinnedConns, connID) + ap.signalChangedLocked() return } ap.pinnedConns[connID] = count - 1 + ap.signalChangedLocked() } func (p *openAIWSConnPool) dialConn(ctx context.Context, req openAIWSAcquireRequest) (*openAIWSConn, error) { @@ -1501,7 +1597,9 @@ func (p *openAIWSConnPool) dialConn(ctx context.Context, req openAIWSAcquireRequ } } id := p.nextConnID(req.Account.ID) - return newOpenAIWSConn(id, req.Account.ID, conn, handshakeHeaders), nil + pooledConn := newOpenAIWSConn(id, req.Account.ID, conn, handshakeHeaders) + pooledConn.betaFeatures = normalizeOpenAIWSBetaFeatures(req.Headers) + return pooledConn, nil } func (p *openAIWSConnPool) nextConnID(accountID int64) string { @@ -1679,6 +1777,31 @@ func cloneOpenAIWSAcquireRequestPtr(req *openAIWSAcquireRequest) *openAIWSAcquir return &copied } +func normalizeOpenAIWSBetaFeatures(headers http.Header) string { + features := make(map[string]struct{}) + for name, values := range headers { + if !strings.EqualFold(strings.TrimSpace(name), "x-codex-beta-features") { + continue + } + for _, value := range values { + for _, feature := range strings.Split(value, ",") { + if feature = strings.TrimSpace(feature); feature != "" { + features[feature] = struct{}{} + } + } + } + } + if len(features) == 0 { + return "" + } + normalized := make([]string, 0, len(features)) + for feature := range features { + normalized = append(normalized, feature) + } + sort.Strings(normalized) + return strings.Join(normalized, ",") +} + func cloneHeader(src http.Header) http.Header { if src == nil { return nil diff --git a/backend/internal/service/openai_ws_pool_test.go b/backend/internal/service/openai_ws_pool_test.go index b2683ee041..ae9b94ce4a 100644 --- a/backend/internal/service/openai_ws_pool_test.go +++ b/backend/internal/service/openai_ws_pool_test.go @@ -342,6 +342,171 @@ func TestOpenAIWSConnPool_ForceNewConnSkipsReuse(t *testing.T) { require.Equal(t, 2, dialer.DialCount(), "ForceNewConn=true 时应跳过空闲连接复用并新建连接") } +func TestOpenAIWSConnPool_AcquireReusesOnlyMatchingBetaFeatures(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: 128, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + baseReq := openAIWSAcquireRequest{ + Account: account, + WSURL: "wss://example.com/v1/responses", + } + + plainLease, err := pool.Acquire(context.Background(), baseReq) + require.NoError(t, err) + plainConnID := plainLease.ConnID() + plainLease.Release() + + betaReq := baseReq + betaReq.Headers = http.Header{"X-Codex-Beta-Features": {" remote_compaction_v2 ", " responses_websockets_v2 "}} + betaLease, err := pool.Acquire(context.Background(), betaReq) + require.NoError(t, err) + require.False(t, betaLease.Reused()) + require.NotEqual(t, plainConnID, betaLease.ConnID()) + betaConnID := betaLease.ConnID() + betaLease.Release() + + reorderedReq := baseReq + reorderedReq.Headers = http.Header{"X-Codex-Beta-Features": {"responses_websockets_v2,remote_compaction_v2"}} + reorderedLease, err := pool.Acquire(context.Background(), reorderedReq) + require.NoError(t, err) + require.True(t, reorderedLease.Reused()) + require.Equal(t, betaConnID, reorderedLease.ConnID()) + reorderedLease.Release() + + _, err = pool.Acquire(context.Background(), openAIWSAcquireRequest{ + Account: account, + WSURL: baseReq.WSURL, + Headers: betaReq.Headers, + PreferredConnID: plainConnID, + ForcePreferredConn: true, + }) + require.ErrorIs(t, err, errOpenAIWSPreferredConnUnavailable) + + plainLease, err = pool.Acquire(context.Background(), baseReq) + require.NoError(t, err) + require.True(t, plainLease.Reused()) + require.Equal(t, plainConnID, plainLease.ConnID()) + plainLease.Release() + + require.Equal(t, 2, dialer.DialCount()) +} + +func TestOpenAIWSConnPool_AcquireReplacesIdleConnWithDifferentBetaFeatures(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: 129, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + plainLease, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{ + Account: account, + WSURL: "wss://example.com/v1/responses", + }) + require.NoError(t, err) + plainConnID := plainLease.ConnID() + plainLease.Release() + + betaLease, err := pool.Acquire(context.Background(), openAIWSAcquireRequest{ + Account: account, + WSURL: "wss://example.com/v1/responses", + Headers: http.Header{"X-Codex-Beta-Features": {"remote_compaction_v2"}}, + }) + require.NoError(t, err) + require.False(t, betaLease.Reused()) + require.NotEqual(t, plainConnID, betaLease.ConnID()) + betaLease.Release() + + require.Equal(t, 2, dialer.DialCount()) +} + +func TestOpenAIWSConnPool_AcquireWaitsForBusyIncompatibleConnection(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: 130, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + baseReq := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"} + + plainLease, err := pool.Acquire(context.Background(), baseReq) + require.NoError(t, err) + plainConnID := plainLease.ConnID() + + type acquireResult struct { + lease *openAIWSConnLease + err error + } + resultCh := make(chan acquireResult, 1) + var done atomic.Bool + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + go func() { + betaReq := baseReq + betaReq.Headers = http.Header{"X-Codex-Beta-Features": {"remote_compaction_v2"}} + lease, acquireErr := pool.Acquire(ctx, betaReq) + resultCh <- acquireResult{lease: lease, err: acquireErr} + done.Store(true) + }() + + require.Never(t, done.Load, 50*time.Millisecond, 5*time.Millisecond) + plainLease.Release() + + result := <-resultCh + require.NoError(t, result.err) + require.NotNil(t, result.lease) + require.False(t, result.lease.Reused()) + require.NotEqual(t, plainConnID, result.lease.ConnID()) + result.lease.Release() + require.Equal(t, 2, dialer.DialCount()) +} + +func TestOpenAIWSConnPool_AcquireReplacesIncompatibleIdleWhenMatchingBusy(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: 131, Platform: PlatformOpenAI, Type: AccountTypeAPIKey} + baseReq := openAIWSAcquireRequest{Account: account, WSURL: "wss://example.com/v1/responses"} + + plainLease, err := pool.Acquire(context.Background(), baseReq) + require.NoError(t, err) + plainConnID := plainLease.ConnID() + plainLease.Release() + + betaReq := baseReq + betaReq.Headers = http.Header{"X-Codex-Beta-Features": {"remote_compaction_v2"}} + busyBetaLease, err := pool.Acquire(context.Background(), betaReq) + require.NoError(t, err) + + secondBetaLease, err := pool.Acquire(context.Background(), betaReq) + require.NoError(t, err) + require.False(t, secondBetaLease.Reused()) + require.NotEqual(t, plainConnID, secondBetaLease.ConnID()) + require.NotEqual(t, busyBetaLease.ConnID(), secondBetaLease.ConnID()) + + secondBetaLease.Release() + busyBetaLease.Release() + require.Equal(t, 3, dialer.DialCount()) +} + func TestOpenAIWSConnPool_AcquireForcePreferredConnUnavailable(t *testing.T) { cfg := &config.Config{} cfg.Gateway.OpenAIWS.MaxConnsPerAccount = 2 diff --git a/backend/migrations/174_group_web_search_price_per_call.sql b/backend/migrations/174_group_web_search_price_per_call.sql new file mode 100644 index 0000000000..da9c90fed8 --- /dev/null +++ b/backend/migrations/174_group_web_search_price_per_call.sql @@ -0,0 +1,3 @@ +-- Codex alpha/search 网页搜索按次计费：分组级单次价格覆盖。 +-- NULL 表示使用内置默认价 0.01 USD/次（OpenAI 官方 web search 定价 $10/1000 次）。 +ALTER TABLE groups ADD COLUMN IF NOT EXISTS web_search_price_per_call DECIMAL(20,8); diff --git a/frontend/src/components/account/AccountUsageCell.vue b/frontend/src/components/account/AccountUsageCell.vue index 81d97efb9c..5d7af74fc6 100644 --- a/frontend/src/components/account/AccountUsageCell.vue +++ b/frontend/src/components/account/AccountUsageCell.vue @@ -386,6 +386,7 @@ :label="t('admin.accounts.usageWindow.grokRequests')" :utilization="grokRequestQuotaBar.utilization" :resets-at="grokRequestQuotaBar.resetsAt" + :remaining-capacity="true" color="indigo" />
@@ -1036,9 +1038,9 @@ interface GrokQuotaBarInfo { const makeGrokQuotaBar = (quota?: { limit?: number | null; remaining?: number | null; reset_at?: string | null } | null): GrokQuotaBarInfo | null => { if (!quota || quota.limit == null || quota.remaining == null || quota.limit <= 0) return null - const used = Math.max(0, quota.limit - quota.remaining) + const remaining = Math.min(quota.limit, Math.max(0, quota.remaining)) return { - utilization: (used / quota.limit) * 100, + utilization: (remaining / quota.limit) * 100, resetsAt: quota.reset_at || null } } diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index d7e72f9e8a..985ceaa731 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -352,7 +352,7 @@
- +
@@ -381,10 +381,34 @@ {{ t('admin.accounts.types.grokOauth') }}
+ +
-

- {{ t('admin.accounts.oauth.grok.oauthOnlyHint') }} -

@@ -1087,10 +1111,12 @@ ? 'https://api.openai.com' : form.platform === 'gemini' ? 'https://generativelanguage.googleapis.com' - : 'https://api.anthropic.com' + : form.platform === 'grok' + ? 'https://api.x.ai/v1' + : 'https://api.anthropic.com' " /> -

{{ baseUrlHint }}

+

{{ baseUrlHint }}

@@ -1104,10 +1130,12 @@ ? 'sk-proj-...' : form.platform === 'gemini' ? 'AIza...' - : 'sk-ant-...' + : form.platform === 'grok' + ? 'xai-...' + : 'sk-ant-...' " /> -

{{ apiKeyHint }}

+

{{ apiKeyHint }}

@@ -3515,14 +3543,14 @@ const oauthStepTitle = computed(() => { const baseUrlHint = computed(() => { if (form.platform === 'openai') return t('admin.accounts.openai.baseUrlHint') if (form.platform === 'gemini') return t('admin.accounts.gemini.baseUrlHint') - if (form.platform === 'grok') return t('admin.accounts.grok.baseUrlHint') + if (form.platform === 'grok') return '' return t('admin.accounts.baseUrlHint') }) const apiKeyHint = computed(() => { if (form.platform === 'openai') return t('admin.accounts.openai.apiKeyHint') if (form.platform === 'gemini') return t('admin.accounts.gemini.apiKeyHint') - if (form.platform === 'grok') return t('admin.accounts.grok.apiKeyHint') + if (form.platform === 'grok') return '' return t('admin.accounts.apiKeyHint') }) @@ -4926,7 +4954,9 @@ const handleSubmit = async () => { ? 'https://api.openai.com' : form.platform === 'gemini' ? 'https://generativelanguage.googleapis.com' - : 'https://api.anthropic.com' + : form.platform === 'grok' + ? 'https://api.x.ai/v1' + : 'https://api.anthropic.com' // Build credentials with optional model mapping const credentials: Record = { diff --git a/frontend/src/components/account/EditAccountModal.vue b/frontend/src/components/account/EditAccountModal.vue index f282fd75b7..0433997156 100644 --- a/frontend/src/components/account/EditAccountModal.vue +++ b/frontend/src/components/account/EditAccountModal.vue @@ -41,10 +41,12 @@ ? 'https://generativelanguage.googleapis.com' : account.platform === 'antigravity' ? 'https://cloudcode-pa.googleapis.com' - : 'https://api.anthropic.com' + : account.platform === 'grok' + ? 'https://api.x.ai/v1' + : 'https://api.anthropic.com' " /> -

{{ baseUrlHint }}

+

{{ baseUrlHint }}

@@ -63,7 +65,9 @@ ? 'AIza...' : account.platform === 'antigravity' ? 'sk-...' - : 'sk-ant-...' + : account.platform === 'grok' + ? 'xai-...' + : 'sk-ant-...' " />

{{ t('admin.accounts.leaveEmptyToKeep') }}

@@ -2626,6 +2630,7 @@ const baseUrlHint = computed(() => { if (!props.account) return t('admin.accounts.baseUrlHint') if (props.account.platform === 'openai') return t('admin.accounts.openai.baseUrlHint') if (props.account.platform === 'gemini') return t('admin.accounts.gemini.baseUrlHint') + if (props.account.platform === 'grok') return '' return t('admin.accounts.baseUrlHint') }) @@ -3076,6 +3081,7 @@ const tempUnschedPresets = computed(() => [ const defaultBaseUrl = computed(() => { if (props.account?.platform === 'openai') return 'https://api.openai.com' if (props.account?.platform === 'gemini') return 'https://generativelanguage.googleapis.com' + if (props.account?.platform === 'grok') return 'https://api.x.ai/v1' return 'https://api.anthropic.com' }) @@ -3370,7 +3376,9 @@ const syncFormFromAccount = (newAccount: Account | null) => { ? 'https://api.openai.com' : newAccount.platform === 'gemini' ? 'https://generativelanguage.googleapis.com' - : 'https://api.anthropic.com' + : newAccount.platform === 'grok' + ? 'https://api.x.ai/v1' + : 'https://api.anthropic.com' editBaseUrl.value = (credentials.base_url as string) || platformDefaultUrl // Load model mappings and detect mode @@ -3446,7 +3454,9 @@ const syncFormFromAccount = (newAccount: Account | null) => { ? 'https://api.openai.com' : newAccount.platform === 'gemini' ? 'https://generativelanguage.googleapis.com' - : 'https://api.anthropic.com' + : newAccount.platform === 'grok' + ? 'https://api.x.ai/v1' + : 'https://api.anthropic.com' editBaseUrl.value = platformDefaultUrl // Load model mappings for OpenAI/Grok OAuth accounts diff --git a/frontend/src/components/account/UsageProgressBar.vue b/frontend/src/components/account/UsageProgressBar.vue index 6a69357318..2f8b9a88a7 100644 --- a/frontend/src/components/account/UsageProgressBar.vue +++ b/frontend/src/components/account/UsageProgressBar.vue @@ -69,6 +69,7 @@ const props = defineProps<{ color: 'indigo' | 'emerald' | 'purple' | 'amber' windowStats?: WindowStats | null showNowWhenIdle?: boolean + remainingCapacity?: boolean }>() const { t } = useI18n() @@ -109,6 +110,14 @@ const labelClass = computed(() => { // Progress bar color based on utilization const barClass = computed(() => { + if (props.remainingCapacity) { + if (props.utilization <= 20) { + return 'bg-red-500' + } else if (props.utilization <= 50) { + return 'bg-amber-500' + } + return 'bg-green-500' + } if (props.utilization >= 100) { return 'bg-red-500' } else if (props.utilization >= 80) { @@ -120,6 +129,14 @@ const barClass = computed(() => { // Text color based on utilization const textClass = computed(() => { + if (props.remainingCapacity) { + if (props.utilization <= 20) { + return 'text-red-600 dark:text-red-400' + } else if (props.utilization <= 50) { + return 'text-amber-600 dark:text-amber-400' + } + return 'text-gray-600 dark:text-gray-400' + } if (props.utilization >= 100) { return 'text-red-600 dark:text-red-400' } else if (props.utilization >= 80) { @@ -131,12 +148,16 @@ const textClass = computed(() => { // Bar width (capped at 100%) const barWidth = computed(() => { - return `${Math.min(props.utilization, 100)}%` + return `${Math.min(Math.max(props.utilization, 0), 100)}%` }) // Display percentage (cap at 999% for readability) const displayPercent = computed(() => { - const percent = Math.round(props.utilization) + const percent = Math.round( + props.remainingCapacity + ? Math.min(Math.max(props.utilization, 0), 100) + : props.utilization + ) return percent > 999 ? '>999%' : `${percent}%` }) diff --git a/frontend/src/components/account/__tests__/AccountStatusIndicator.spec.ts b/frontend/src/components/account/__tests__/AccountStatusIndicator.spec.ts index f758e6b0f6..545c2abac4 100644 --- a/frontend/src/components/account/__tests__/AccountStatusIndicator.spec.ts +++ b/frontend/src/components/account/__tests__/AccountStatusIndicator.spec.ts @@ -13,6 +13,14 @@ vi.mock('vue-i18n', async () => { } }) +vi.mock('@/utils/format', async () => { + const actual = await vi.importActual('@/utils/format') + return { + ...actual, + formatCountdown: () => '1h' + } +}) + function makeAccount(overrides: Partial): Account { return { id: 1, @@ -43,6 +51,31 @@ function makeAccount(overrides: Partial): Account { } describe('AccountStatusIndicator', () => { + it('Grok 账号额度限流时显示自动恢复时间而非临时不可调度', () => { + const wrapper = mount(AccountStatusIndicator, { + props: { + account: makeAccount({ + id: 5, + name: 'grok-free-1', + platform: 'grok', + rate_limited_at: '2026-07-11T12:00:00Z', + rate_limit_reset_at: '2099-07-11T13:00:00Z', + temp_unschedulable_until: '2099-07-11T12:30:00Z', + temp_unschedulable_reason: 'legacy grok rate limited' + }) + }, + global: { + stubs: { + Icon: true + } + } + }) + + expect(wrapper.find('.badge-warning').text()).toBe('admin.accounts.status.rateLimited') + expect(wrapper.text()).toContain('admin.accounts.status.rateLimitedAutoResume') + expect(wrapper.text()).not.toContain('admin.accounts.status.tempUnschedulable') + }) + it('模型限流 + overages 启用 + 无 AICredits key → 显示 ⚡ (credits_active)', () => { const wrapper = mount(AccountStatusIndicator, { props: { diff --git a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts index 2df3cc07e2..2abf6513ca 100644 --- a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts +++ b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts @@ -566,7 +566,7 @@ describe('AccountUsageCell', () => { expect(badges.some(node => node.attributes('title') === 'usage.userBilled')).toBe(true) }) - it('Grok OAuth 会展示本地 user billed 用量并保留超限百分比', async () => { + it('Grok OAuth 会展示本地 user billed 用量并把耗尽配额显示为 0% 剩余', async () => { getUsage.mockResolvedValue({ grok_local_usage: { requests: 4, @@ -611,13 +611,55 @@ describe('AccountUsageCell', () => { expect(wrapper.text()).toContain('1.2K') expect(wrapper.text()).toContain('A $0.12') expect(wrapper.text()).toContain('U $0.34') - expect(wrapper.text()).toContain('admin.accounts.usageWindow.grokRequests|120|2026-07-09T16:00:00Z') + expect(wrapper.text()).toContain('admin.accounts.usageWindow.grokRequests|0|2026-07-09T16:00:00Z') const badges = wrapper.findAll('span[title]') expect(badges.some(node => node.attributes('title') === 'usage.accountBilled')).toBe(true) expect(badges.some(node => node.attributes('title') === 'usage.userBilled')).toBe(true) }) + it('Grok OAuth 配额条按剩余容量显示 100% 满格和 25% 低量', async () => { + getUsage.mockResolvedValue({ + grok_request_quota: { + limit: 100, + remaining: 100, + reset_at: '2026-07-09T16:00:00Z' + }, + grok_token_quota: { + limit: 1000, + remaining: 250, + reset_at: '2026-07-09T16:00:00Z' + }, + grok_quota_snapshot_state: 'observed' + }) + + const wrapper = mount(AccountUsageCell, { + props: { + account: makeAccount({ + id: 4073, + platform: 'grok', + type: 'oauth', + extra: {} + }) + }, + global: { + stubs: { + UsageProgressBar: { + props: ['label', 'utilization', 'resetsAt', 'color', 'remainingCapacity'], + template: '
{{ label }}|{{ utilization }}|{{ remainingCapacity }}
' + }, + AccountQuotaInfo: true, + GrokQuotaProbeCell: true + } + } + }) + + await flushPromises() + + expect(wrapper.text()).toContain('admin.accounts.usageWindow.grokRequests|100|true') + expect(wrapper.text()).toContain('admin.accounts.usageWindow.grokTokens|25|true') + }) + it('Key 账号在 today stats loading 时显示骨架屏', async () => { const wrapper = mount(AccountUsageCell, { props: { diff --git a/frontend/src/components/account/__tests__/CreateAccountModal.grok.spec.ts b/frontend/src/components/account/__tests__/CreateAccountModal.grok.spec.ts new file mode 100644 index 0000000000..0df3663f43 --- /dev/null +++ b/frontend/src/components/account/__tests__/CreateAccountModal.grok.spec.ts @@ -0,0 +1,19 @@ +import { readFileSync } from 'node:fs' +import { resolve } from 'node:path' +import { describe, expect, it } from 'vitest' + +const source = readFileSync( + resolve(process.cwd(), 'src/components/account/CreateAccountModal.vue'), + 'utf8' +) + +describe('CreateAccountModal Grok account types', () => { + it('offers API-key setup alongside OAuth with the official xAI default', () => { + expect(source).toContain('data-testid="grok-account-type-api-key"') + expect(source).toContain("@click=\"accountCategory = 'apikey'\"") + expect(source).toContain("newPlatform === 'grok'") + expect(source).toContain("? 'https://api.x.ai/v1'") + expect(source).toContain("form.platform === 'grok'") + expect(source).toContain("? 'xai-...'") + }) +}) diff --git a/frontend/src/components/account/__tests__/EditAccountModal.spec.ts b/frontend/src/components/account/__tests__/EditAccountModal.spec.ts index 394b7e9228..7ae9afc2e5 100644 --- a/frontend/src/components/account/__tests__/EditAccountModal.spec.ts +++ b/frontend/src/components/account/__tests__/EditAccountModal.spec.ts @@ -267,6 +267,18 @@ function buildGrokOAuthAccount() { } as any } +function buildGrokAPIKeyAccount() { + return { + ...buildAccount(), + id: 6, + name: 'Grok API Key', + platform: 'grok', + credentials: {}, + credentials_status: { has_api_key: true }, + concurrency: 2 + } as any +} + function buildOpenAISetupTokenAccount() { return { ...buildAccount(), @@ -433,6 +445,24 @@ describe('EditAccountModal', () => { }) }) + it('uses the official xAI base URL when a Grok API-key account omits base_url', async () => { + const account = buildGrokAPIKeyAccount() + updateAccountMock.mockReset() + checkMixedChannelRiskMock.mockReset() + checkMixedChannelRiskMock.mockResolvedValue({ has_risk: false }) + updateAccountMock.mockResolvedValue(account) + + const wrapper = mountModal(account) + + expect((wrapper.get('input[placeholder="https://api.x.ai/v1"]').element as HTMLInputElement).value) + .toBe('https://api.x.ai/v1') + + await wrapper.get('form#edit-account-form').trigger('submit.prevent') + + expect(updateAccountMock).toHaveBeenCalledTimes(1) + expect(updateAccountMock.mock.calls[0]?.[1]?.credentials?.base_url).toBe('https://api.x.ai/v1') + }) + it('only submits model mapping credentials when saving an OpenAI spark shadow account', async () => { authIsSimpleMode.value = false const account = buildOpenAISparkShadowAccount() diff --git a/frontend/src/components/account/__tests__/UsageProgressBar.spec.ts b/frontend/src/components/account/__tests__/UsageProgressBar.spec.ts index 6fa6575f54..af5fc5d66d 100644 --- a/frontend/src/components/account/__tests__/UsageProgressBar.spec.ts +++ b/frontend/src/components/account/__tests__/UsageProgressBar.spec.ts @@ -96,4 +96,54 @@ describe('UsageProgressBar', () => { expect(wrapper.text()).toContain('usage.resetNow') expect(wrapper.text()).not.toContain('usage.resetPending') }) + + it('剩余容量模式在 100% 时显示满格绿色', () => { + const wrapper = mount(UsageProgressBar, { + props: { + label: 'Req', + utilization: 100, + remainingCapacity: true, + color: 'indigo' + } + }) + + expect(wrapper.text()).toContain('100%') + expect(wrapper.get('.h-1\\.5 > div').attributes('style')).toContain('width: 100%') + expect(wrapper.get('.h-1\\.5 > div').classes()).toContain('bg-green-500') + }) + + it('剩余容量模式在低量和耗尽时缩短并变红', async () => { + const wrapper = mount(UsageProgressBar, { + props: { + label: 'Req', + utilization: 15, + remainingCapacity: true, + color: 'indigo' + } + }) + + expect(wrapper.text()).toContain('15%') + expect(wrapper.get('.h-1\\.5 > div').attributes('style')).toContain('width: 15%') + expect(wrapper.get('.h-1\\.5 > div').classes()).toContain('bg-red-500') + + await wrapper.setProps({ utilization: 0 }) + + expect(wrapper.text()).toContain('0%') + expect(wrapper.get('.h-1\\.5 > div').attributes('style')).toContain('width: 0%') + expect(wrapper.get('.h-1\\.5 > div').classes()).toContain('bg-red-500') + }) + + it('默认利用率模式仍把超限显示为满格红色', () => { + const wrapper = mount(UsageProgressBar, { + props: { + label: '5h', + utilization: 120, + color: 'indigo' + } + }) + + expect(wrapper.text()).toContain('120%') + expect(wrapper.get('.h-1\\.5 > div').attributes('style')).toContain('width: 100%') + expect(wrapper.get('.h-1\\.5 > div').classes()).toContain('bg-red-500') + }) }) diff --git a/frontend/src/components/keys/UseKeyModal.vue b/frontend/src/components/keys/UseKeyModal.vue index d3bca5188a..e6bce556da 100644 --- a/frontend/src/components/keys/UseKeyModal.vue +++ b/frontend/src/components/keys/UseKeyModal.vue @@ -181,6 +181,8 @@ const defaultClientTab = computed(() => { switch (props.platform) { case 'openai': return 'codex' + case 'grok': + return 'grok' case 'gemini': return 'gemini' case 'antigravity': @@ -288,6 +290,11 @@ const clientTabs = computed((): TabConfig[] => { { id: 'gemini', label: t('keys.useKeyModal.cliTabs.geminiCli'), icon: SparkleIcon }, { id: 'opencode', label: t('keys.useKeyModal.cliTabs.opencode'), icon: TerminalIcon } ] + case 'grok': + return [ + { id: 'grok', label: t('keys.useKeyModal.cliTabs.grokCli'), icon: TerminalIcon }, + { id: 'opencode', label: t('keys.useKeyModal.cliTabs.opencode'), icon: TerminalIcon } + ] default: return [ { id: 'claude', label: t('keys.useKeyModal.cliTabs.claudeCode'), icon: TerminalIcon }, @@ -313,7 +320,7 @@ const showShellTabs = computed(() => activeClientTab.value !== 'opencode') const currentTabs = computed(() => { if (!showShellTabs.value) return [] - if (activeClientTab.value === 'codex' || activeClientTab.value === 'codex-ws') { + if (activeClientTab.value === 'codex' || activeClientTab.value === 'codex-ws' || activeClientTab.value === 'grok') { return openaiTabs } return shellTabs @@ -330,6 +337,8 @@ const platformDescription = computed(() => { return t('keys.useKeyModal.gemini.description') case 'antigravity': return t('keys.useKeyModal.antigravity.description') + case 'grok': + return t('keys.useKeyModal.grok.description') default: return t('keys.useKeyModal.description') } @@ -350,6 +359,10 @@ const platformNote = computed(() => { return activeClientTab.value === 'claude' ? t('keys.useKeyModal.antigravity.claudeNote') : t('keys.useKeyModal.antigravity.geminiNote') + case 'grok': + return activeTab.value === 'windows' + ? t('keys.useKeyModal.grok.noteWindows') + : t('keys.useKeyModal.grok.note') default: return t('keys.useKeyModal.note') } @@ -407,6 +420,8 @@ const currentFiles = computed((): FileConfig[] => { generateOpenCodeConfig('antigravity-claude', antigravityBase, apiKey, 'opencode.json (Claude)'), generateOpenCodeConfig('antigravity-gemini', antigravityGeminiBase, apiKey, 'opencode.json (Gemini)') ] + case 'grok': + return [generateOpenCodeConfig('grok', apiBase, apiKey)] default: return [generateOpenCodeConfig('openai', apiBase, apiKey)] } @@ -428,6 +443,8 @@ const currentFiles = computed((): FileConfig[] => { return [generateGeminiCliContent(`${baseUrl}/antigravity`, apiKey)] } return generateAnthropicFiles(`${baseUrl}/antigravity`, apiKey) + case 'grok': + return generateGrokFiles(apiBase, apiKey) default: return generateAnthropicFiles(baseUrl, apiKey) } @@ -568,6 +585,30 @@ goals = true` ] } +function generateGrokFiles(baseUrl: string, apiKey: string): FileConfig[] { + const isWindows = activeTab.value === 'windows' + const configDir = isWindows ? '%userprofile%\\.grok' : '~/.grok' + const configContent = `[models] +default = "sub2api-grok" +web_search = "sub2api-grok" + +[model."sub2api-grok"] +model = "grok-4.5" +base_url = "${baseUrl}" +name = "Grok 4.5 via Sub2API" +description = "Grok 4.5 through a Sub2API Grok group" +api_key = "${apiKey}" +api_backend = "responses" +context_window = 1000000 +supports_backend_search = true` + + return [{ + path: `${configDir}/config.toml`, + content: configContent, + hint: t('keys.useKeyModal.grok.configTomlHint') + }] +} + function generateOpenAIWsFiles(baseUrl: string, apiKey: string): FileConfig[] { const isWindows = activeTab.value === 'windows' const configDir = isWindows ? '%userprofile%\\.codex' : '~/.codex' @@ -1072,6 +1113,24 @@ function generateOpenCodeConfig(platform: string, baseUrl: string, apiKey: strin } } } + const grokModels = { + 'grok-4.5': { + name: 'Grok 4.5', + limit: { context: 1000000, output: 128000 } + }, + 'grok-4.3': { + name: 'Grok 4.3', + limit: { context: 1000000, output: 128000 } + }, + 'grok-build-0.1': { + name: 'Grok Build 0.1', + limit: { context: 256000, output: 128000 } + }, + 'grok-composer-2.5-fast': { + name: 'Grok Composer 2.5 Fast', + limit: { context: 500000, output: 128000 } + } + } if (platform === 'gemini') { provider[platform].npm = '@ai-sdk/google' @@ -1088,6 +1147,10 @@ function generateOpenCodeConfig(platform: string, baseUrl: string, apiKey: strin provider[platform].models = antigravityGeminiModels } else if (platform === 'openai') { provider[platform].models = openaiModels + } else if (platform === 'grok') { + provider[platform].npm = '@ai-sdk/openai' + provider[platform].name = 'Grok via Sub2API' + provider[platform].models = grokModels } const agent = diff --git a/frontend/src/components/keys/__tests__/UseKeyModal.spec.ts b/frontend/src/components/keys/__tests__/UseKeyModal.spec.ts index f8395dc75f..24581d9f6b 100644 --- a/frontend/src/components/keys/__tests__/UseKeyModal.spec.ts +++ b/frontend/src/components/keys/__tests__/UseKeyModal.spec.ts @@ -17,6 +17,67 @@ vi.mock('@/composables/useClipboard', () => ({ import UseKeyModal from '../UseKeyModal.vue' describe('UseKeyModal', () => { + it('renders Grok Build and OpenCode setup for Grok groups', async () => { + const wrapper = mount(UseKeyModal, { + props: { + show: true, + apiKey: 'sk-grok-test', + baseUrl: 'https://example.com/v1', + platform: 'grok' + }, + global: { + stubs: { + BaseDialog: { + template: '
' + }, + Icon: { + template: '' + } + } + } + }) + + const grokTab = wrapper.findAll('button').find((button) => + button.text().includes('keys.useKeyModal.cliTabs.grokCli') + ) + expect(grokTab).toBeDefined() + + const grokConfig = wrapper.findAll('pre code') + .map((code) => code.text()) + .find((content) => content.includes('[model."sub2api-grok"]')) + expect(grokConfig).toBeDefined() + expect(grokConfig).toContain('model = "grok-4.5"') + expect(grokConfig).toContain('base_url = "https://example.com/v1"') + expect(grokConfig).toContain('api_key = "sk-grok-test"') + expect(grokConfig).toContain('api_backend = "responses"') + + const windowsTab = wrapper.findAll('button').find( + (button) => button.text().trim() === 'Windows' + ) + expect(windowsTab).toBeDefined() + await windowsTab!.trigger('click') + await nextTick() + expect(wrapper.text()).toContain('%userprofile%\\.grok/config.toml') + + const opencodeTab = wrapper.findAll('button').find((button) => + button.text().includes('keys.useKeyModal.cliTabs.opencode') + ) + expect(opencodeTab).toBeDefined() + await opencodeTab!.trigger('click') + await nextTick() + + const parsed = JSON.parse(wrapper.find('pre code').text()) + expect(parsed.provider.grok.npm).toBe('@ai-sdk/openai') + expect(parsed.provider.grok.options).toEqual({ + baseURL: 'https://example.com/v1', + apiKey: 'sk-grok-test' + }) + expect(parsed.provider.grok.models['grok-4.5']).toBeDefined() + expect(parsed.provider.grok.models['grok-build-0.1']).toBeDefined() + expect(parsed.provider.grok.models['grok-composer-2.5-fast']).toBeDefined() + expect(parsed.provider.grok.models['gpt-5.6']).toBeUndefined() + }) + it('renders GPT-5.5 and goals feature in OpenAI Codex config', () => { const wrapper = mount(UseKeyModal, { props: { diff --git a/frontend/src/composables/__tests__/useGrokOAuth.spec.ts b/frontend/src/composables/__tests__/useGrokOAuth.spec.ts index 0d7c93d9c6..958a71e67c 100644 --- a/frontend/src/composables/__tests__/useGrokOAuth.spec.ts +++ b/frontend/src/composables/__tests__/useGrokOAuth.spec.ts @@ -53,3 +53,20 @@ describe('useGrokOAuth.exchangeAuthCode', () => { ) }) }) + +describe('useGrokOAuth.buildCredentials', () => { + it('persists the Grok CLI subscription proxy for OAuth inference', () => { + const oauth = useGrokOAuth() + + const credentials = oauth.buildCredentials({ + access_token: 'access-token', + token_type: 'Bearer', + expires_at: 1_900_000_000, + client_id: 'client-id', + scope: 'openid grok-cli:access', + email: 'grok@example.com' + }) + + expect(credentials.base_url).toBe('https://cli-chat-proxy.grok.com/v1') + }) +}) diff --git a/frontend/src/composables/useGrokOAuth.ts b/frontend/src/composables/useGrokOAuth.ts index ad746bfecf..d02934a818 100644 --- a/frontend/src/composables/useGrokOAuth.ts +++ b/frontend/src/composables/useGrokOAuth.ts @@ -122,7 +122,8 @@ export function useGrokOAuth() { scope: tokenInfo.scope, email: tokenInfo.email, subscription_tier: tokenInfo.subscription_tier, - entitlement_status: tokenInfo.entitlement_status + entitlement_status: tokenInfo.entitlement_status, + base_url: 'https://cli-chat-proxy.grok.com/v1' } if (tokenInfo.refresh_token) credentials.refresh_token = tokenInfo.refresh_token if (tokenInfo.id_token) credentials.id_token = tokenInfo.id_token diff --git a/frontend/src/i18n/__tests__/openaiFastPolicyLocales.spec.ts b/frontend/src/i18n/__tests__/openaiFastPolicyLocales.spec.ts new file mode 100644 index 0000000000..7fd550c1ad --- /dev/null +++ b/frontend/src/i18n/__tests__/openaiFastPolicyLocales.spec.ts @@ -0,0 +1,30 @@ +import { describe, expect, it } from 'vitest' + +import en from '../locales/en' +import zh from '../locales/zh' + +describe('OpenAI Fast/Flex policy locale keys', () => { + it('exposes user scope copy at the runtime zh path', () => { + expect(zh.admin.settings.openaiFastPolicy).toMatchObject({ + userIds: '指定用户', + userIdsHint: '输入任意邮箱关键词进行模糊搜索。留空表示对全部 Sub2API 用户生效；选中用户的 API Key 请求优先匹配用户规则。', + userSearchPlaceholder: '输入用户邮箱搜索', + userSearchEmpty: '未找到匹配用户', + userDeleted: '（已删除）', + userIdFallback: '用户 #{id}', + removeUser: '移除用户' + }) + }) + + it('exposes user scope copy at the runtime en path', () => { + expect(en.admin.settings.openaiFastPolicy).toMatchObject({ + userIds: 'Specific users', + userIdsHint: 'Type any part of a user email to search. Leave empty to apply to all Sub2API users. Selected users match requests from their API keys and take precedence over global rules.', + userSearchPlaceholder: 'Search by user email', + userSearchEmpty: 'No matching users found', + userDeleted: '(deleted)', + userIdFallback: 'User #{id}', + removeUser: 'Remove user' + }) + }) +}) diff --git a/frontend/src/i18n/__tests__/opsLocaleKeys.spec.ts b/frontend/src/i18n/__tests__/opsLocaleKeys.spec.ts index c396991a0c..d44a47dbc8 100644 --- a/frontend/src/i18n/__tests__/opsLocaleKeys.spec.ts +++ b/frontend/src/i18n/__tests__/opsLocaleKeys.spec.ts @@ -1,5 +1,6 @@ import { describe, expect, it } from 'vitest' import en from '@/i18n/locales/en' +import zh from '@/i18n/locales/zh' function flattenKeys(obj: Record, prefix = ''): string[] { const keys: string[] = [] @@ -35,4 +36,18 @@ describe('groups locale key completeness', () => { const enKeys = flattenKeys(en) expect(enKeys).toContain('admin.groups.failedToSave') }) + + const webSearchPricingKeys = [ + 'admin.groups.webSearchPricing.title', + 'admin.groups.webSearchPricing.pricePerCall', + 'admin.groups.webSearchPricing.pricePerCallHint', + 'admin.groups.webSearchPricing.finalPricePreview', + ] + + for (const key of webSearchPricingKeys) { + it(`en and zh locales both have ${key}`, () => { + expect(flattenKeys(en)).toContain(key) + expect(flattenKeys(zh)).toContain(key) + }) + } }) diff --git a/frontend/src/i18n/locales/en/admin/overview.ts b/frontend/src/i18n/locales/en/admin/overview.ts index 7012911065..2a643ae192 100644 --- a/frontend/src/i18n/locales/en/admin/overview.ts +++ b/frontend/src/i18n/locales/en/admin/overview.ts @@ -939,6 +939,13 @@ export default { finalPricePreview: 'Final per-second price preview', notConfigured: 'Not configured' }, + webSearchPricing: { + title: 'Codex Web Search Pricing', + pricePerCall: 'Price per search call (USD)', + pricePerCallHint: + 'Leave empty to use the default $0.01 per call (official pricing: $10 per 1,000 calls); 0 means free. The group rate multiplier is applied on top.', + finalPricePreview: 'Per-call price after current multiplier: {price}' + }, peakRate: { enable: 'Enable peak rate multiplier', peakStart: 'Peak start', diff --git a/frontend/src/i18n/locales/en/admin/settings.ts b/frontend/src/i18n/locales/en/admin/settings.ts index 37dc8aa1c8..e08be9f591 100644 --- a/frontend/src/i18n/locales/en/admin/settings.ts +++ b/frontend/src/i18n/locales/en/admin/settings.ts @@ -979,11 +979,6 @@ export default { scopeOAuth: 'OAuth only', scopeAPIKey: 'API Key only', scopeBedrock: 'Bedrock only', - userIds: 'Specific user IDs', - userIdsHint: 'Leave empty to apply to all Sub2API users. Specified users match requests from their API keys and take precedence over global rules.', - userIdPlaceholder: 'e.g., 1001', - addUserId: 'Add user ID', - removeUserId: 'Remove user ID', errorMessage: 'Error message', errorMessagePlaceholder: 'Custom error message when blocked', errorMessageHint: 'Leave empty for default message', @@ -1024,6 +1019,13 @@ export default { scopeOAuth: 'OAuth only', scopeAPIKey: 'API Key only', scopeBedrock: 'Bedrock only', + userIds: 'Specific users', + userIdsHint: 'Type any part of a user email to search. Leave empty to apply to all Sub2API users. Selected users match requests from their API keys and take precedence over global rules.', + userSearchPlaceholder: 'Search by user email', + userSearchEmpty: 'No matching users found', + userDeleted: '(deleted)', + userIdFallback: 'User #{id}', + removeUser: 'Remove user', errorMessage: 'Error message', errorMessagePlaceholder: 'Custom error message when blocked', errorMessageHint: 'Leave empty for the default message.', diff --git a/frontend/src/i18n/locales/en/dashboard.ts b/frontend/src/i18n/locales/en/dashboard.ts index 0a62bac7cf..3edf99158a 100644 --- a/frontend/src/i18n/locales/en/dashboard.ts +++ b/frontend/src/i18n/locales/en/dashboard.ts @@ -146,6 +146,7 @@ export default { geminiCli: 'Gemini CLI', codexCli: 'Codex CLI', codexCliWs: 'Codex CLI (WebSocket)', + grokCli: 'Grok CLI', opencode: 'OpenCode', }, antigravity: { @@ -160,6 +161,12 @@ export default { modelComment: 'If you have Gemini 3 access, you can use: gemini-3-pro-preview', note: 'These environment variables will be active in the current terminal session. For permanent configuration, add them to ~/.bashrc, ~/.zshrc, or the appropriate configuration file.', }, + grok: { + description: 'Configure Grok Build or OpenCode to send Responses API traffic through your Sub2API Grok group.', + configTomlHint: 'Back up an existing config.toml before merging this model entry. Run grok inspect after saving to verify the effective configuration.', + note: 'Save the file as ~/.grok/config.toml, then run grok inspect and select sub2api-grok from /model.', + noteWindows: 'Save the file as %USERPROFILE%\\.grok\\config.toml, then run grok inspect and select sub2api-grok from /model.', + }, opencode: { title: 'OpenCode Example', subtitle: 'opencode.json', diff --git a/frontend/src/i18n/locales/zh/admin/overview.ts b/frontend/src/i18n/locales/zh/admin/overview.ts index a40fe4afe2..d2638bb79d 100644 --- a/frontend/src/i18n/locales/zh/admin/overview.ts +++ b/frontend/src/i18n/locales/zh/admin/overview.ts @@ -930,6 +930,13 @@ export default { finalPricePreview: '最终每秒价格预览', notConfigured: '未配置' }, + webSearchPricing: { + title: 'Codex 网页搜索计费', + pricePerCall: '搜索单次价格（USD/次）', + pricePerCallHint: + '留空使用默认价 $0.01/次（官方定价 $10/1000 次）；填 0 表示免费。实际扣费会叠加分组费率倍数。', + finalPricePreview: '应用当前倍率后的单次价格：{price}' + }, peakRate: { enable: '启用高峰倍率', peakStart: '高峰开始', diff --git a/frontend/src/i18n/locales/zh/admin/settings.ts b/frontend/src/i18n/locales/zh/admin/settings.ts index 5c0d874b57..08c0dbcd53 100644 --- a/frontend/src/i18n/locales/zh/admin/settings.ts +++ b/frontend/src/i18n/locales/zh/admin/settings.ts @@ -974,11 +974,6 @@ export default { scopeOAuth: '仅 OAuth 账号', scopeAPIKey: '仅 API Key 账号', scopeBedrock: '仅 Bedrock 账号', - userIds: '指定用户 ID', - userIdsHint: '留空表示对全部 Sub2API 用户生效。指定后仅匹配这些用户的 API Key 请求，且优先于全局规则。', - userIdPlaceholder: '例如: 1001', - addUserId: '添加用户 ID', - removeUserId: '移除用户 ID', errorMessage: '错误消息', errorMessagePlaceholder: '拦截时返回的自定义错误消息', errorMessageHint: '留空则使用默认错误消息', @@ -1019,6 +1014,13 @@ export default { scopeOAuth: '仅 OAuth 账号', scopeAPIKey: '仅 API Key 账号', scopeBedrock: '仅 Bedrock 账号', + userIds: '指定用户', + userIdsHint: '输入任意邮箱关键词进行模糊搜索。留空表示对全部 Sub2API 用户生效；选中用户的 API Key 请求优先匹配用户规则。', + userSearchPlaceholder: '输入用户邮箱搜索', + userSearchEmpty: '未找到匹配用户', + userDeleted: '（已删除）', + userIdFallback: '用户 #{id}', + removeUser: '移除用户', errorMessage: '错误消息', errorMessagePlaceholder: '拦截时返回的自定义错误消息', errorMessageHint: '留空则使用默认错误消息。', diff --git a/frontend/src/i18n/locales/zh/dashboard.ts b/frontend/src/i18n/locales/zh/dashboard.ts index f457f10e38..6b689cdf18 100644 --- a/frontend/src/i18n/locales/zh/dashboard.ts +++ b/frontend/src/i18n/locales/zh/dashboard.ts @@ -147,6 +147,7 @@ export default { geminiCli: 'Gemini CLI', codexCli: 'Codex CLI', codexCliWs: 'Codex CLI (WebSocket)', + grokCli: 'Grok CLI', opencode: 'OpenCode' }, antigravity: { @@ -164,6 +165,12 @@ export default { modelComment: '如果你有 Gemini 3 权限可以填：gemini-3-pro-preview', note: '这些环境变量将在当前终端会话中生效。如需永久配置，请将其添加到 ~/.bashrc、~/.zshrc 或相应的配置文件中。' }, + grok: { + description: '配置 Grok Build 或 OpenCode，让 Responses API 请求通过当前 Sub2API Grok 分组发送。', + configTomlHint: '如已有 config.toml，请先备份再合并此模型配置。保存后运行 grok inspect 验证生效配置。', + note: '保存为 ~/.grok/config.toml，然后运行 grok inspect，并在 /model 中选择 sub2api-grok。', + noteWindows: '保存为 %USERPROFILE%\\.grok\\config.toml，然后运行 grok inspect，并在 /model 中选择 sub2api-grok。' + }, opencode: { title: 'OpenCode 配置示例', subtitle: 'opencode.json', diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 0e79563cc6..bd8137829e 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -531,6 +531,8 @@ export interface Group { video_price_480p: number | null video_price_720p: number | null video_price_1080p: number | null + // Codex 网页搜索单次价格（USD/次）；null 表示使用默认价 0.01 + web_search_price_per_call: number | null // 高峰时段倍率配置 peak_rate_enabled: boolean peak_start: string @@ -664,6 +666,7 @@ export interface CreateGroupRequest { video_price_480p?: number | null video_price_720p?: number | null video_price_1080p?: number | null + web_search_price_per_call?: number | null peak_rate_enabled?: boolean peak_start?: string peak_end?: string @@ -711,6 +714,7 @@ export interface UpdateGroupRequest { video_price_480p?: number | null video_price_720p?: number | null video_price_1080p?: number | null + web_search_price_per_call?: number | null peak_rate_enabled?: boolean peak_start?: string peak_end?: string diff --git a/frontend/src/views/admin/GroupsView.vue b/frontend/src/views/admin/GroupsView.vue index 1f22216100..d7cf9dd717 100644 --- a/frontend/src/views/admin/GroupsView.vue +++ b/frontend/src/views/admin/GroupsView.vue @@ -1315,6 +1315,41 @@
+ +
+

+ {{ t("admin.groups.webSearchPricing.title") }} +

+
+ + +

+ {{ t("admin.groups.webSearchPricing.pricePerCallHint") }} +

+
+ {{ + t("admin.groups.webSearchPricing.finalPricePreview", { + price: createWebSearchFinalPricePreview, + }) + }} +
+
+
+
+ +
+

+ {{ t("admin.groups.webSearchPricing.title") }} +

+
+ + +

+ {{ t("admin.groups.webSearchPricing.pricePerCallHint") }} +

+
+ {{ + t("admin.groups.webSearchPricing.finalPricePreview", { + price: editWebSearchFinalPricePreview, + }) + }} +
+
+
+
buildVideoFinalPricePreview(editForm), ); +// Codex 网页搜索单次默认价（与后端 defaultWebSearchPricePerCall 一致，官方 $10/1000 次） +const DEFAULT_WEB_SEARCH_PRICE_PER_CALL = 0.01; + +const buildWebSearchFinalPricePreview = (form: { + web_search_price_per_call: number | string | null; + rate_multiplier: number | string | null; +}) => { + const basePrice = + parsePreviewPrice(form.web_search_price_per_call) ?? + DEFAULT_WEB_SEARCH_PRICE_PER_CALL; + const multiplier = normalizePreviewNumber(form.rate_multiplier, 1); + return formatImagePricePreview(basePrice * multiplier); +}; + +const createWebSearchFinalPricePreview = computed(() => + buildWebSearchFinalPricePreview(createForm), +); +const editWebSearchFinalPricePreview = computed(() => + buildWebSearchFinalPricePreview(editForm), +); + const resetDisabledBatchImagePricing = ( form: Pick< ImagePricingFormState, @@ -4615,6 +4710,7 @@ const closeCreateModal = () => { createForm.video_price_480p = null; createForm.video_price_720p = null; createForm.video_price_1080p = null; + createForm.web_search_price_per_call = null; createForm.peak_rate_enabled = false; createForm.peak_start = ""; createForm.peak_end = ""; @@ -4726,6 +4822,9 @@ const handleCreateGroup = async () => { requestData.video_price_480p = emptyToNull(requestData.video_price_480p); requestData.video_price_720p = emptyToNull(requestData.video_price_720p); requestData.video_price_1080p = emptyToNull(requestData.video_price_1080p); + requestData.web_search_price_per_call = emptyToNull( + requestData.web_search_price_per_call, + ); requestData.peak_rate_enabled = createForm.peak_rate_enabled; requestData.peak_start = createForm.peak_start; requestData.peak_end = createForm.peak_end; @@ -4779,6 +4878,7 @@ const handleEdit = async (group: AdminGroup) => { editForm.video_price_480p = group.video_price_480p; editForm.video_price_720p = group.video_price_720p; editForm.video_price_1080p = group.video_price_1080p; + editForm.web_search_price_per_call = group.web_search_price_per_call ?? null; editForm.peak_rate_enabled = group.peak_rate_enabled ?? false; editForm.peak_start = group.peak_start ?? ""; editForm.peak_end = group.peak_end ?? ""; @@ -4836,6 +4936,7 @@ const closeEditModal = () => { editForm.video_price_480p = null; editForm.video_price_720p = null; editForm.video_price_1080p = null; + editForm.web_search_price_per_call = null; resetMessagesDispatchFormState(editForm); resetModelsListState(editModelsListState); }; @@ -4914,6 +5015,9 @@ const handleUpdateGroup = async () => { payload.video_price_480p = emptyPriceToClear(payload.video_price_480p); payload.video_price_720p = emptyPriceToClear(payload.video_price_720p); payload.video_price_1080p = emptyPriceToClear(payload.video_price_1080p); + payload.web_search_price_per_call = emptyPriceToClear( + payload.web_search_price_per_call, + ); payload.peak_rate_enabled = editForm.peak_rate_enabled; payload.peak_start = editForm.peak_start; payload.peak_end = editForm.peak_end; diff --git a/frontend/src/views/admin/SettingsView.vue b/frontend/src/views/admin/SettingsView.vue index a87cea9060..26253c7eb8 100644 --- a/frontend/src/views/admin/SettingsView.vue +++ b/frontend/src/views/admin/SettingsView.vue @@ -1199,60 +1199,10 @@

{{ t("admin.settings.openaiFastPolicy.userIdsHint") }}

-
- - -
- +
@@ -7431,6 +7381,7 @@ import ProxySelector from "@/components/common/ProxySelector.vue"; import ImageUpload from "@/components/common/ImageUpload.vue"; import BackupSettings from "@/views/admin/BackupView.vue"; import EmailTemplateEditor from "@/views/admin/settings/EmailTemplateEditor.vue"; +import OpenAIFastPolicyUserSelector from "@/views/admin/settings/OpenAIFastPolicyUserSelector.vue"; import { useClipboard } from "@/composables/useClipboard"; import { affiliatesAPI, type AffiliateAdminEntry, type SimpleUser as AffiliateSimpleUser } from "@/api/admin/affiliates"; import { extractApiErrorMessage, extractI18nErrorMessage } from "@/utils/apiError"; @@ -10226,18 +10177,6 @@ function removeOpenAIFastPolicyRule(index: number) { openaiFastPolicyForm.rules.splice(index, 1); } -function addOpenAIFastPolicyUserID(rule: OpenAIFastPolicyRule) { - if (!rule.user_ids) rule.user_ids = []; - rule.user_ids.push(0); -} - -function removeOpenAIFastPolicyUserID( - rule: OpenAIFastPolicyRule, - idx: number, -) { - rule.user_ids?.splice(idx, 1); -} - function addOpenAIFastPolicyModelPattern(rule: OpenAIFastPolicyRule) { if (!rule.model_whitelist) rule.model_whitelist = []; rule.model_whitelist.push(""); diff --git a/frontend/src/views/admin/settings/OpenAIFastPolicyUserSelector.vue b/frontend/src/views/admin/settings/OpenAIFastPolicyUserSelector.vue new file mode 100644 index 0000000000..0b25c876ad --- /dev/null +++ b/frontend/src/views/admin/settings/OpenAIFastPolicyUserSelector.vue @@ -0,0 +1,229 @@ + + + diff --git a/frontend/src/views/admin/settings/__tests__/OpenAIFastPolicyUserSelector.spec.ts b/frontend/src/views/admin/settings/__tests__/OpenAIFastPolicyUserSelector.spec.ts new file mode 100644 index 0000000000..c8e1f0e1cf --- /dev/null +++ b/frontend/src/views/admin/settings/__tests__/OpenAIFastPolicyUserSelector.spec.ts @@ -0,0 +1,112 @@ +import { flushPromises, mount } from '@vue/test-utils' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +import OpenAIFastPolicyUserSelector from '../OpenAIFastPolicyUserSelector.vue' + +const messages: Record = { + 'admin.settings.openaiFastPolicy.userDeleted': '(deleted)', + 'admin.settings.openaiFastPolicy.userIdFallback': 'User #{id}', + 'admin.settings.openaiFastPolicy.removeUser': 'Remove user', + 'admin.settings.openaiFastPolicy.userSearchPlaceholder': 'Search users', + 'admin.settings.openaiFastPolicy.userSearchEmpty': 'No users found', + 'common.loading': 'Loading', +} + +vi.mock('vue-i18n', () => ({ + useI18n: () => ({ + t: (key: string, params?: Record) => { + const message = messages[key] ?? key + return params + ? Object.entries(params).reduce( + (value, [name, replacement]) => value.replace(`{${name}}`, String(replacement)), + message, + ) + : message + }, + }), +})) + +const mockSearchUsers = vi.fn() +const mockGetUserById = vi.fn() + +vi.mock('@/api/admin', () => ({ + adminAPI: { + usage: { + searchUsers: (...args: unknown[]) => mockSearchUsers(...args), + }, + users: { + getById: (...args: unknown[]) => mockGetUserById(...args), + }, + }, +})) + +describe('OpenAIFastPolicyUserSelector', () => { + beforeEach(() => { + vi.useFakeTimers() + mockSearchUsers.mockReset() + mockGetUserById.mockReset() + }) + + afterEach(() => { + vi.useRealTimers() + }) + + it('hydrates existing IDs to email labels without changing the saved IDs', async () => { + mockGetUserById.mockResolvedValue({ + id: 7, + email: 'existing@example.com', + deleted_at: null, + }) + + const wrapper = mount(OpenAIFastPolicyUserSelector, { + props: { modelValue: [7] }, + global: { stubs: { Icon: true } }, + }) + await flushPromises() + + expect(mockGetUserById).toHaveBeenCalledWith(7, true) + expect(wrapper.text()).toContain('existing@example.com') + expect(wrapper.text()).toContain('#7') + expect(wrapper.emitted('update:modelValue')).toBeUndefined() + }) + + it('searches after one character and adds the selected user ID', async () => { + mockSearchUsers.mockResolvedValue([ + { id: 9, email: 'alice@example.com', deleted: false }, + ]) + + const wrapper = mount(OpenAIFastPolicyUserSelector, { + props: { modelValue: [] }, + global: { stubs: { Icon: true } }, + }) + const input = wrapper.get('input') + await input.trigger('focus') + await input.setValue('a') + await input.trigger('input') + vi.advanceTimersByTime(300) + await flushPromises() + + expect(mockSearchUsers).toHaveBeenCalledWith('a') + const result = wrapper.findAll('button').find((button) => + button.text().includes('alice@example.com'), + ) + expect(result).toBeDefined() + await result!.trigger('click') + + expect(wrapper.emitted('update:modelValue')).toEqual([[[9]]]) + }) + + it('keeps an unresolved saved ID visible and removable', async () => { + mockGetUserById.mockRejectedValue(new Error('not found')) + + const wrapper = mount(OpenAIFastPolicyUserSelector, { + props: { modelValue: [42] }, + global: { stubs: { Icon: true } }, + }) + await flushPromises() + + expect(wrapper.text()).toContain('User #42') + await wrapper.get('button[aria-label="Remove user"]').trigger('click') + expect(wrapper.emitted('update:modelValue')).toEqual([[[]]]) + }) +})