chore(openai): 同步上游并解决认证冲突

This commit is contained in:
cat committed 2026-07-14 19:04:04 +08:00
commit d68a2aac1a
388 files changed
+27587 -1591

No files matched your search

+9
View File
@@ -8,6 +8,15 @@ permissions:
contents: read
jobs:
shell:
runs-on: macos-15
steps:
- uses: actions/checkout@v6
- name: Check deployment scripts
run: |
/bin/bash -n deploy/apple-container.sh
/bin/bash deploy/tests/apple-container-test.sh
test:
runs-on: ubuntu-latest
steps:
+2
View File
@@ -116,6 +116,8 @@ backend/.installed
# 其他
# ===================
tests
!deploy/tests/
!deploy/tests/**
CLAUDE.md
.claude
scripts
+91 -12
View File
@@ -163,6 +163,12 @@ Model authenticity: no content intervention or secondary filtering — experienc
</td>
</tr>
<tr>
<td width="180"><a href="http://aimzoon.com"><img src="assets/partners/logos/aimzoon.jpg" alt="aimzoon" width="150"></a></td>
<td>Thanks to Aimzoon for sponsoring this project! <a href="http://aimzoon.com">Aimzoon</a> provides stable, cost-effective AI API access services, enabling developers to quickly connect popular AI services to coding tools such as Codex, Claude Code, and Gemini CLI. No complex configuration — faster onboarding, more stable calls, and lower costs. Ongoing promotions including discounted Codex rates and special pricing, with free trial credits upon registration, bringing AI coding into your daily workflow. <a href="http://aimzoon.com">Click here</a> to register and try it out!
</td>
</tr>
</table>
## Overview
@@ -329,6 +335,7 @@ cd sub2api/deploy
# 2. Copy environment configuration
cp .env.example .env
chmod 600 .env
# 3. Edit configuration (generate secure passwords)
nano .env
@@ -448,7 +455,23 @@ rm -rf data/ postgres_data/ redis_data/
---
### Method 3: Build from Source
### Method 3: Apple container (macOS)
Apple-silicon Macs running macOS 26 can run the full Sub2API, PostgreSQL, and Redis stack with Apple `container` 1.1.0 or newer:
```bash
git clone https://github.com/Wei-Shaw/sub2api.git
cd sub2api/deploy
./apple-container.sh init
./apple-container.sh up
./apple-container.sh status
```
This is an operator-managed local workflow; Docker Compose remains the recommended production path. See [deploy/APPLE_CONTAINER.md](deploy/APPLE_CONTAINER.md) for lifecycle commands, persistence, upgrades, and runtime limitations.
---
### Method 4: Build from Source
Build and run from source code for development or customization.
@@ -579,6 +602,27 @@ If you disable URL validation or response header filtering, harden your network
- Enforce TLS-only outbound traffic
- Strip sensitive upstream response headers at the proxy
#### OpenAI Responses WebSocket ingress limits
`gateway.openai_ws` bounds the lifetime and aggregate count of client-facing
Responses WebSocket sessions. These safeguards apply independently from
per-turn user and account concurrency slots, which are released between turns.
```yaml
gateway:
openai_ws:
# Close a client socket idle between completed turns; 0 disables this safeguard.
ingress_inter_turn_idle_timeout_seconds: 300
# Distributed API-key limit for live client ingress sessions; 0 disables it.
max_ingress_connections_per_api_key: 64
```
The connection cap is coordinated through Redis using a 60-second lease that
is refreshed every 20 seconds. A process that cannot confirm a lease for a
full lease lifetime closes its local WebSocket rather than continuing outside
the global cap. Use `http_bridge` for client-WebSocket/upstream-HTTP operation
when rolling out or mitigating upstream WebSocket issues.
#### ⚠️ Important: Creating the Admin Account
The initial admin account is **only created via the setup wizard** (served at `http://<host>:8080` on first run). The `default.admin_email` / `default.admin_password` fields in `config.yaml` are **not used** to create it — they exist in the template for historical reasons.
@@ -637,20 +681,20 @@ 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`
- 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.
- 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/edits`, `/videos/edits`, `/v1/videos/extensions`, `/videos/extensions`, `/v1/videos/{request_id}`, and `/videos/{request_id}`. Generation, editing, and extension 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 +709,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 +721,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.
---
+23 -1
View File
@@ -166,6 +166,11 @@
</td>
</tr>
<tr>
<td width="180"><a href="http://aimzoon.com"><img src="assets/partners/logos/aimzoon.jpg" alt="aimzoon" width="150"></a></td>
<td>感谢 Aimzoon 对本项目的赞助! <a href="http://aimzoon.com">Aimzoon</a> 提供稳定、高性价比的 AI API 接入服务,支持开发者将常用 AI 服务快速接入 Codex、Claude Code、Gemini CLI 等编程工具。无需复杂配置,更快接入,更稳调用,更省成本。codex倍率优惠,特价倍率等促销不断,注册即送免费体验额度,让 AI 编程真正进入日常工作流。<a href="http://aimzoon.com">点击这里</a>注册体验!
</td>
</tr>
</table>
@@ -333,6 +338,7 @@ cd sub2api/deploy
# 2. 复制环境配置文件
cp .env.example .env
chmod 600 .env
# 3. 编辑配置(生成安全密码)
nano .env
@@ -464,7 +470,23 @@ rm -rf data/ postgres_data/ redis_data/
---
### 方式三:源码编译
### 方式三:Apple container(macOS)
Apple 芯片 Mac 在 macOS 26 上可使用 Apple `container` 1.1.0 或更高版本运行完整的 Sub2API、PostgreSQL 和 Redis:
```bash
git clone https://github.com/Wei-Shaw/sub2api.git
cd sub2api/deploy
./apple-container.sh init
./apple-container.sh up
./apple-container.sh status
```
该方式面向本地开发和人工运维,不提供持续重启监管;生产部署仍推荐 Docker Compose。生命周期命令、持久化、升级和运行时限制见 [deploy/APPLE_CONTAINER.md](deploy/APPLE_CONTAINER.md)。
---
### 方式四:源码编译
从源码编译安装,适合开发或定制需求。
+24 -1
View File
@@ -161,6 +161,12 @@
</td>
</tr>
<tr>
<td width="180"><a href="http://aimzoon.com"><img src="assets/partners/logos/aimzoon.jpg" alt="aimzoon" width="150"></a></td>
<td>Aimzoon のご支援に感謝します!<a href="http://aimzoon.com">Aimzoon</a> は安定してコストパフォーマンスに優れた AI API 接続サービスを提供し、開発者が主要な AI サービスを Codex、Claude Code、Gemini CLI などのコーディングツールへ素早く接続できるようにします。複雑な設定は不要で、より速い接続、より安定した呼び出し、より低いコストを実現。Codex レート割引や特価レートなどのキャンペーンも随時開催中、登録するだけで無料お試しクレジットをプレゼント。AI コーディングを日常のワークフローへ。<a href="http://aimzoon.com">こちら</a>から登録してお試しください!
</td>
</tr>
</table>
## 概要
@@ -327,6 +333,7 @@ cd sub2api/deploy
# 2. 環境設定ファイルをコピー
cp .env.example .env
chmod 600 .env
# 3. 設定を編集(セキュアなパスワードを生成)
nano .env
@@ -446,7 +453,23 @@ rm -rf data/ postgres_data/ redis_data/
---
### 方法3: ソースからビルド
### 方法3: Apple container(macOS)
Apple シリコン搭載 Mac と macOS 26 では、Apple `container` 1.1.0 以降を使用して Sub2API、PostgreSQL、Redis の完全なスタックを実行できます:
```bash
git clone https://github.com/Wei-Shaw/sub2api.git
cd sub2api/deploy
./apple-container.sh init
./apple-container.sh up
./apple-container.sh status
```
これはローカル開発および手動運用向けです。本番環境では引き続き Docker Compose を推奨します。ライフサイクル、永続化、アップグレード、制限については [deploy/APPLE_CONTAINER.md](deploy/APPLE_CONTAINER.md) を参照してください。
---
### 方法4: ソースからビルド
開発やカスタマイズのためにソースコードからビルドして実行します。
Binary file not shown.

After

Width:  |  Height:  |  Size: 25 KiB

+1 -1
View File
@@ -1 +1 @@
0.1.150
0.1.155
+4 -4
View File
@@ -182,12 +182,13 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
claudeUsageFetcher := repository.NewClaudeUsageFetcher(httpUpstream)
antigravityQuotaFetcher := service.NewAntigravityQuotaFetcher(proxyRepository)
grokQuotaFetcher := service.NewGrokQuotaFetcher()
grokQuotaService := service.ProvideGrokQuotaService(accountRepository, proxyRepository, grokTokenProvider, httpUpstream, usageLogRepository)
openAIQuotaService := service.ProvideOpenAIQuotaService(accountRepository, proxyRepository, openAITokenProvider, privacyClientFactory, openAIGatewayService)
usageCache := service.NewUsageCache()
accountUsageService := service.ProvideAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService, openAIGatewayService)
accountUsageService := service.ProvideAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher, grokQuotaFetcher, grokQuotaService, openAIQuotaService, usageCache, identityCache, tlsFingerprintProfileService, openAIGatewayService)
accountTestService := service.ProvideAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService, openAIGatewayService)
crsSyncService := service.NewCRSSyncService(accountRepository, proxyRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig)
accountHandler := admin.NewAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator)
accountHandler := admin.ProvideAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator, grokQuotaService)
adminAnnouncementHandler := admin.NewAnnouncementHandler(announcementService)
dataManagementService := service.NewDataManagementService()
dataManagementHandler := admin.NewDataManagementHandler(dataManagementService)
@@ -199,7 +200,6 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
openAIOAuthHandler := admin.NewOpenAIOAuthHandler(openAIOAuthService, adminService, openAIQuotaService)
geminiOAuthHandler := admin.NewGeminiOAuthHandler(geminiOAuthService)
antigravityOAuthHandler := admin.NewAntigravityOAuthHandler(antigravityOAuthService)
grokQuotaService := service.ProvideGrokQuotaService(accountRepository, proxyRepository, grokTokenProvider, httpUpstream)
grokOAuthHandler := admin.NewGrokOAuthHandler(grokOAuthService, adminService, grokQuotaService)
proxyHandler := admin.NewProxyHandler(adminService)
adminRedeemHandler := admin.NewRedeemHandler(adminService, redeemService)
@@ -256,7 +256,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
openAIGatewayHandler := handler.NewOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, configConfig)
handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService)
totpHandler := handler.NewTotpHandler(totpService)
handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService, channelService)
handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService)
paymentWebhookHandler := handler.NewPaymentWebhookHandler(paymentService, registry)
availableChannelHandler := handler.NewAvailableChannelHandler(channelService, apiKeyService, settingService)
batchImageRepository := repository.NewBatchImageRepository(db)
+2 -1
View File
@@ -167,6 +167,7 @@ const (
ProviderOpenai Provider = "openai"
ProviderAnthropic Provider = "anthropic"
ProviderGemini Provider = "gemini"
ProviderGrok Provider = "grok"
)
func (pr Provider) String() string {
@@ -176,7 +177,7 @@ func (pr Provider) String() string {
// ProviderValidator is a validator for the "provider" field enum values. It is called by the builders before save.
func ProviderValidator(pr Provider) error {
switch pr {
case ProviderOpenai, ProviderAnthropic, ProviderGemini:
case ProviderOpenai, ProviderAnthropic, ProviderGemini, ProviderGrok:
return nil
default:
return fmt.Errorf("channelmonitor: invalid enum value for provider field: %q", pr)
@@ -103,6 +103,7 @@ const (
ProviderOpenai Provider = "openai"
ProviderAnthropic Provider = "anthropic"
ProviderGemini Provider = "gemini"
ProviderGrok Provider = "grok"
)
func (pr Provider) String() string {
@@ -112,7 +113,7 @@ func (pr Provider) String() string {
// ProviderValidator is a validator for the "provider" field enum values. It is called by the builders before save.
func ProviderValidator(pr Provider) error {
switch pr {
case ProviderOpenai, ProviderAnthropic, ProviderGemini:
case ProviderOpenai, ProviderAnthropic, ProviderGemini, ProviderGrok:
return nil
default:
return fmt.Errorf("channelmonitorrequesttemplate: invalid enum value for provider field: %q", pr)
+15 -1
View File
@@ -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(", ")
+8
View File
@@ -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()
+55
View File
@@ -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))
+98
View File
@@ -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) {
+72
View File
@@ -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)
}
+19 -17
View File
@@ -623,7 +623,7 @@ var (
{Name: "created_at", Type: field.TypeTime, SchemaType: map[string]string{"postgres": "timestamptz"}},
{Name: "updated_at", Type: field.TypeTime, SchemaType: map[string]string{"postgres": "timestamptz"}},
{Name: "name", Type: field.TypeString, Size: 100},
{Name: "provider", Type: field.TypeEnum, Enums: []string{"openai", "anthropic", "gemini"}},
{Name: "provider", Type: field.TypeEnum, Enums: []string{"openai", "anthropic", "gemini", "grok"}},
{Name: "api_mode", Type: field.TypeString, Size: 32, Default: "chat_completions"},
{Name: "endpoint", Type: field.TypeString, Size: 500},
{Name: "api_key_encrypted", Type: field.TypeString},
@@ -768,7 +768,7 @@ var (
{Name: "created_at", Type: field.TypeTime, SchemaType: map[string]string{"postgres": "timestamptz"}},
{Name: "updated_at", Type: field.TypeTime, SchemaType: map[string]string{"postgres": "timestamptz"}},
{Name: "name", Type: field.TypeString, Size: 100},
{Name: "provider", Type: field.TypeEnum, Enums: []string{"openai", "anthropic", "gemini"}},
{Name: "provider", Type: field.TypeEnum, Enums: []string{"openai", "anthropic", "gemini", "grok"}},
{Name: "api_mode", Type: field.TypeString, Size: 32, Default: "chat_completions"},
{Name: "description", Type: field.TypeString, Nullable: true, Size: 500, Default: ""},
{Name: "extra_headers", Type: field.TypeJSON},
@@ -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]},
},
},
}
@@ -1559,6 +1560,7 @@ var (
{Name: "total_cost", Type: field.TypeFloat64, Default: 0, SchemaType: map[string]string{"postgres": "decimal(20,10)"}},
{Name: "actual_cost", Type: field.TypeFloat64, Default: 0, SchemaType: map[string]string{"postgres": "decimal(20,10)"}},
{Name: "rate_multiplier", Type: field.TypeFloat64, Default: 1, SchemaType: map[string]string{"postgres": "decimal(10,4)"}},
{Name: "long_context_billing_applied", Type: field.TypeBool, Default: false},
{Name: "account_rate_multiplier", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(10,4)"}},
{Name: "billing_type", Type: field.TypeInt8, Default: 0},
{Name: "stream", Type: field.TypeBool, Default: false},
@@ -1591,31 +1593,31 @@ var (
ForeignKeys: []*schema.ForeignKey{
{
Symbol: "usage_logs_api_keys_usage_logs",
Columns: []*schema.Column{UsageLogsColumns[40]},
Columns: []*schema.Column{UsageLogsColumns[41]},
RefColumns: []*schema.Column{APIKeysColumns[0]},
OnDelete: schema.NoAction,
},
{
Symbol: "usage_logs_accounts_usage_logs",
Columns: []*schema.Column{UsageLogsColumns[41]},
Columns: []*schema.Column{UsageLogsColumns[42]},
RefColumns: []*schema.Column{AccountsColumns[0]},
OnDelete: schema.NoAction,
},
{
Symbol: "usage_logs_groups_usage_logs",
Columns: []*schema.Column{UsageLogsColumns[42]},
Columns: []*schema.Column{UsageLogsColumns[43]},
RefColumns: []*schema.Column{GroupsColumns[0]},
OnDelete: schema.SetNull,
},
{
Symbol: "usage_logs_users_usage_logs",
Columns: []*schema.Column{UsageLogsColumns[43]},
Columns: []*schema.Column{UsageLogsColumns[44]},
RefColumns: []*schema.Column{UsersColumns[0]},
OnDelete: schema.NoAction,
},
{
Symbol: "usage_logs_user_subscriptions_usage_logs",
Columns: []*schema.Column{UsageLogsColumns[44]},
Columns: []*schema.Column{UsageLogsColumns[45]},
RefColumns: []*schema.Column{UserSubscriptionsColumns[0]},
OnDelete: schema.SetNull,
},
@@ -1624,32 +1626,32 @@ var (
{
Name: "usagelog_user_id",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[43]},
Columns: []*schema.Column{UsageLogsColumns[44]},
},
{
Name: "usagelog_api_key_id",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[40]},
Columns: []*schema.Column{UsageLogsColumns[41]},
},
{
Name: "usagelog_account_id",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[41]},
Columns: []*schema.Column{UsageLogsColumns[42]},
},
{
Name: "usagelog_group_id",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[42]},
Columns: []*schema.Column{UsageLogsColumns[43]},
},
{
Name: "usagelog_subscription_id",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[44]},
Columns: []*schema.Column{UsageLogsColumns[45]},
},
{
Name: "usagelog_created_at",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[39]},
Columns: []*schema.Column{UsageLogsColumns[40]},
},
{
Name: "usagelog_model",
@@ -1669,17 +1671,17 @@ var (
{
Name: "usagelog_user_id_created_at",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[43], UsageLogsColumns[39]},
Columns: []*schema.Column{UsageLogsColumns[44], UsageLogsColumns[40]},
},
{
Name: "usagelog_api_key_id_created_at",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[40], UsageLogsColumns[39]},
Columns: []*schema.Column{UsageLogsColumns[41], UsageLogsColumns[40]},
},
{
Name: "usagelog_group_id_created_at",
Unique: false,
Columns: []*schema.Column{UsageLogsColumns[42], UsageLogsColumns[39]},
Columns: []*schema.Column{UsageLogsColumns[43], UsageLogsColumns[40]},
},
},
}
+240 -79
View File
@@ -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
@@ -41656,83 +41763,84 @@ func (m *UsageCleanupTaskMutation) ResetEdge(name string) error {
// UsageLogMutation represents an operation that mutates the UsageLog nodes in the graph.
type UsageLogMutation struct {
config
op Op
typ string
id *int64
request_id *string
model *string
requested_model *string
upstream_model *string
channel_id *int64
addchannel_id *int64
model_mapping_chain *string
billing_tier *string
billing_mode *string
input_tokens *int
addinput_tokens *int
output_tokens *int
addoutput_tokens *int
cache_creation_tokens *int
addcache_creation_tokens *int
cache_read_tokens *int
addcache_read_tokens *int
cache_creation_5m_tokens *int
addcache_creation_5m_tokens *int
cache_creation_1h_tokens *int
addcache_creation_1h_tokens *int
input_cost *float64
addinput_cost *float64
output_cost *float64
addoutput_cost *float64
cache_creation_cost *float64
addcache_creation_cost *float64
cache_read_cost *float64
addcache_read_cost *float64
total_cost *float64
addtotal_cost *float64
actual_cost *float64
addactual_cost *float64
rate_multiplier *float64
addrate_multiplier *float64
account_rate_multiplier *float64
addaccount_rate_multiplier *float64
billing_type *int8
addbilling_type *int8
stream *bool
duration_ms *int
addduration_ms *int
first_token_ms *int
addfirst_token_ms *int
user_agent *string
ip_address *string
image_count *int
addimage_count *int
image_size *string
image_input_size *string
image_output_size *string
image_size_source *string
image_size_breakdown *map[string]int
video_count *int
addvideo_count *int
video_resolution *string
video_duration_seconds *int
addvideo_duration_seconds *int
cache_ttl_overridden *bool
created_at *time.Time
clearedFields map[string]struct{}
user *int64
cleareduser bool
api_key *int64
clearedapi_key bool
account *int64
clearedaccount bool
group *int64
clearedgroup bool
subscription *int64
clearedsubscription bool
done bool
oldValue func(context.Context) (*UsageLog, error)
predicates []predicate.UsageLog
op Op
typ string
id *int64
request_id *string
model *string
requested_model *string
upstream_model *string
channel_id *int64
addchannel_id *int64
model_mapping_chain *string
billing_tier *string
billing_mode *string
input_tokens *int
addinput_tokens *int
output_tokens *int
addoutput_tokens *int
cache_creation_tokens *int
addcache_creation_tokens *int
cache_read_tokens *int
addcache_read_tokens *int
cache_creation_5m_tokens *int
addcache_creation_5m_tokens *int
cache_creation_1h_tokens *int
addcache_creation_1h_tokens *int
input_cost *float64
addinput_cost *float64
output_cost *float64
addoutput_cost *float64
cache_creation_cost *float64
addcache_creation_cost *float64
cache_read_cost *float64
addcache_read_cost *float64
total_cost *float64
addtotal_cost *float64
actual_cost *float64
addactual_cost *float64
rate_multiplier *float64
addrate_multiplier *float64
long_context_billing_applied *bool
account_rate_multiplier *float64
addaccount_rate_multiplier *float64
billing_type *int8
addbilling_type *int8
stream *bool
duration_ms *int
addduration_ms *int
first_token_ms *int
addfirst_token_ms *int
user_agent *string
ip_address *string
image_count *int
addimage_count *int
image_size *string
image_input_size *string
image_output_size *string
image_size_source *string
image_size_breakdown *map[string]int
video_count *int
addvideo_count *int
video_resolution *string
video_duration_seconds *int
addvideo_duration_seconds *int
cache_ttl_overridden *bool
created_at *time.Time
clearedFields map[string]struct{}
user *int64
cleareduser bool
api_key *int64
clearedapi_key bool
account *int64
clearedaccount bool
group *int64
clearedgroup bool
subscription *int64
clearedsubscription bool
done bool
oldValue func(context.Context) (*UsageLog, error)
predicates []predicate.UsageLog
}
var _ ent.Mutation = (*UsageLogMutation)(nil)
@@ -43154,6 +43262,42 @@ func (m *UsageLogMutation) ResetRateMultiplier() {
m.addrate_multiplier = nil
}
// SetLongContextBillingApplied sets the "long_context_billing_applied" field.
func (m *UsageLogMutation) SetLongContextBillingApplied(b bool) {
m.long_context_billing_applied = &b
}
// LongContextBillingApplied returns the value of the "long_context_billing_applied" field in the mutation.
func (m *UsageLogMutation) LongContextBillingApplied() (r bool, exists bool) {
v := m.long_context_billing_applied
if v == nil {
return
}
return *v, true
}
// OldLongContextBillingApplied returns the old "long_context_billing_applied" field's value of the UsageLog entity.
// If the UsageLog object wasn't provided to the builder, the object is fetched from the database.
// An error is returned if the mutation operation is not UpdateOne, or the database query fails.
func (m *UsageLogMutation) OldLongContextBillingApplied(ctx context.Context) (v bool, err error) {
if !m.op.Is(OpUpdateOne) {
return v, errors.New("OldLongContextBillingApplied is only allowed on UpdateOne operations")
}
if m.id == nil || m.oldValue == nil {
return v, errors.New("OldLongContextBillingApplied requires an ID field in the mutation")
}
oldValue, err := m.oldValue(ctx)
if err != nil {
return v, fmt.Errorf("querying old value for OldLongContextBillingApplied: %w", err)
}
return oldValue.LongContextBillingApplied, nil
}
// ResetLongContextBillingApplied resets all changes to the "long_context_billing_applied" field.
func (m *UsageLogMutation) ResetLongContextBillingApplied() {
m.long_context_billing_applied = nil
}
// SetAccountRateMultiplier sets the "account_rate_multiplier" field.
func (m *UsageLogMutation) SetAccountRateMultiplier(f float64) {
m.account_rate_multiplier = &f
@@ -44271,7 +44415,7 @@ func (m *UsageLogMutation) Type() string {
// order to get all numeric fields that were incremented/decremented, call
// AddedFields().
func (m *UsageLogMutation) Fields() []string {
fields := make([]string, 0, 44)
fields := make([]string, 0, 45)
if m.user != nil {
fields = append(fields, usagelog.FieldUserID)
}
@@ -44350,6 +44494,9 @@ func (m *UsageLogMutation) Fields() []string {
if m.rate_multiplier != nil {
fields = append(fields, usagelog.FieldRateMultiplier)
}
if m.long_context_billing_applied != nil {
fields = append(fields, usagelog.FieldLongContextBillingApplied)
}
if m.account_rate_multiplier != nil {
fields = append(fields, usagelog.FieldAccountRateMultiplier)
}
@@ -44464,6 +44611,8 @@ func (m *UsageLogMutation) Field(name string) (ent.Value, bool) {
return m.ActualCost()
case usagelog.FieldRateMultiplier:
return m.RateMultiplier()
case usagelog.FieldLongContextBillingApplied:
return m.LongContextBillingApplied()
case usagelog.FieldAccountRateMultiplier:
return m.AccountRateMultiplier()
case usagelog.FieldBillingType:
@@ -44561,6 +44710,8 @@ func (m *UsageLogMutation) OldField(ctx context.Context, name string) (ent.Value
return m.OldActualCost(ctx)
case usagelog.FieldRateMultiplier:
return m.OldRateMultiplier(ctx)
case usagelog.FieldLongContextBillingApplied:
return m.OldLongContextBillingApplied(ctx)
case usagelog.FieldAccountRateMultiplier:
return m.OldAccountRateMultiplier(ctx)
case usagelog.FieldBillingType:
@@ -44788,6 +44939,13 @@ func (m *UsageLogMutation) SetField(name string, value ent.Value) error {
}
m.SetRateMultiplier(v)
return nil
case usagelog.FieldLongContextBillingApplied:
v, ok := value.(bool)
if !ok {
return fmt.Errorf("unexpected type %T for field %s", value, name)
}
m.SetLongContextBillingApplied(v)
return nil
case usagelog.FieldAccountRateMultiplier:
v, ok := value.(float64)
if !ok {
@@ -45419,6 +45577,9 @@ func (m *UsageLogMutation) ResetField(name string) error {
case usagelog.FieldRateMultiplier:
m.ResetRateMultiplier()
return nil
case usagelog.FieldLongContextBillingApplied:
m.ResetLongContextBillingApplied()
return nil
case usagelog.FieldAccountRateMultiplier:
m.ResetAccountRateMultiplier()
return nil
+29 -25
View File
@@ -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()
@@ -1940,56 +1940,60 @@ func init() {
usagelogDescRateMultiplier := usagelogFields[25].Descriptor()
// usagelog.DefaultRateMultiplier holds the default value on creation for the rate_multiplier field.
usagelog.DefaultRateMultiplier = usagelogDescRateMultiplier.Default.(float64)
// usagelogDescLongContextBillingApplied is the schema descriptor for long_context_billing_applied field.
usagelogDescLongContextBillingApplied := usagelogFields[26].Descriptor()
// usagelog.DefaultLongContextBillingApplied holds the default value on creation for the long_context_billing_applied field.
usagelog.DefaultLongContextBillingApplied = usagelogDescLongContextBillingApplied.Default.(bool)
// usagelogDescBillingType is the schema descriptor for billing_type field.
usagelogDescBillingType := usagelogFields[27].Descriptor()
usagelogDescBillingType := usagelogFields[28].Descriptor()
// usagelog.DefaultBillingType holds the default value on creation for the billing_type field.
usagelog.DefaultBillingType = usagelogDescBillingType.Default.(int8)
// usagelogDescStream is the schema descriptor for stream field.
usagelogDescStream := usagelogFields[28].Descriptor()
usagelogDescStream := usagelogFields[29].Descriptor()
// usagelog.DefaultStream holds the default value on creation for the stream field.
usagelog.DefaultStream = usagelogDescStream.Default.(bool)
// usagelogDescUserAgent is the schema descriptor for user_agent field.
usagelogDescUserAgent := usagelogFields[31].Descriptor()
usagelogDescUserAgent := usagelogFields[32].Descriptor()
// usagelog.UserAgentValidator is a validator for the "user_agent" field. It is called by the builders before save.
usagelog.UserAgentValidator = usagelogDescUserAgent.Validators[0].(func(string) error)
// usagelogDescIPAddress is the schema descriptor for ip_address field.
usagelogDescIPAddress := usagelogFields[32].Descriptor()
usagelogDescIPAddress := usagelogFields[33].Descriptor()
// usagelog.IPAddressValidator is a validator for the "ip_address" field. It is called by the builders before save.
usagelog.IPAddressValidator = usagelogDescIPAddress.Validators[0].(func(string) error)
// usagelogDescImageCount is the schema descriptor for image_count field.
usagelogDescImageCount := usagelogFields[33].Descriptor()
usagelogDescImageCount := usagelogFields[34].Descriptor()
// usagelog.DefaultImageCount holds the default value on creation for the image_count field.
usagelog.DefaultImageCount = usagelogDescImageCount.Default.(int)
// usagelogDescImageSize is the schema descriptor for image_size field.
usagelogDescImageSize := usagelogFields[34].Descriptor()
usagelogDescImageSize := usagelogFields[35].Descriptor()
// usagelog.ImageSizeValidator is a validator for the "image_size" field. It is called by the builders before save.
usagelog.ImageSizeValidator = usagelogDescImageSize.Validators[0].(func(string) error)
// usagelogDescImageInputSize is the schema descriptor for image_input_size field.
usagelogDescImageInputSize := usagelogFields[35].Descriptor()
usagelogDescImageInputSize := usagelogFields[36].Descriptor()
// usagelog.ImageInputSizeValidator is a validator for the "image_input_size" field. It is called by the builders before save.
usagelog.ImageInputSizeValidator = usagelogDescImageInputSize.Validators[0].(func(string) error)
// usagelogDescImageOutputSize is the schema descriptor for image_output_size field.
usagelogDescImageOutputSize := usagelogFields[36].Descriptor()
usagelogDescImageOutputSize := usagelogFields[37].Descriptor()
// usagelog.ImageOutputSizeValidator is a validator for the "image_output_size" field. It is called by the builders before save.
usagelog.ImageOutputSizeValidator = usagelogDescImageOutputSize.Validators[0].(func(string) error)
// usagelogDescImageSizeSource is the schema descriptor for image_size_source field.
usagelogDescImageSizeSource := usagelogFields[37].Descriptor()
usagelogDescImageSizeSource := usagelogFields[38].Descriptor()
// usagelog.ImageSizeSourceValidator is a validator for the "image_size_source" field. It is called by the builders before save.
usagelog.ImageSizeSourceValidator = usagelogDescImageSizeSource.Validators[0].(func(string) error)
// usagelogDescVideoCount is the schema descriptor for video_count field.
usagelogDescVideoCount := usagelogFields[39].Descriptor()
usagelogDescVideoCount := usagelogFields[40].Descriptor()
// usagelog.DefaultVideoCount holds the default value on creation for the video_count field.
usagelog.DefaultVideoCount = usagelogDescVideoCount.Default.(int)
// usagelogDescVideoResolution is the schema descriptor for video_resolution field.
usagelogDescVideoResolution := usagelogFields[40].Descriptor()
usagelogDescVideoResolution := usagelogFields[41].Descriptor()
// usagelog.VideoResolutionValidator is a validator for the "video_resolution" field. It is called by the builders before save.
usagelog.VideoResolutionValidator = usagelogDescVideoResolution.Validators[0].(func(string) error)
// usagelogDescCacheTTLOverridden is the schema descriptor for cache_ttl_overridden field.
usagelogDescCacheTTLOverridden := usagelogFields[42].Descriptor()
usagelogDescCacheTTLOverridden := usagelogFields[43].Descriptor()
// usagelog.DefaultCacheTTLOverridden holds the default value on creation for the cache_ttl_overridden field.
usagelog.DefaultCacheTTLOverridden = usagelogDescCacheTTLOverridden.Default.(bool)
// usagelogDescCreatedAt is the schema descriptor for created_at field.
usagelogDescCreatedAt := usagelogFields[43].Descriptor()
usagelogDescCreatedAt := usagelogFields[44].Descriptor()
// usagelog.DefaultCreatedAt holds the default value on creation for the created_at field.
usagelog.DefaultCreatedAt = usagelogDescCreatedAt.Default.(func() time.Time)
userMixin := schema.User{}.Mixin()
+1 -1
View File
@@ -35,7 +35,7 @@ func (ChannelMonitor) Fields() []ent.Field {
NotEmpty().
MaxLen(100),
field.Enum("provider").
Values("openai", "anthropic", "gemini"),
Values("openai", "anthropic", "gemini", "grok"),
field.String("api_mode").
Default("chat_completions").
MaxLen(32).
@@ -39,7 +39,7 @@ func (ChannelMonitorRequestTemplate) Fields() []ent.Field {
NotEmpty().
MaxLen(100),
field.Enum("provider").
Values("openai", "anthropic", "gemini"),
Values("openai", "anthropic", "gemini", "grok"),
field.String("api_mode").
Default("chat_completions").
MaxLen(32).
+5
View File
@@ -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").
+3
View File
@@ -100,6 +100,9 @@ func (UsageLog) Fields() []ent.Field {
field.Float("rate_multiplier").
Default(1).
SchemaType(map[string]string{dialect.Postgres: "decimal(10,4)"}),
field.Bool("long_context_billing_applied").
Default(false).
Comment("Whether long-context pricing changed token prices for this request"),
// account_rate_multiplier: 账号计费倍率快照(NULL 表示按 1.0 处理)
field.Float("account_rate_multiplier").
+12 -1
View File
@@ -75,6 +75,8 @@ type UsageLog struct {
ActualCost float64 `json:"actual_cost,omitempty"`
// RateMultiplier holds the value of the "rate_multiplier" field.
RateMultiplier float64 `json:"rate_multiplier,omitempty"`
// Whether long-context pricing changed token prices for this request
LongContextBillingApplied bool `json:"long_context_billing_applied,omitempty"`
// AccountRateMultiplier holds the value of the "account_rate_multiplier" field.
AccountRateMultiplier *float64 `json:"account_rate_multiplier,omitempty"`
// BillingType holds the value of the "billing_type" field.
@@ -196,7 +198,7 @@ func (*UsageLog) scanValues(columns []string) ([]any, error) {
switch columns[i] {
case usagelog.FieldImageSizeBreakdown:
values[i] = new([]byte)
case usagelog.FieldStream, usagelog.FieldCacheTTLOverridden:
case usagelog.FieldLongContextBillingApplied, usagelog.FieldStream, usagelog.FieldCacheTTLOverridden:
values[i] = new(sql.NullBool)
case usagelog.FieldInputCost, usagelog.FieldOutputCost, usagelog.FieldCacheCreationCost, usagelog.FieldCacheReadCost, usagelog.FieldTotalCost, usagelog.FieldActualCost, usagelog.FieldRateMultiplier, usagelog.FieldAccountRateMultiplier:
values[i] = new(sql.NullFloat64)
@@ -391,6 +393,12 @@ func (_m *UsageLog) assignValues(columns []string, values []any) error {
} else if value.Valid {
_m.RateMultiplier = value.Float64
}
case usagelog.FieldLongContextBillingApplied:
if value, ok := values[i].(*sql.NullBool); !ok {
return fmt.Errorf("unexpected type %T for field long_context_billing_applied", values[i])
} else if value.Valid {
_m.LongContextBillingApplied = value.Bool
}
case usagelog.FieldAccountRateMultiplier:
if value, ok := values[i].(*sql.NullFloat64); !ok {
return fmt.Errorf("unexpected type %T for field account_rate_multiplier", values[i])
@@ -667,6 +675,9 @@ func (_m *UsageLog) String() string {
builder.WriteString("rate_multiplier=")
builder.WriteString(fmt.Sprintf("%v", _m.RateMultiplier))
builder.WriteString(", ")
builder.WriteString("long_context_billing_applied=")
builder.WriteString(fmt.Sprintf("%v", _m.LongContextBillingApplied))
builder.WriteString(", ")
if v := _m.AccountRateMultiplier; v != nil {
builder.WriteString("account_rate_multiplier=")
builder.WriteString(fmt.Sprintf("%v", *v))
+10
View File
@@ -66,6 +66,8 @@ const (
FieldActualCost = "actual_cost"
// FieldRateMultiplier holds the string denoting the rate_multiplier field in the database.
FieldRateMultiplier = "rate_multiplier"
// FieldLongContextBillingApplied holds the string denoting the long_context_billing_applied field in the database.
FieldLongContextBillingApplied = "long_context_billing_applied"
// FieldAccountRateMultiplier holds the string denoting the account_rate_multiplier field in the database.
FieldAccountRateMultiplier = "account_rate_multiplier"
// FieldBillingType holds the string denoting the billing_type field in the database.
@@ -180,6 +182,7 @@ var Columns = []string{
FieldTotalCost,
FieldActualCost,
FieldRateMultiplier,
FieldLongContextBillingApplied,
FieldAccountRateMultiplier,
FieldBillingType,
FieldStream,
@@ -251,6 +254,8 @@ var (
DefaultActualCost float64
// DefaultRateMultiplier holds the default value on creation for the "rate_multiplier" field.
DefaultRateMultiplier float64
// DefaultLongContextBillingApplied holds the default value on creation for the "long_context_billing_applied" field.
DefaultLongContextBillingApplied bool
// DefaultBillingType holds the default value on creation for the "billing_type" field.
DefaultBillingType int8
// DefaultStream holds the default value on creation for the "stream" field.
@@ -417,6 +422,11 @@ func ByRateMultiplier(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldRateMultiplier, opts...).ToFunc()
}
// ByLongContextBillingApplied orders the results by the long_context_billing_applied field.
func ByLongContextBillingApplied(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldLongContextBillingApplied, opts...).ToFunc()
}
// ByAccountRateMultiplier orders the results by the account_rate_multiplier field.
func ByAccountRateMultiplier(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldAccountRateMultiplier, opts...).ToFunc()
+15
View File
@@ -185,6 +185,11 @@ func RateMultiplier(v float64) predicate.UsageLog {
return predicate.UsageLog(sql.FieldEQ(FieldRateMultiplier, v))
}
// LongContextBillingApplied applies equality check predicate on the "long_context_billing_applied" field. It's identical to LongContextBillingAppliedEQ.
func LongContextBillingApplied(v bool) predicate.UsageLog {
return predicate.UsageLog(sql.FieldEQ(FieldLongContextBillingApplied, v))
}
// AccountRateMultiplier applies equality check predicate on the "account_rate_multiplier" field. It's identical to AccountRateMultiplierEQ.
func AccountRateMultiplier(v float64) predicate.UsageLog {
return predicate.UsageLog(sql.FieldEQ(FieldAccountRateMultiplier, v))
@@ -1465,6 +1470,16 @@ func RateMultiplierLTE(v float64) predicate.UsageLog {
return predicate.UsageLog(sql.FieldLTE(FieldRateMultiplier, v))
}
// LongContextBillingAppliedEQ applies the EQ predicate on the "long_context_billing_applied" field.
func LongContextBillingAppliedEQ(v bool) predicate.UsageLog {
return predicate.UsageLog(sql.FieldEQ(FieldLongContextBillingApplied, v))
}
// LongContextBillingAppliedNEQ applies the NEQ predicate on the "long_context_billing_applied" field.
func LongContextBillingAppliedNEQ(v bool) predicate.UsageLog {
return predicate.UsageLog(sql.FieldNEQ(FieldLongContextBillingApplied, v))
}
// AccountRateMultiplierEQ applies the EQ predicate on the "account_rate_multiplier" field.
func AccountRateMultiplierEQ(v float64) predicate.UsageLog {
return predicate.UsageLog(sql.FieldEQ(FieldAccountRateMultiplier, v))
+65
View File
@@ -351,6 +351,20 @@ func (_c *UsageLogCreate) SetNillableRateMultiplier(v *float64) *UsageLogCreate
return _c
}
// SetLongContextBillingApplied sets the "long_context_billing_applied" field.
func (_c *UsageLogCreate) SetLongContextBillingApplied(v bool) *UsageLogCreate {
_c.mutation.SetLongContextBillingApplied(v)
return _c
}
// SetNillableLongContextBillingApplied sets the "long_context_billing_applied" field if the given value is not nil.
func (_c *UsageLogCreate) SetNillableLongContextBillingApplied(v *bool) *UsageLogCreate {
if v != nil {
_c.SetLongContextBillingApplied(*v)
}
return _c
}
// SetAccountRateMultiplier sets the "account_rate_multiplier" field.
func (_c *UsageLogCreate) SetAccountRateMultiplier(v float64) *UsageLogCreate {
_c.mutation.SetAccountRateMultiplier(v)
@@ -707,6 +721,10 @@ func (_c *UsageLogCreate) defaults() {
v := usagelog.DefaultRateMultiplier
_c.mutation.SetRateMultiplier(v)
}
if _, ok := _c.mutation.LongContextBillingApplied(); !ok {
v := usagelog.DefaultLongContextBillingApplied
_c.mutation.SetLongContextBillingApplied(v)
}
if _, ok := _c.mutation.BillingType(); !ok {
v := usagelog.DefaultBillingType
_c.mutation.SetBillingType(v)
@@ -824,6 +842,9 @@ func (_c *UsageLogCreate) check() error {
if _, ok := _c.mutation.RateMultiplier(); !ok {
return &ValidationError{Name: "rate_multiplier", err: errors.New(`ent: missing required field "UsageLog.rate_multiplier"`)}
}
if _, ok := _c.mutation.LongContextBillingApplied(); !ok {
return &ValidationError{Name: "long_context_billing_applied", err: errors.New(`ent: missing required field "UsageLog.long_context_billing_applied"`)}
}
if _, ok := _c.mutation.BillingType(); !ok {
return &ValidationError{Name: "billing_type", err: errors.New(`ent: missing required field "UsageLog.billing_type"`)}
}
@@ -997,6 +1018,10 @@ func (_c *UsageLogCreate) createSpec() (*UsageLog, *sqlgraph.CreateSpec) {
_spec.SetField(usagelog.FieldRateMultiplier, field.TypeFloat64, value)
_node.RateMultiplier = value
}
if value, ok := _c.mutation.LongContextBillingApplied(); ok {
_spec.SetField(usagelog.FieldLongContextBillingApplied, field.TypeBool, value)
_node.LongContextBillingApplied = value
}
if value, ok := _c.mutation.AccountRateMultiplier(); ok {
_spec.SetField(usagelog.FieldAccountRateMultiplier, field.TypeFloat64, value)
_node.AccountRateMultiplier = &value
@@ -1650,6 +1675,18 @@ func (u *UsageLogUpsert) AddRateMultiplier(v float64) *UsageLogUpsert {
return u
}
// SetLongContextBillingApplied sets the "long_context_billing_applied" field.
func (u *UsageLogUpsert) SetLongContextBillingApplied(v bool) *UsageLogUpsert {
u.Set(usagelog.FieldLongContextBillingApplied, v)
return u
}
// UpdateLongContextBillingApplied sets the "long_context_billing_applied" field to the value that was provided on create.
func (u *UsageLogUpsert) UpdateLongContextBillingApplied() *UsageLogUpsert {
u.SetExcluded(usagelog.FieldLongContextBillingApplied)
return u
}
// SetAccountRateMultiplier sets the "account_rate_multiplier" field.
func (u *UsageLogUpsert) SetAccountRateMultiplier(v float64) *UsageLogUpsert {
u.Set(usagelog.FieldAccountRateMultiplier, v)
@@ -2531,6 +2568,20 @@ func (u *UsageLogUpsertOne) UpdateRateMultiplier() *UsageLogUpsertOne {
})
}
// SetLongContextBillingApplied sets the "long_context_billing_applied" field.
func (u *UsageLogUpsertOne) SetLongContextBillingApplied(v bool) *UsageLogUpsertOne {
return u.Update(func(s *UsageLogUpsert) {
s.SetLongContextBillingApplied(v)
})
}
// UpdateLongContextBillingApplied sets the "long_context_billing_applied" field to the value that was provided on create.
func (u *UsageLogUpsertOne) UpdateLongContextBillingApplied() *UsageLogUpsertOne {
return u.Update(func(s *UsageLogUpsert) {
s.UpdateLongContextBillingApplied()
})
}
// SetAccountRateMultiplier sets the "account_rate_multiplier" field.
func (u *UsageLogUpsertOne) SetAccountRateMultiplier(v float64) *UsageLogUpsertOne {
return u.Update(func(s *UsageLogUpsert) {
@@ -3631,6 +3682,20 @@ func (u *UsageLogUpsertBulk) UpdateRateMultiplier() *UsageLogUpsertBulk {
})
}
// SetLongContextBillingApplied sets the "long_context_billing_applied" field.
func (u *UsageLogUpsertBulk) SetLongContextBillingApplied(v bool) *UsageLogUpsertBulk {
return u.Update(func(s *UsageLogUpsert) {
s.SetLongContextBillingApplied(v)
})
}
// UpdateLongContextBillingApplied sets the "long_context_billing_applied" field to the value that was provided on create.
func (u *UsageLogUpsertBulk) UpdateLongContextBillingApplied() *UsageLogUpsertBulk {
return u.Update(func(s *UsageLogUpsert) {
s.UpdateLongContextBillingApplied()
})
}
// SetAccountRateMultiplier sets the "account_rate_multiplier" field.
func (u *UsageLogUpsertBulk) SetAccountRateMultiplier(v float64) *UsageLogUpsertBulk {
return u.Update(func(s *UsageLogUpsert) {
+34
View File
@@ -542,6 +542,20 @@ func (_u *UsageLogUpdate) AddRateMultiplier(v float64) *UsageLogUpdate {
return _u
}
// SetLongContextBillingApplied sets the "long_context_billing_applied" field.
func (_u *UsageLogUpdate) SetLongContextBillingApplied(v bool) *UsageLogUpdate {
_u.mutation.SetLongContextBillingApplied(v)
return _u
}
// SetNillableLongContextBillingApplied sets the "long_context_billing_applied" field if the given value is not nil.
func (_u *UsageLogUpdate) SetNillableLongContextBillingApplied(v *bool) *UsageLogUpdate {
if v != nil {
_u.SetLongContextBillingApplied(*v)
}
return _u
}
// SetAccountRateMultiplier sets the "account_rate_multiplier" field.
func (_u *UsageLogUpdate) SetAccountRateMultiplier(v float64) *UsageLogUpdate {
_u.mutation.ResetAccountRateMultiplier()
@@ -1199,6 +1213,9 @@ func (_u *UsageLogUpdate) sqlSave(ctx context.Context) (_node int, err error) {
if value, ok := _u.mutation.AddedRateMultiplier(); ok {
_spec.AddField(usagelog.FieldRateMultiplier, field.TypeFloat64, value)
}
if value, ok := _u.mutation.LongContextBillingApplied(); ok {
_spec.SetField(usagelog.FieldLongContextBillingApplied, field.TypeBool, value)
}
if value, ok := _u.mutation.AccountRateMultiplier(); ok {
_spec.SetField(usagelog.FieldAccountRateMultiplier, field.TypeFloat64, value)
}
@@ -1982,6 +1999,20 @@ func (_u *UsageLogUpdateOne) AddRateMultiplier(v float64) *UsageLogUpdateOne {
return _u
}
// SetLongContextBillingApplied sets the "long_context_billing_applied" field.
func (_u *UsageLogUpdateOne) SetLongContextBillingApplied(v bool) *UsageLogUpdateOne {
_u.mutation.SetLongContextBillingApplied(v)
return _u
}
// SetNillableLongContextBillingApplied sets the "long_context_billing_applied" field if the given value is not nil.
func (_u *UsageLogUpdateOne) SetNillableLongContextBillingApplied(v *bool) *UsageLogUpdateOne {
if v != nil {
_u.SetLongContextBillingApplied(*v)
}
return _u
}
// SetAccountRateMultiplier sets the "account_rate_multiplier" field.
func (_u *UsageLogUpdateOne) SetAccountRateMultiplier(v float64) *UsageLogUpdateOne {
_u.mutation.ResetAccountRateMultiplier()
@@ -2669,6 +2700,9 @@ func (_u *UsageLogUpdateOne) sqlSave(ctx context.Context) (_node *UsageLog, err
if value, ok := _u.mutation.AddedRateMultiplier(); ok {
_spec.AddField(usagelog.FieldRateMultiplier, field.TypeFloat64, value)
}
if value, ok := _u.mutation.LongContextBillingApplied(); ok {
_spec.SetField(usagelog.FieldLongContextBillingApplied, field.TypeBool, value)
}
if value, ok := _u.mutation.AccountRateMultiplier(); ok {
_spec.SetField(usagelog.FieldAccountRateMultiplier, field.TypeFloat64, value)
}
+1 -1
View File
@@ -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
+8
View File
@@ -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=
+29
View File
@@ -601,6 +601,7 @@ type ServerConfig struct {
Host string `mapstructure:"host"`
Port int `mapstructure:"port"`
Mode string `mapstructure:"mode"` // debug/release
EnableServerTiming bool `mapstructure:"enable_server_timing"` // Admin UI Server-Timing response header
FrontendURL string `mapstructure:"frontend_url"` // 前端基础 URL,用于生成邮件中的外部链接
ReadHeaderTimeout int `mapstructure:"read_header_timeout"` // 读取请求头超时(秒)
IdleTimeout int `mapstructure:"idle_timeout"` // 空闲连接超时(秒)
@@ -818,6 +819,8 @@ type GatewayConfig struct {
ImageStreamDataIntervalTimeout int `mapstructure:"image_stream_data_interval_timeout"`
// ImageStreamKeepaliveInterval: 图片流式 keepalive 间隔(秒),0表示禁用
ImageStreamKeepaliveInterval int `mapstructure:"image_stream_keepalive_interval"`
// ImageNonstreamKeepaliveInterval: 图片非流式 JSON keepalive 间隔(秒),0表示禁用
ImageNonstreamKeepaliveInterval int `mapstructure:"image_nonstream_keepalive_interval"`
// MaxLineSize: 上游 SSE 单行最大字节数(0使用默认值)
MaxLineSize int `mapstructure:"max_line_size"`
@@ -923,6 +926,12 @@ type GatewayOpenAIWSConfig struct {
ModeRouterV2Enabled bool `mapstructure:"mode_router_v2_enabled"`
// IngressModeDefault: ingress 默认模式(off/ctx_pool/passthrough/http_bridge)
IngressModeDefault string `mapstructure:"ingress_mode_default"`
// IngressInterTurnIdleTimeoutSeconds bounds the time a client may remain idle
// between completed ingress turns. Zero disables this protection.
IngressInterTurnIdleTimeoutSeconds int `mapstructure:"ingress_inter_turn_idle_timeout_seconds"`
// MaxIngressConnectionsPerAPIKey bounds live client WebSocket ingress sessions
// per API key across all instances. Zero disables this protection.
MaxIngressConnectionsPerAPIKey int `mapstructure:"max_ingress_connections_per_api_key"`
// Enabled: 全局总开关(默认 true)
Enabled bool `mapstructure:"enabled"`
// OAuthEnabled: 是否允许 OpenAI OAuth 账号使用 WS
@@ -1453,6 +1462,9 @@ func load(allowMissingJWTSecret bool) (*Config, error) {
// 环境变量支持
viper.AutomaticEnv()
viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
if err := viper.BindEnv("server.enable_server_timing", "ENABLE_SERVER_TIMING"); err != nil {
return nil, fmt.Errorf("bind ENABLE_SERVER_TIMING: %w", err)
}
// 默认值
setDefaults()
@@ -1608,6 +1620,7 @@ func setDefaults() {
viper.SetDefault("server.host", "0.0.0.0")
viper.SetDefault("server.port", 8080)
viper.SetDefault("server.mode", "release")
viper.SetDefault("server.enable_server_timing", false)
viper.SetDefault("server.frontend_url", "")
viper.SetDefault("server.read_header_timeout", 30) // 30秒读取请求头
viper.SetDefault("server.idle_timeout", 120) // 120秒空闲超时
@@ -1945,6 +1958,8 @@ func setDefaults() {
viper.SetDefault("gateway.openai_ws.enabled", true)
viper.SetDefault("gateway.openai_ws.mode_router_v2_enabled", false)
viper.SetDefault("gateway.openai_ws.ingress_mode_default", "ctx_pool")
viper.SetDefault("gateway.openai_ws.ingress_inter_turn_idle_timeout_seconds", 300)
viper.SetDefault("gateway.openai_ws.max_ingress_connections_per_api_key", 64)
viper.SetDefault("gateway.openai_ws.oauth_enabled", true)
viper.SetDefault("gateway.openai_ws.apikey_enabled", true)
viper.SetDefault("gateway.openai_ws.force_http", false)
@@ -2024,6 +2039,7 @@ func setDefaults() {
viper.SetDefault("gateway.stream_keepalive_interval", 10)
viper.SetDefault("gateway.image_stream_data_interval_timeout", 900)
viper.SetDefault("gateway.image_stream_keepalive_interval", 10)
viper.SetDefault("gateway.image_nonstream_keepalive_interval", 0)
viper.SetDefault("gateway.max_line_size", 500*1024*1024)
viper.SetDefault("gateway.scheduling.sticky_session_max_waiting", 3)
viper.SetDefault("gateway.scheduling.sticky_session_wait_timeout", 120*time.Second)
@@ -2710,6 +2726,13 @@ func (c *Config) Validate() error {
(c.Gateway.ImageStreamKeepaliveInterval < 5 || c.Gateway.ImageStreamKeepaliveInterval > 60) {
return fmt.Errorf("gateway.image_stream_keepalive_interval must be 0 or between 5-60 seconds")
}
if c.Gateway.ImageNonstreamKeepaliveInterval < 0 {
return fmt.Errorf("gateway.image_nonstream_keepalive_interval must be non-negative")
}
if c.Gateway.ImageNonstreamKeepaliveInterval != 0 &&
(c.Gateway.ImageNonstreamKeepaliveInterval < 5 || c.Gateway.ImageNonstreamKeepaliveInterval > 60) {
return fmt.Errorf("gateway.image_nonstream_keepalive_interval must be 0 or between 5-60 seconds")
}
// 兼容旧键 sticky_previous_response_ttl_seconds
if c.Gateway.OpenAIWS.StickyResponseIDTTLSeconds <= 0 && c.Gateway.OpenAIWS.StickyPreviousResponseTTLSeconds > 0 {
c.Gateway.OpenAIWS.StickyResponseIDTTLSeconds = c.Gateway.OpenAIWS.StickyPreviousResponseTTLSeconds
@@ -2717,6 +2740,12 @@ func (c *Config) Validate() error {
if c.Gateway.OpenAIWS.MaxConnsPerAccount <= 0 {
return fmt.Errorf("gateway.openai_ws.max_conns_per_account must be positive")
}
if c.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds < 0 {
return fmt.Errorf("gateway.openai_ws.ingress_inter_turn_idle_timeout_seconds must be non-negative")
}
if c.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey < 0 {
return fmt.Errorf("gateway.openai_ws.max_ingress_connections_per_api_key must be non-negative")
}
if c.Gateway.OpenAIWS.MinIdlePerAccount < 0 {
return fmt.Errorf("gateway.openai_ws.min_idle_per_account must be non-negative")
}
+55
View File
@@ -17,6 +17,23 @@ func resetViperWithJWTSecret(t *testing.T) {
t.Setenv("JWT_SECRET", strings.Repeat("x", 32))
}
func TestLoadServerTimingConfig(t *testing.T) {
t.Run("disabled by default", func(t *testing.T) {
resetViperWithJWTSecret(t)
cfg, err := Load()
require.NoError(t, err)
require.False(t, cfg.Server.EnableServerTiming)
})
t.Run("enabled by exact environment variable", func(t *testing.T) {
resetViperWithJWTSecret(t)
t.Setenv("ENABLE_SERVER_TIMING", "true")
cfg, err := Load()
require.NoError(t, err)
require.True(t, cfg.Server.EnableServerTiming)
})
}
func TestLoadForBootstrapAllowsMissingJWTSecret(t *testing.T) {
viper.Reset()
t.Setenv("JWT_SECRET", "")
@@ -182,6 +199,12 @@ func TestLoadDefaultOpenAIWSConfig(t *testing.T) {
if cfg.Gateway.OpenAIWS.IngressModeDefault != "ctx_pool" {
t.Fatalf("Gateway.OpenAIWS.IngressModeDefault = %q, want %q", cfg.Gateway.OpenAIWS.IngressModeDefault, "ctx_pool")
}
if cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds != 300 {
t.Fatalf("Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = %d, want 300", cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds)
}
if cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey != 64 {
t.Fatalf("Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = %d, want 64", cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey)
}
}
func TestLoadDefaultOpenAICompactModel(t *testing.T) {
@@ -236,6 +259,15 @@ func TestLoadOpenAIResponseHeaderTimeoutFromEnv(t *testing.T) {
require.Equal(t, 1800, cfg.Gateway.OpenAIResponseHeaderTimeout)
}
func TestLoadImageNonstreamKeepaliveFromEnv(t *testing.T) {
resetViperWithJWTSecret(t)
t.Setenv("GATEWAY_IMAGE_NONSTREAM_KEEPALIVE_INTERVAL", "15")
cfg, err := Load()
require.NoError(t, err)
require.Equal(t, 15, cfg.Gateway.ImageNonstreamKeepaliveInterval)
}
func TestLoadOpenAIWSStickyTTLCompatibility(t *testing.T) {
resetViperWithJWTSecret(t)
t.Setenv("GATEWAY_OPENAI_WS_STICKY_RESPONSE_ID_TTL_SECONDS", "0")
@@ -1406,6 +1438,16 @@ func TestValidateConfigErrors(t *testing.T) {
mutate: func(c *Config) { c.Gateway.ImageStreamKeepaliveInterval = -1 },
wantErr: "gateway.image_stream_keepalive_interval must be non-negative",
},
{
name: "gateway image nonstream keepalive range",
mutate: func(c *Config) { c.Gateway.ImageNonstreamKeepaliveInterval = 4 },
wantErr: "gateway.image_nonstream_keepalive_interval",
},
{
name: "gateway image nonstream keepalive negative",
mutate: func(c *Config) { c.Gateway.ImageNonstreamKeepaliveInterval = -1 },
wantErr: "gateway.image_nonstream_keepalive_interval must be non-negative",
},
{
name: "gateway image stream data interval range",
mutate: func(c *Config) { c.Gateway.ImageStreamDataIntervalTimeout = 30 },
@@ -1640,6 +1682,16 @@ func TestValidateConfig_OpenAIWSRules(t *testing.T) {
mutate: func(c *Config) { c.Gateway.OpenAIWS.MaxConnsPerAccount = 0 },
wantErr: "gateway.openai_ws.max_conns_per_account",
},
{
name: "ingress_inter_turn_idle_timeout_seconds 不能为负数",
mutate: func(c *Config) { c.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = -1 },
wantErr: "gateway.openai_ws.ingress_inter_turn_idle_timeout_seconds",
},
{
name: "max_ingress_connections_per_api_key 不能为负数",
mutate: func(c *Config) { c.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = -1 },
wantErr: "gateway.openai_ws.max_ingress_connections_per_api_key",
},
{
name: "min_idle_per_account 不能为负数",
mutate: func(c *Config) { c.Gateway.OpenAIWS.MinIdlePerAccount = -1 },
@@ -1964,6 +2016,9 @@ func TestLoad_DefaultGatewayImageStreamConfig(t *testing.T) {
if cfg.Gateway.ImageStreamKeepaliveInterval != 10 {
t.Fatalf("image_stream_keepalive_interval = %d, want 10", cfg.Gateway.ImageStreamKeepaliveInterval)
}
if cfg.Gateway.ImageNonstreamKeepaliveInterval != 0 {
t.Fatalf("image_nonstream_keepalive_interval = %d, want 0", cfg.Gateway.ImageNonstreamKeepaliveInterval)
}
if cfg.Gateway.ImageConcurrency.Enabled {
t.Fatalf("image_concurrency.enabled = true, want false")
}
@@ -120,6 +120,10 @@ func (h *AccountHandler) ImportCodexSession(c *gin.Context) {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
if err := service.ValidateOpenAILongContextBillingExtra(service.PlatformOpenAI, req.Extra); err != nil {
response.ErrorFrom(c, err)
return
}
if req.Concurrency != nil && *req.Concurrency < 0 {
response.BadRequest(c, "concurrency must be >= 0")
return
@@ -630,6 +630,7 @@ func TestImportCodexSessionsAccessTokenOnlySameUserUpdatesExisting(t *testing.T)
"chatgpt_user_id": "user-1",
"access_token": existingToken,
},
Extra: map[string]any{"openai_long_context_billing_enabled": false},
}})
handler := NewAccountHandler(svc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
req := CodexSessionImportRequest{SkipDefaultGroupBind: boolPtr(true)}
@@ -650,6 +651,9 @@ func TestImportCodexSessionsAccessTokenOnlySameUserUpdatesExisting(t *testing.T)
if len(svc.updatedAccounts) != 1 || svc.updatedAccounts[0].id != 10 {
t.Fatalf("updated accounts = %+v, want account 10", svc.updatedAccounts)
}
if got := svc.updatedAccounts[0].input.Extra["openai_long_context_billing_enabled"]; got != false {
t.Fatalf("openai_long_context_billing_enabled = %v, want false", got)
}
}
func TestImportCodexSessionsUpgradesAccessTokenOnlyAccountWithRefreshToken(t *testing.T) {
@@ -460,6 +460,7 @@ func (h *AccountHandler) importData(ctx context.Context, req DataImportRequest)
if created.Platform == service.PlatformAntigravity && created.Type == service.AccountTypeOAuth {
privacyAccounts = append(privacyAccounts, created)
}
h.scheduleGrokImportProbe(created)
result.AccountCreated++
}
@@ -60,6 +60,7 @@ type AccountHandler struct {
sessionLimitCache service.SessionLimitCache
rpmCache service.RPMCache
tokenCacheInvalidator service.TokenCacheInvalidator
grokImportProber grokUsageProber
}
// NewAccountHandler creates a new admin account handler
@@ -784,6 +785,10 @@ func (h *AccountHandler) Create(c *gin.Context) {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
if err := service.ValidateOpenAILongContextBillingExtra(req.Platform, req.Extra); err != nil {
response.ErrorFrom(c, err)
return
}
if req.RateMultiplier != nil && *req.RateMultiplier < 0 {
response.BadRequest(c, "rate_multiplier must be >= 0")
return
@@ -851,6 +856,7 @@ func (h *AccountHandler) Create(c *gin.Context) {
// OpenAI APIKey 账号创建后异步探测上游 /v1/responses 能力。
// 探测失败不影响账号创建响应。
h.scheduleOpenAIResponsesProbe(createdAccount)
h.scheduleGrokImportProbe(createdAccount)
response.Success(c, result.Data)
}
@@ -1299,6 +1305,10 @@ func (h *AccountHandler) ApplyOAuthCredentials(c *gin.Context) {
response.ErrorFrom(c, infraerrors.BadRequest("NOT_OAUTH", "cannot apply oauth credentials to non-OAuth account"))
return
}
if err := service.ValidateOpenAILongContextBillingExtra(existing.Platform, req.Extra); err != nil {
response.ErrorFrom(c, err)
return
}
updatedAccount, err := h.adminService.UpdateAccount(ctx, accountID, &service.UpdateAccountInput{
Type: req.Type,
@@ -1592,6 +1602,12 @@ func (h *AccountHandler) BatchCreate(c *gin.Context) {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
for _, item := range req.Accounts {
if err := service.ValidateOpenAILongContextBillingExtra(item.Platform, item.Extra); err != nil {
response.ErrorFrom(c, err)
return
}
}
executeAdminIdempotentJSON(c, "admin.accounts.batch_create", req, service.DefaultWriteIdempotencyTTL(), func(ctx context.Context) (any, error) {
success := 0
@@ -1653,6 +1669,7 @@ func (h *AccountHandler) BatchCreate(c *gin.Context) {
}
// OpenAI APIKey 账号异步探测 /v1/responses 能力。
h.scheduleOpenAIResponsesProbe(account)
h.scheduleGrokImportProbe(account)
success++
results = append(results, gin.H{
"name": item.Name,
@@ -0,0 +1,165 @@
package admin
import (
"bytes"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestAccountAdminBoundariesRejectMalformedOpenAILongContextBillingValue(t *testing.T) {
const malformedExtra = `"extra":{"openai_long_context_billing_enabled":"true"}`
tests := []struct {
name string
method string
path string
body string
mount func(*gin.Engine, *AccountHandler)
setup func(*stubAdminService)
}{
{
name: "create",
method: http.MethodPost,
path: "/accounts",
body: `{"name":"account","platform":"openai","type":"apikey","credentials":{"api_key":"test"},` + malformedExtra + `}`,
mount: func(router *gin.Engine, handler *AccountHandler) { router.POST("/accounts", handler.Create) },
},
{
name: "update",
method: http.MethodPut,
path: "/accounts/1",
body: `{` + malformedExtra + `}`,
mount: func(router *gin.Engine, handler *AccountHandler) { router.PUT("/accounts/:id", handler.Update) },
setup: func(stub *stubAdminService) {
stub.updateAccountErr = infraerrors.BadRequest("OPENAI_LONG_CONTEXT_BILLING_INVALID", "invalid")
},
},
{
name: "bulk update",
method: http.MethodPost,
path: "/accounts/bulk-update",
body: `{"account_ids":[1],` + malformedExtra + `}`,
mount: func(router *gin.Engine, handler *AccountHandler) {
router.POST("/accounts/bulk-update", handler.BulkUpdate)
},
setup: func(stub *stubAdminService) {
stub.bulkUpdateAccountErr = infraerrors.BadRequest("OPENAI_LONG_CONTEXT_BILLING_INVALID", "invalid")
},
},
{
name: "batch create",
method: http.MethodPost,
path: "/accounts/batch",
body: `{"accounts":[{"name":"account","platform":"openai","type":"apikey","credentials":{"api_key":"test"},` + malformedExtra + `}]}`,
mount: func(router *gin.Engine, handler *AccountHandler) { router.POST("/accounts/batch", handler.BatchCreate) },
},
{
name: "Codex session import",
method: http.MethodPost,
path: "/accounts/import-codex-session",
body: `{"content":"token",` + malformedExtra + `}`,
mount: func(router *gin.Engine, handler *AccountHandler) {
router.POST("/accounts/import-codex-session", handler.ImportCodexSession)
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
gin.SetMode(gin.TestMode)
stub := newStubAdminService()
if tt.setup != nil {
tt.setup(stub)
}
handler := NewAccountHandler(stub, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
router := gin.New()
tt.mount(router, handler)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(tt.method, tt.path, bytes.NewBufferString(tt.body))
request.Header.Set("Content-Type", "application/json")
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusBadRequest, recorder.Code)
var responseBody struct {
Reason string `json:"reason"`
}
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &responseBody))
require.Equal(t, "OPENAI_LONG_CONTEXT_BILLING_INVALID", responseBody.Reason)
})
}
}
func TestAccountCreateBoundaryDoesNotApplyOpenAIValidationToOtherPlatforms(t *testing.T) {
gin.SetMode(gin.TestMode)
handler := NewAccountHandler(newStubAdminService(), nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
router := gin.New()
router.POST("/accounts", handler.Create)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPost, "/accounts", bytes.NewBufferString(
`{"name":"account","platform":"anthropic","type":"apikey","credentials":{"api_key":"test"},"extra":{"openai_long_context_billing_enabled":"provider-owned"}}`,
))
request.Header.Set("Content-Type", "application/json")
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
}
func TestApplyOAuthCredentialsRejectsMalformedOpenAILongContextBillingBeforeMutation(t *testing.T) {
gin.SetMode(gin.TestMode)
stub := newStubAdminService()
stub.getAccountResult = &service.Account{
ID: 1,
Platform: service.PlatformOpenAI,
Type: service.AccountTypeOAuth,
}
handler := NewAccountHandler(stub, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
router := gin.New()
router.POST("/accounts/:id/apply-oauth-credentials", handler.ApplyOAuthCredentials)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPost, "/accounts/1/apply-oauth-credentials", bytes.NewBufferString(
`{"type":"oauth","credentials":{"access_token":"new-token"},"extra":{"openai_long_context_billing_enabled":"true"}}`,
))
request.Header.Set("Content-Type", "application/json")
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusBadRequest, recorder.Code)
var responseBody struct {
Reason string `json:"reason"`
}
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &responseBody))
require.Equal(t, "OPENAI_LONG_CONTEXT_BILLING_INVALID", responseBody.Reason)
require.Zero(t, stub.updateAccountCalls)
require.Zero(t, stub.updateAccountExtraCalls)
}
func TestOpenAIOAuthCodexPATBoundaryRejectsMalformedOpenAILongContextBillingValueBeforeTokenValidation(t *testing.T) {
gin.SetMode(gin.TestMode)
handler := NewOpenAIOAuthHandler(nil, newStubAdminService(), nil)
router := gin.New()
router.Use(gin.Recovery())
router.POST("/openai/create-from-codex-pat", handler.CreateAccountFromCodexPAT)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPost, "/openai/create-from-codex-pat", bytes.NewBufferString(
`{"access_token":"token","extra":{"openai_long_context_billing_enabled":1}}`,
))
request.Header.Set("Content-Type", "application/json")
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusBadRequest, recorder.Code)
var responseBody struct {
Reason string `json:"reason"`
}
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &responseBody))
require.Equal(t, "OPENAI_LONG_CONTEXT_BILLING_INVALID", responseBody.Reason)
}
@@ -33,6 +33,9 @@ type stubAdminService struct {
createSparkShadowErr error
updateAccountErr error
bulkUpdateAccountErr error
getAccountResult *service.Account
updateAccountCalls int
updateAccountExtraCalls int
checkMixedErr error
lastMixedCheck struct {
accountID int64
@@ -388,6 +391,9 @@ func (s *stubAdminService) ListOpenAISchedulableAccountsForSchedulerScore(_ cont
}
func (s *stubAdminService) GetAccount(ctx context.Context, id int64) (*service.Account, error) {
if s.getAccountResult != nil {
return s.getAccountResult, nil
}
account := service.Account{ID: id, Name: "account", Status: service.StatusActive}
return &account, nil
}
@@ -413,6 +419,7 @@ func (s *stubAdminService) CreateAccount(ctx context.Context, input *service.Cre
}
func (s *stubAdminService) UpdateAccount(ctx context.Context, id int64, input *service.UpdateAccountInput) (*service.Account, error) {
s.updateAccountCalls++
if s.updateAccountErr != nil {
return nil, s.updateAccountErr
}
@@ -421,6 +428,7 @@ func (s *stubAdminService) UpdateAccount(ctx context.Context, id int64, input *s
}
func (s *stubAdminService) UpdateAccountExtra(ctx context.Context, id int64, updates map[string]any) error {
s.updateAccountExtraCalls++
return nil
}
@@ -37,11 +37,11 @@ func NewChannelMonitorHandler(monitorService *service.ChannelMonitorService) *Ch
type channelMonitorCreateRequest struct {
Name string `json:"name" binding:"required,max=100"`
Provider string `json:"provider" binding:"required,oneof=openai anthropic gemini"`
Provider string `json:"provider" binding:"required,oneof=openai anthropic gemini grok"`
APIMode string `json:"api_mode" binding:"omitempty,oneof=chat_completions responses"`
Endpoint string `json:"endpoint" binding:"required,max=500"`
APIKey string `json:"api_key" binding:"required,max=2000"`
PrimaryModel string `json:"primary_model" binding:"required,max=200"`
PrimaryModel string `json:"primary_model" binding:"max=200"`
ExtraModels []string `json:"extra_models"`
GroupName string `json:"group_name" binding:"max=100"`
Enabled *bool `json:"enabled"`
@@ -55,7 +55,7 @@ type channelMonitorCreateRequest struct {
type channelMonitorUpdateRequest struct {
Name *string `json:"name" binding:"omitempty,max=100"`
Provider *string `json:"provider" binding:"omitempty,oneof=openai anthropic gemini"`
Provider *string `json:"provider" binding:"omitempty,oneof=openai anthropic gemini grok"`
APIMode *string `json:"api_mode" binding:"omitempty,oneof=chat_completions responses"`
Endpoint *string `json:"endpoint" binding:"omitempty,max=500"`
APIKey *string `json:"api_key" binding:"omitempty,max=2000"`
@@ -26,7 +26,7 @@ func NewChannelMonitorRequestTemplateHandler(templateService *service.ChannelMon
type channelMonitorTemplateCreateRequest struct {
Name string `json:"name" binding:"required,max=100"`
Provider string `json:"provider" binding:"required,oneof=openai anthropic gemini"`
Provider string `json:"provider" binding:"required,oneof=openai anthropic gemini grok"`
APIMode string `json:"api_mode" binding:"omitempty,oneof=chat_completions responses"`
Description string `json:"description" binding:"max=500"`
ExtraHeaders map[string]string `json:"extra_headers"`
@@ -0,0 +1,203 @@
package admin
import (
"context"
"log/slog"
"sync"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/service"
)
const (
grokImportProbeConcurrency = 3
grokImportProbeTimeout = 25 * time.Second
)
type grokUsageProber interface {
ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error)
}
type grokImportProbeTask struct {
prober grokUsageProber
accountID int64
}
type grokImportProbeScheduler struct {
mu sync.Mutex
queue []grokImportProbeTask
concurrency int
workers int
maxWorkers int
timeout time.Duration
}
var defaultGrokImportProbeScheduler = newGrokImportProbeScheduler(
grokImportProbeConcurrency,
grokImportProbeTimeout,
)
func newGrokImportProbeScheduler(concurrency int, timeout time.Duration) *grokImportProbeScheduler {
if concurrency <= 0 {
concurrency = 1
}
if timeout <= 0 {
timeout = grokImportProbeTimeout
}
return &grokImportProbeScheduler{
concurrency: concurrency,
timeout: timeout,
}
}
func (s *grokImportProbeScheduler) schedule(prober grokUsageProber, account *service.Account) {
if s == nil || prober == nil || account == nil || account.ID <= 0 {
return
}
if account.Platform != service.PlatformGrok || account.Type != service.AccountTypeOAuth {
return
}
s.mu.Lock()
s.queue = append(s.queue, grokImportProbeTask{prober: prober, accountID: account.ID})
if s.workers < s.concurrency {
s.workers++
if s.workers > s.maxWorkers {
s.maxWorkers = s.workers
}
go s.worker()
}
s.mu.Unlock()
}
func (s *grokImportProbeScheduler) worker() {
for {
task, ok := s.nextTask()
if !ok {
return
}
s.run(task.prober, task.accountID)
}
}
func (s *grokImportProbeScheduler) nextTask() (grokImportProbeTask, bool) {
s.mu.Lock()
defer s.mu.Unlock()
if len(s.queue) == 0 {
s.workers--
return grokImportProbeTask{}, false
}
task := s.queue[0]
s.queue[0] = grokImportProbeTask{}
s.queue = s.queue[1:]
if len(s.queue) == 0 {
s.queue = nil
}
return task, true
}
func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64) {
defer func() {
if recovered := recover(); recovered != nil {
slog.Error(
"grok_import_active_probe_panic",
"account_id", accountID,
"recovery_type", panicType(recovered),
)
}
}()
// Queue time is intentionally excluded: every imported account is probed,
// while this timeout only bounds the actual upstream probe execution.
ctx, cancel := context.WithTimeout(context.Background(), s.timeout)
defer cancel()
result, err := prober.ProbeUsage(ctx, accountID)
if err != nil {
slog.Warn(
"grok_import_active_probe_failed",
"account_id", accountID,
"status", infraerrors.Code(err),
"reason", infraerrors.Reason(err),
)
return
}
if result == nil {
slog.Warn(
"grok_import_active_probe_failed",
"account_id", accountID,
"reason", "empty_result",
)
return
}
slog.Info(
"grok_import_active_probe_completed",
"account_id", accountID,
"model", result.Model,
"status", result.StatusCode,
"headers_observed", result.HeadersObserved,
)
}
func panicType(value any) string {
switch value.(type) {
case string:
return "string"
case error:
return "error"
default:
return "unknown"
}
}
func (h *AccountHandler) scheduleGrokImportProbe(account *service.Account) {
if h == nil {
return
}
defaultGrokImportProbeScheduler.schedule(h.grokImportProber, account)
}
func (h *GrokOAuthHandler) scheduleGrokImportProbe(account *service.Account) {
if h == nil {
return
}
defaultGrokImportProbeScheduler.schedule(h.importProber, account)
}
// ProvideAccountHandler injects the Grok active prober for production while
// keeping NewAccountHandler convenient for focused unit tests.
func ProvideAccountHandler(
adminService service.AdminService,
oauthService *service.OAuthService,
openaiOAuthService *service.OpenAIOAuthService,
geminiOAuthService *service.GeminiOAuthService,
antigravityOAuthService *service.AntigravityOAuthService,
rateLimitService *service.RateLimitService,
accountUsageService *service.AccountUsageService,
accountTestService *service.AccountTestService,
concurrencyService *service.ConcurrencyService,
crsSyncService *service.CRSSyncService,
sessionLimitCache service.SessionLimitCache,
rpmCache service.RPMCache,
tokenCacheInvalidator service.TokenCacheInvalidator,
grokQuotaService *service.GrokQuotaService,
) *AccountHandler {
handler := NewAccountHandler(
adminService,
oauthService,
openaiOAuthService,
geminiOAuthService,
antigravityOAuthService,
rateLimitService,
accountUsageService,
accountTestService,
concurrencyService,
crsSyncService,
sessionLimitCache,
rpmCache,
tokenCacheInvalidator,
)
handler.grokImportProber = grokQuotaService
return handler
}
@@ -0,0 +1,119 @@
//go:build unit
package admin
import (
"context"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type grokImportAdminService struct {
*stubAdminService
mu sync.Mutex
nextID int64
}
func newGrokImportAdminService() *grokImportAdminService {
return &grokImportAdminService{
stubAdminService: newStubAdminService(),
nextID: 500,
}
}
func (s *grokImportAdminService) CreateAccount(_ context.Context, input *service.CreateAccountInput) (*service.Account, error) {
s.mu.Lock()
s.nextID++
id := s.nextID
s.mu.Unlock()
return &service.Account{
ID: id,
Name: input.Name,
Platform: input.Platform,
Type: input.Type,
Credentials: input.Credentials,
Extra: input.Extra,
ProxyID: input.ProxyID,
Concurrency: input.Concurrency,
Status: service.StatusActive,
Schedulable: true,
CreatedAt: time.Now(),
UpdatedAt: time.Now(),
}, nil
}
type grokImportOAuthClientStub struct{}
func (grokImportOAuthClientStub) ExchangeCode(context.Context, string, string, string, string, string) (*xai.TokenResponse, error) {
return &xai.TokenResponse{AccessToken: "access-token", RefreshToken: "refresh-token", ExpiresIn: 3600}, nil
}
func (grokImportOAuthClientStub) RefreshToken(context.Context, string, string, string) (*xai.TokenResponse, error) {
return &xai.TokenResponse{AccessToken: "access-token", RefreshToken: "refresh-token", ExpiresIn: 3600}, nil
}
func (grokImportOAuthClientStub) ConvertSSOToBuild(context.Context, string, string) (*xai.TokenResponse, error) {
return &xai.TokenResponse{AccessToken: "access-token", RefreshToken: "refresh-token", ExpiresIn: 3600}, nil
}
func TestGrokSSOBatchImportKeepsCreatedAccountsWhenOneAutomaticProbeFails(t *testing.T) {
gin.SetMode(gin.TestMode)
adminService := newGrokImportAdminService()
oauthService := service.NewGrokOAuthService(nil, grokImportOAuthClientStub{})
defer oauthService.Stop()
prober := newGrokImportProbeStub(3)
prober.failures[502] = infraerrors.New(502, "GROK_TEST_PROBE_FAILED", "sensitive-upstream-body")
handler := NewGrokOAuthHandler(oauthService, adminService, nil)
handler.importProber = prober
router := gin.New()
router.POST("/api/v1/admin/grok/sso-to-oauth", handler.CreateAccountsFromSSO)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(
http.MethodPost,
"/api/v1/admin/grok/sso-to-oauth",
strings.NewReader(`{"sso_tokens":["sso-one","sso-two","sso-three"]}`),
)
request.Header.Set("Content-Type", "application/json")
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
require.Contains(t, recorder.Body.String(), `"created"`)
require.NotContains(t, recorder.Body.String(), `GROK_TEST_PROBE_FAILED`)
for i := 0; i < 3; i++ {
awaitGrokProbeSignal(t, prober.done)
}
calls, _, _ := prober.snapshot()
require.Equal(t, map[int64]int{501: 1, 502: 1, 503: 1}, calls)
}
func TestAccountCreateWithoutAutomaticGrokProbeServiceStillSucceeds(t *testing.T) {
gin.SetMode(gin.TestMode)
handler := NewAccountHandler(
newGrokImportAdminService(),
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
)
router := gin.New()
router.POST("/api/v1/admin/accounts", handler.Create)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(
http.MethodPost,
"/api/v1/admin/accounts",
strings.NewReader(`{"name":"grok-rt","platform":"grok","type":"oauth","credentials":{"refresh_token":"secret"}}`),
)
request.Header.Set("Content-Type", "application/json")
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
}
@@ -0,0 +1,232 @@
//go:build unit
package admin
import (
"bytes"
"context"
"log/slog"
"sync"
"testing"
"time"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
type grokImportProbeStub struct {
mu sync.Mutex
calls map[int64]int
failures map[int64]error
active int
maxActive int
deadlineSeen bool
block <-chan struct{}
started chan int64
done chan int64
}
func newGrokImportProbeStub(buffer int) *grokImportProbeStub {
return &grokImportProbeStub{
calls: make(map[int64]int),
failures: make(map[int64]error),
started: make(chan int64, buffer),
done: make(chan int64, buffer),
}
}
func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) {
_, deadlineSeen := ctx.Deadline()
s.mu.Lock()
s.calls[accountID]++
s.active++
if s.active > s.maxActive {
s.maxActive = s.active
}
s.deadlineSeen = s.deadlineSeen || deadlineSeen
s.mu.Unlock()
s.started <- accountID
var ctxErr error
if s.block != nil {
select {
case <-s.block:
case <-ctx.Done():
ctxErr = ctx.Err()
}
}
s.mu.Lock()
s.active--
failure := s.failures[accountID]
s.mu.Unlock()
s.done <- accountID
if ctxErr != nil {
return nil, ctxErr
}
if failure != nil {
return nil, failure
}
return &service.GrokQuotaProbeResult{
Source: "active_probe",
Model: "grok-4.5",
StatusCode: 200,
ResetSupported: false,
}, nil
}
func (s *grokImportProbeStub) snapshot() (map[int64]int, int, bool) {
s.mu.Lock()
defer s.mu.Unlock()
calls := make(map[int64]int, len(s.calls))
for id, count := range s.calls {
calls[id] = count
}
return calls, s.maxActive, s.deadlineSeen
}
type grokImportProbeSchedulerTestSnapshot struct {
queued int
workers int
maxWorkers int
}
func snapshotGrokImportProbeScheduler(s *grokImportProbeScheduler) grokImportProbeSchedulerTestSnapshot {
if s == nil {
return grokImportProbeSchedulerTestSnapshot{}
}
s.mu.Lock()
defer s.mu.Unlock()
return grokImportProbeSchedulerTestSnapshot{
queued: len(s.queue),
workers: s.workers,
maxWorkers: s.maxWorkers,
}
}
func newGrokOAuthImportAccount(id int64) *service.Account {
return &service.Account{
ID: id,
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
}
}
func awaitGrokProbeSignal(t *testing.T, signals <-chan int64) int64 {
t.Helper()
select {
case id := <-signals:
return id
case <-time.After(time.Second):
t.Fatal("timed out waiting for Grok import probe")
return 0
}
}
func TestGrokImportProbeSchedulerProbesSingleAccountOnce(t *testing.T) {
scheduler := newGrokImportProbeScheduler(1, time.Second)
prober := newGrokImportProbeStub(1)
scheduler.schedule(prober, newGrokOAuthImportAccount(101))
require.Equal(t, int64(101), awaitGrokProbeSignal(t, prober.done))
calls, maxActive, deadlineSeen := prober.snapshot()
require.Equal(t, map[int64]int{101: 1}, calls)
require.Equal(t, 1, maxActive)
require.True(t, deadlineSeen)
require.Eventually(t, func() bool {
snapshot := snapshotGrokImportProbeScheduler(scheduler)
return snapshot.queued == 0 && snapshot.workers == 0
}, time.Second, 10*time.Millisecond)
}
func TestGrokImportProbeSchedulerQueuesBatchWithoutPerTaskGoroutines(t *testing.T) {
const taskCount = 100
release := make(chan struct{})
scheduler := newGrokImportProbeScheduler(3, time.Second)
prober := newGrokImportProbeStub(taskCount)
prober.block = release
prober.failures[150] = infraerrors.New(502, "GROK_TEST_PROBE_FAILED", "sensitive-upstream-body")
for id := int64(101); id < 101+taskCount; id++ {
scheduler.schedule(prober, newGrokOAuthImportAccount(id))
}
for i := 0; i < 3; i++ {
awaitGrokProbeSignal(t, prober.started)
}
snapshot := snapshotGrokImportProbeScheduler(scheduler)
require.Equal(t, 97, snapshot.queued)
require.Equal(t, 3, snapshot.workers)
require.Equal(t, 3, snapshot.maxWorkers)
select {
case id := <-prober.started:
t.Fatalf("probe %d started before a concurrency slot was released", id)
case <-time.After(75 * time.Millisecond):
}
close(release)
for i := 0; i < taskCount; i++ {
awaitGrokProbeSignal(t, prober.done)
}
calls, maxActive, _ := prober.snapshot()
require.Len(t, calls, taskCount)
for id := int64(101); id < 101+taskCount; id++ {
require.Equal(t, 1, calls[id])
}
require.Equal(t, 3, maxActive)
require.Eventually(t, func() bool {
snapshot = snapshotGrokImportProbeScheduler(scheduler)
return snapshot.queued == 0 && snapshot.workers == 0
}, time.Second, 10*time.Millisecond)
require.Equal(t, 3, snapshot.maxWorkers)
}
func TestGrokImportProbeSchedulerTimeoutCancelsProbe(t *testing.T) {
neverRelease := make(chan struct{})
scheduler := newGrokImportProbeScheduler(1, 20*time.Millisecond)
prober := newGrokImportProbeStub(1)
prober.block = neverRelease
scheduler.schedule(prober, newGrokOAuthImportAccount(201))
require.Equal(t, int64(201), awaitGrokProbeSignal(t, prober.done))
calls, _, _ := prober.snapshot()
require.Equal(t, 1, calls[201])
}
func TestGrokImportProbeSchedulerSkipsMissingServiceAndNonGrokAccounts(t *testing.T) {
scheduler := newGrokImportProbeScheduler(1, time.Second)
prober := newGrokImportProbeStub(1)
scheduler.schedule(nil, newGrokOAuthImportAccount(301))
scheduler.schedule(prober, &service.Account{ID: 302, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth})
scheduler.schedule(prober, &service.Account{ID: 303, Platform: service.PlatformGrok, Type: service.AccountTypeAPIKey})
select {
case id := <-prober.started:
t.Fatalf("unexpected probe for account %d", id)
case <-time.After(50 * time.Millisecond):
}
calls, _, _ := prober.snapshot()
require.Empty(t, calls)
}
func TestGrokImportProbeFailureLogDoesNotIncludeErrorMessage(t *testing.T) {
var logs bytes.Buffer
previousLogger := slog.Default()
slog.SetDefault(slog.New(slog.NewTextHandler(&logs, nil)))
defer slog.SetDefault(previousLogger)
scheduler := newGrokImportProbeScheduler(1, time.Second)
prober := newGrokImportProbeStub(1)
prober.failures[401] = infraerrors.New(502, "GROK_TEST_PROBE_FAILED", "refresh-token-secret")
scheduler.schedule(prober, newGrokOAuthImportAccount(401))
awaitGrokProbeSignal(t, prober.done)
require.Eventually(t, func() bool {
return bytes.Contains(logs.Bytes(), []byte("grok_import_active_probe_failed"))
}, time.Second, 10*time.Millisecond)
require.Contains(t, logs.String(), "GROK_TEST_PROBE_FAILED")
require.NotContains(t, logs.String(), "refresh-token-secret")
}
@@ -1,20 +1,28 @@
package admin
import (
"context"
"fmt"
"log/slog"
"strconv"
"strings"
"sync"
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
const grokSSOImportConcurrency = 3
type GrokOAuthHandler struct {
grokOAuthService *service.GrokOAuthService
adminService service.AdminService
quotaService *service.GrokQuotaService
importProber grokUsageProber
}
func NewGrokOAuthHandler(
@@ -26,6 +34,7 @@ func NewGrokOAuthHandler(
grokOAuthService: grokOAuthService,
adminService: adminService,
quotaService: quotaService,
importProber: quotaService,
}
}
@@ -202,9 +211,243 @@ func (h *GrokOAuthHandler) CreateAccountFromOAuth(c *gin.Context) {
response.ErrorFrom(c, err)
return
}
h.scheduleGrokImportProbe(account)
response.Success(c, dto.AccountFromService(account))
}
type GrokSSOToOAuthRequest struct {
SSOTokens []string `json:"sso_tokens"`
SSOToken string `json:"sso_token"`
Name string `json:"name"`
Notes *string `json:"notes"`
ProxyID *int64 `json:"proxy_id"`
GroupIDs []int64 `json:"group_ids"`
Credentials map[string]any `json:"credentials"`
Extra map[string]any `json:"extra"`
Concurrency int `json:"concurrency"`
LoadFactor *int `json:"load_factor"`
Priority int `json:"priority"`
RateMultiplier *float64 `json:"rate_multiplier"`
ExpiresAt *int64 `json:"expires_at"`
AutoPauseOnExpired *bool `json:"auto_pause_on_expired"`
}
type GrokSSOToOAuthItemResult struct {
Index int `json:"index"`
Name string `json:"name,omitempty"`
Email string `json:"email,omitempty"`
Account *dto.Account `json:"account,omitempty"`
Error string `json:"error,omitempty"`
}
type GrokSSOToOAuthResponse struct {
Created []GrokSSOToOAuthItemResult `json:"created"`
Failed []GrokSSOToOAuthItemResult `json:"failed"`
}
type grokSSOImportJob struct {
index int
token string
}
type grokSSOImportWorkerResult struct {
created bool
item GrokSSOToOAuthItemResult
}
func (h *GrokOAuthHandler) CreateAccountsFromSSO(c *gin.Context) {
var req GrokSSOToOAuthRequest
if err := c.ShouldBindJSON(&req); err != nil {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
tokens := normalizeSSOImportTokens(req.SSOTokens, req.SSOToken)
if len(tokens) == 0 {
response.BadRequest(c, "sso_tokens is required")
return
}
ctx := c.Request.Context()
workerCount := grokSSOImportConcurrency
if len(tokens) < workerCount {
workerCount = len(tokens)
}
jobs := make(chan grokSSOImportJob)
items := make([]grokSSOImportWorkerResult, len(tokens))
var wg sync.WaitGroup
for i := 0; i < workerCount; i++ {
wg.Add(1)
go func() {
defer wg.Done()
for job := range jobs {
items[job.index] = h.safeCreateAccountFromSSOToken(ctx, req, job.token, job.index+1, len(tokens))
}
}()
}
for i, token := range tokens {
jobs <- grokSSOImportJob{index: i, token: token}
}
close(jobs)
wg.Wait()
result := GrokSSOToOAuthResponse{
Created: make([]GrokSSOToOAuthItemResult, 0, len(tokens)),
Failed: make([]GrokSSOToOAuthItemResult, 0),
}
for _, item := range items {
if item.created {
result.Created = append(result.Created, item.item)
} else {
result.Failed = append(result.Failed, item.item)
}
}
response.Success(c, result)
}
func (h *GrokOAuthHandler) safeCreateAccountFromSSOToken(ctx context.Context, req GrokSSOToOAuthRequest, token string, index, total int) (result grokSSOImportWorkerResult) {
defer func() {
if recovered := recover(); recovered != nil {
slog.Error("grok_sso_import_worker_panic", "index", index, "recover", recovered)
result = grokSSOImportWorkerResult{
item: GrokSSOToOAuthItemResult{
Index: index,
Error: fmt.Sprintf("internal worker panic: %v", recovered),
},
}
}
}()
return h.createAccountFromSSOToken(ctx, req, token, index, total)
}
func (h *GrokOAuthHandler) createAccountFromSSOToken(ctx context.Context, req GrokSSOToOAuthRequest, token string, index, total int) grokSSOImportWorkerResult {
tokenInfo, err := h.grokOAuthService.ConvertFromSSO(ctx, token, req.ProxyID)
if err != nil {
return grokSSOImportWorkerResult{item: GrokSSOToOAuthItemResult{Index: index, Error: grokSSOImportErrorMessage(err)}}
}
credentials := h.grokOAuthService.BuildAccountCredentials(tokenInfo)
credentials = service.MergeCredentials(cloneGrokSSOMap(req.Credentials), credentials)
name := grokSSOImportAccountName(req.Name, tokenInfo, index, total)
expiresAt, autoPauseOnExpired := grokSSOImportExpiry(req.ExpiresAt, req.AutoPauseOnExpired, tokenInfo)
account, err := h.adminService.CreateAccount(ctx, &service.CreateAccountInput{
Name: name,
Notes: req.Notes,
Platform: service.PlatformGrok,
Type: service.AccountTypeOAuth,
Credentials: credentials,
Extra: cloneGrokSSOMap(req.Extra),
ProxyID: req.ProxyID,
Concurrency: req.Concurrency,
LoadFactor: req.LoadFactor,
Priority: req.Priority,
RateMultiplier: req.RateMultiplier,
GroupIDs: append([]int64(nil), req.GroupIDs...),
ExpiresAt: expiresAt,
AutoPauseOnExpired: autoPauseOnExpired,
})
if err != nil {
return grokSSOImportWorkerResult{item: GrokSSOToOAuthItemResult{Index: index, Name: name, Email: tokenInfo.Email, Error: grokSSOImportErrorMessage(err)}}
}
h.scheduleGrokImportProbe(account)
return grokSSOImportWorkerResult{
created: true,
item: GrokSSOToOAuthItemResult{
Index: index,
Name: name,
Email: tokenInfo.Email,
Account: dto.AccountFromService(account),
},
}
}
func grokSSOImportExpiry(requestExpiresAt *int64, requestAutoPause *bool, tokenInfo *service.GrokTokenInfo) (*int64, *bool) {
if tokenInfo == nil || strings.TrimSpace(tokenInfo.RefreshToken) != "" || tokenInfo.ExpiresAt <= 0 {
return requestExpiresAt, requestAutoPause
}
expiresAt := tokenInfo.ExpiresAt
if requestExpiresAt != nil && *requestExpiresAt > 0 && *requestExpiresAt < expiresAt {
expiresAt = *requestExpiresAt
}
autoPause := true
return &expiresAt, &autoPause
}
func cloneGrokSSOMap(source map[string]any) map[string]any {
if source == nil {
return nil
}
clone := make(map[string]any, len(source))
for key, value := range source {
clone[key] = cloneGrokSSOValue(value)
}
return clone
}
func cloneGrokSSOValue(value any) any {
switch v := value.(type) {
case map[string]any:
return cloneGrokSSOMap(v)
case []any:
clone := make([]any, len(v))
for i, item := range v {
clone[i] = cloneGrokSSOValue(item)
}
return clone
default:
return value
}
}
func normalizeSSOImportTokens(tokens []string, single string) []string {
items := make([]string, 0, len(tokens)+1)
if strings.TrimSpace(single) != "" {
items = append(items, single)
}
items = append(items, tokens...)
seen := make(map[string]struct{}, len(items))
result := make([]string, 0, len(items))
for _, item := range items {
parts := strings.Split(strings.NewReplacer(",", "\n", "\r", "\n").Replace(item), "\n")
for _, token := range parts {
if token = xai.NormalizeSSOToken(token); token == "" {
continue
}
if _, ok := seen[token]; ok {
continue
}
seen[token] = struct{}{}
result = append(result, token)
}
}
return result
}
func grokSSOImportAccountName(base string, tokenInfo *service.GrokTokenInfo, index, total int) string {
base = strings.TrimSpace(base)
if base == "" && tokenInfo != nil {
base = strings.TrimSpace(tokenInfo.Email)
}
if base == "" {
base = "Grok OAuth Account"
}
if total > 1 {
return base + " #" + strconv.Itoa(index)
}
return base
}
func grokSSOImportErrorMessage(err error) string {
status := infraerrors.FromError(err)
if status == nil {
return ""
}
if status.Reason != "" {
return status.Reason + ": " + status.Message
}
return status.Message
}
func (h *GrokOAuthHandler) QueryQuota(c *gin.Context) {
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil {
@@ -215,7 +458,7 @@ func (h *GrokOAuthHandler) QueryQuota(c *gin.Context) {
response.BadRequest(c, "grok quota service is not enabled")
return
}
result, err := h.quotaService.ProbeUsage(c.Request.Context(), accountID)
result, err := h.quotaService.QueryQuota(c.Request.Context(), accountID)
if err != nil {
response.ErrorFrom(c, err)
return
@@ -8,6 +8,7 @@ import (
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"time"
@@ -41,17 +42,35 @@ func (r *grokQuotaHandlerAccountRepo) UpdateExtra(_ context.Context, id int64, u
}
type grokQuotaHandlerUpstream struct {
resp *http.Response
lastReq *http.Request
lastBody []byte
mu sync.Mutex
requests []*http.Request
bodies [][]byte
}
func (u *grokQuotaHandlerUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
u.lastReq = req
var body []byte
if req.Body != nil {
u.lastBody, _ = io.ReadAll(req.Body)
body, _ = io.ReadAll(req.Body)
}
return u.resp, nil
u.mu.Lock()
u.requests = append(u.requests, req)
u.bodies = append(u.bodies, body)
u.mu.Unlock()
if req.URL.Path == "/v1/responses" {
return &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"X-Ratelimit-Limit-Requests": []string{"10"},
"X-Ratelimit-Remaining-Requests": []string{"8"},
},
Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
}, nil
}
payload := `{"config":{"billingPeriodStart":"2026-07-01T00:00:00Z","billingPeriodEnd":"2026-08-01T00:00:00Z"}}`
if req.URL.RawQuery == "format=credits" {
payload = `{"config":{"currentPeriod":{"type":"WEEKLY","start":"2026-07-09T03:25:00Z","end":"2026-07-16T03:25:00Z"}}}`
}
return &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(payload))}, nil
}
func (u *grokQuotaHandlerUpstream) DoWithTLS(
@@ -77,14 +96,7 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) {
"expires_at": time.Now().Add(time.Hour).UTC().Format(time.RFC3339),
},
}}
upstream := &grokQuotaHandlerUpstream{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{
"X-Ratelimit-Limit-Requests": []string{"10"},
"X-Ratelimit-Remaining-Requests": []string{"8"},
},
Body: io.NopCloser(strings.NewReader(`{"id":"resp_probe"}`)),
}}
upstream := &grokQuotaHandlerUpstream{}
quotaService := service.NewGrokQuotaService(repo, nil, service.NewGrokTokenProvider(repo, nil), upstream)
handler := NewGrokOAuthHandler(nil, nil, quotaService)
@@ -95,12 +107,23 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) {
router.ServeHTTP(rec, req)
require.Equal(t, http.StatusOK, rec.Code)
require.Contains(t, rec.Body.String(), `"source":"active_probe"`)
require.Contains(t, rec.Body.String(), `"source":"hybrid_probe"`)
require.Contains(t, rec.Body.String(), `"billing":`)
require.Contains(t, rec.Body.String(), `"snapshot":`)
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, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
require.Contains(t, string(upstream.lastBody), `"store":false`)
upstream.mu.Lock()
requests := append([]*http.Request(nil), upstream.requests...)
bodies := append([][]byte(nil), upstream.bodies...)
upstream.mu.Unlock()
require.Len(t, requests, 3)
for i, upstreamReq := range requests {
require.Equal(t, "Bearer access-token", upstreamReq.Header.Get("Authorization"))
if upstreamReq.URL.String() == xai.DefaultCLIBaseURL+"/responses" {
require.Contains(t, string(bodies[i]), `"model":"grok-4.5"`)
require.Contains(t, string(bodies[i]), `"store":false`)
}
}
require.NotNil(t, repo.updates[42])
}
@@ -145,3 +168,51 @@ func TestGrokOAuthHandlerRuntimeSanityDoesNotExposeSecrets(t *testing.T) {
require.NotContains(t, rec.Body.String(), "secret")
require.NotContains(t, rec.Body.String(), "client-secret-like-value")
}
func TestGrokSSOImportExpiryUsesTokenExpiryWithoutRefreshToken(t *testing.T) {
tokenExpiry := time.Now().Add(6 * time.Hour).Unix()
expiresAt, autoPause := grokSSOImportExpiry(nil, nil, &service.GrokTokenInfo{
ExpiresAt: tokenExpiry,
})
require.NotNil(t, expiresAt)
require.Equal(t, tokenExpiry, *expiresAt)
require.NotNil(t, autoPause)
require.True(t, *autoPause)
}
func TestGrokSSOImportExpiryUsesEarlierRequestedExpiryWithoutRefreshToken(t *testing.T) {
requestedExpiry := time.Now().Add(2 * time.Hour).Unix()
tokenExpiry := time.Now().Add(6 * time.Hour).Unix()
requestedAutoPause := false
expiresAt, autoPause := grokSSOImportExpiry(&requestedExpiry, &requestedAutoPause, &service.GrokTokenInfo{
ExpiresAt: tokenExpiry,
})
require.NotNil(t, expiresAt)
require.Equal(t, requestedExpiry, *expiresAt)
require.NotNil(t, autoPause)
require.True(t, *autoPause)
}
func TestGrokSSOImportExpiryPreservesRequestSettingsWithRefreshToken(t *testing.T) {
requestedExpiry := time.Now().Add(2 * time.Hour).Unix()
requestedAutoPause := false
expiresAt, autoPause := grokSSOImportExpiry(&requestedExpiry, &requestedAutoPause, &service.GrokTokenInfo{
RefreshToken: "refresh-token",
ExpiresAt: time.Now().Add(6 * time.Hour).Unix(),
})
require.Same(t, &requestedExpiry, expiresAt)
require.Same(t, &requestedAutoPause, autoPause)
}
func TestGrokSSOImportWorkerRecoversPanic(t *testing.T) {
h := &GrokOAuthHandler{}
result := h.safeCreateAccountFromSSOToken(context.Background(), GrokSSOToOAuthRequest{}, "token", 2, 3)
// Without a service, createAccountFromSSOToken would panic on nil service access.
// Recovery must convert that into a failed item and keep the worker alive.
require.False(t, result.created)
require.Equal(t, 2, result.item.Index)
require.Contains(t, result.item.Error, "internal worker panic")
}
@@ -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,
@@ -304,6 +304,10 @@ func (h *OpenAIOAuthHandler) CreateAccountFromCodexPAT(c *gin.Context) {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
if err := service.ValidateOpenAILongContextBillingExtra(service.PlatformOpenAI, req.Extra); err != nil {
response.ErrorFrom(c, err)
return
}
if req.Concurrency != nil && *req.Concurrency < 0 {
response.BadRequest(c, "concurrency must be >= 0")
return
@@ -15,6 +15,7 @@ import (
type opsSystemLogCleanupRequest struct {
StartTime string `json:"start_time"`
EndTime string `json:"end_time"`
Host string `json:"host"`
Level string `json:"level"`
Component string `json:"component"`
@@ -56,6 +57,7 @@ func (h *OpsHandler) ListSystemLogs(c *gin.Context) {
PageSize: pageSize,
StartTime: &start,
EndTime: &end,
Host: strings.TrimSpace(c.Query("host")),
Level: strings.TrimSpace(c.Query("level")),
Component: strings.TrimSpace(c.Query("component")),
RequestID: strings.TrimSpace(c.Query("request_id")),
@@ -153,6 +155,7 @@ func (h *OpsHandler) CleanupSystemLogs(c *gin.Context) {
filter := &service.OpsSystemLogCleanupFilter{
StartTime: start,
EndTime: end,
Host: strings.TrimSpace(req.Host),
Level: strings.TrimSpace(req.Level),
Component: strings.TrimSpace(req.Component),
RequestID: strings.TrimSpace(req.RequestID),
@@ -2,6 +2,7 @@ package admin
import (
"bytes"
"context"
"encoding/json"
"net/http"
"net/http/httptest"
@@ -19,6 +20,26 @@ type responseEnvelope struct {
Data json.RawMessage `json:"data"`
}
type opsSystemLogCaptureRepo struct {
service.OpsRepository
listFilter *service.OpsSystemLogFilter
cleanupFilter *service.OpsSystemLogCleanupFilter
}
func (r *opsSystemLogCaptureRepo) ListSystemLogs(_ context.Context, filter *service.OpsSystemLogFilter) (*service.OpsSystemLogList, error) {
r.listFilter = filter
return &service.OpsSystemLogList{Logs: []*service.OpsSystemLog{}, Page: filter.Page, PageSize: filter.PageSize}, nil
}
func (r *opsSystemLogCaptureRepo) DeleteSystemLogs(_ context.Context, filter *service.OpsSystemLogCleanupFilter) (int64, error) {
r.cleanupFilter = filter
return 1, nil
}
func (r *opsSystemLogCaptureRepo) InsertSystemLogCleanupAudit(_ context.Context, _ *service.OpsSystemLogCleanupAudit) error {
return nil
}
func newOpsSystemLogTestRouter(handler *OpsHandler, withUser bool) *gin.Engine {
gin.SetMode(gin.TestMode)
r := gin.New()
@@ -121,6 +142,23 @@ func TestOpsSystemLogHandler_ListSuccess(t *testing.T) {
}
}
func TestOpsSystemLogHandler_ListAcceptsHost(t *testing.T) {
repo := &opsSystemLogCaptureRepo{}
svc := service.NewOpsService(repo, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
h := NewOpsHandler(svc)
r := newOpsSystemLogTestRouter(h, false)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/logs?host=api-node-1", nil)
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status=%d, want 200", w.Code)
}
if repo.listFilter == nil || repo.listFilter.Host != "api-node-1" {
t.Fatalf("host filter = %+v, want api-node-1", repo.listFilter)
}
}
func TestOpsSystemLogHandler_CleanupUnauthorized(t *testing.T) {
svc := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
h := NewOpsHandler(svc)
@@ -205,6 +243,24 @@ func TestOpsSystemLogHandler_CleanupAcceptsAPIKeyID(t *testing.T) {
}
}
func TestOpsSystemLogHandler_CleanupAcceptsHost(t *testing.T) {
repo := &opsSystemLogCaptureRepo{}
svc := service.NewOpsService(repo, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
h := NewOpsHandler(svc)
r := newOpsSystemLogTestRouter(h, true)
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/logs/cleanup", bytes.NewBufferString(`{"host":"api-node-1"}`))
req.Header.Set("Content-Type", "application/json")
r.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status=%d, want 200", w.Code)
}
if repo.cleanupFilter == nil || repo.cleanupFilter.Host != "api-node-1" {
t.Fatalf("host filter = %+v, want api-node-1", repo.cleanupFilter)
}
}
func TestOpsSystemLogHandler_CleanupInvalidAPIKeyID(t *testing.T) {
svc := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
h := NewOpsHandler(svc)
@@ -16,6 +16,7 @@ import (
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
@@ -323,7 +324,7 @@ func (h *OpsHandler) QPSWSHandler(c *gin.Context) {
// If realtime monitoring is disabled, prefer a successful WS upgrade followed by a clean close
// with a deterministic close code. This prevents clients from spinning on 404/1006 reconnect loops.
if !h.opsService.IsRealtimeMonitoringEnabled(c.Request.Context()) {
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
conn, err := upgrader.Upgrade(c.Writer, c.Request, servermiddleware.ServerTimingResponseHeader(c))
if err != nil {
c.JSON(http.StatusNotFound, gin.H{"error": "ops realtime monitoring is disabled"})
return
@@ -358,7 +359,7 @@ func (h *OpsHandler) QPSWSHandler(c *gin.Context) {
defer releaseOpsWSIPSlot(clientIP)
}
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
conn, err := upgrader.Upgrade(c.Writer, c.Request, servermiddleware.ServerTimingResponseHeader(c))
if err != nil {
logger.LegacyPrintf("handler.admin.ops_ws", "[OpsWS] upgrade failed: %v", err)
return
+50 -48
View File
@@ -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,
@@ -598,54 +599,55 @@ func usageLogFromServiceUser(l *service.UsageLog) UsageLog {
requestedModel = l.Model
}
return UsageLog{
ID: l.ID,
UserID: l.UserID,
APIKeyID: l.APIKeyID,
AccountID: l.AccountID,
RequestID: l.RequestID,
Model: requestedModel,
ServiceTier: l.ServiceTier,
ReasoningEffort: l.ReasoningEffort,
InboundEndpoint: l.InboundEndpoint,
GroupID: l.GroupID,
SubscriptionID: l.SubscriptionID,
InputTokens: l.InputTokens,
OutputTokens: l.OutputTokens,
CacheCreationTokens: l.CacheCreationTokens,
CacheReadTokens: l.CacheReadTokens,
CacheCreation5mTokens: l.CacheCreation5mTokens,
CacheCreation1hTokens: l.CacheCreation1hTokens,
InputCost: l.InputCost,
OutputCost: l.OutputCost,
CacheCreationCost: l.CacheCreationCost,
CacheReadCost: l.CacheReadCost,
TotalCost: l.TotalCost,
ActualCost: l.ActualCost,
RateMultiplier: l.RateMultiplier,
BillingType: l.BillingType,
RequestType: requestType.String(),
Stream: stream,
OpenAIWSMode: openAIWSMode,
DurationMs: l.DurationMs,
FirstTokenMs: l.FirstTokenMs,
ImageCount: l.ImageCount,
ImageSize: l.ImageSize,
ImageInputSize: l.ImageInputSize,
ImageOutputSize: l.ImageOutputSize,
ImageOutputTokens: l.ImageOutputTokens,
ImageOutputCost: l.ImageOutputCost,
ImageSizeSource: l.ImageSizeSource,
ImageSizeBreakdown: l.ImageSizeBreakdown,
MediaType: l.MediaType,
UserAgent: l.UserAgent,
IPAddress: l.IPAddress,
CacheTTLOverridden: l.CacheTTLOverridden,
BillingMode: l.BillingMode,
CreatedAt: l.CreatedAt,
User: UserFromServiceShallow(l.User),
APIKey: APIKeyFromService(l.APIKey),
Group: GroupFromServiceShallow(l.Group),
Subscription: UserSubscriptionFromService(l.Subscription),
ID: l.ID,
UserID: l.UserID,
APIKeyID: l.APIKeyID,
AccountID: l.AccountID,
RequestID: l.RequestID,
Model: requestedModel,
ServiceTier: l.ServiceTier,
ReasoningEffort: l.ReasoningEffort,
InboundEndpoint: l.InboundEndpoint,
GroupID: l.GroupID,
SubscriptionID: l.SubscriptionID,
InputTokens: l.InputTokens,
OutputTokens: l.OutputTokens,
CacheCreationTokens: l.CacheCreationTokens,
CacheReadTokens: l.CacheReadTokens,
CacheCreation5mTokens: l.CacheCreation5mTokens,
CacheCreation1hTokens: l.CacheCreation1hTokens,
InputCost: l.InputCost,
OutputCost: l.OutputCost,
CacheCreationCost: l.CacheCreationCost,
CacheReadCost: l.CacheReadCost,
TotalCost: l.TotalCost,
ActualCost: l.ActualCost,
RateMultiplier: l.RateMultiplier,
LongContextBillingApplied: l.LongContextBillingApplied,
BillingType: l.BillingType,
RequestType: requestType.String(),
Stream: stream,
OpenAIWSMode: openAIWSMode,
DurationMs: l.DurationMs,
FirstTokenMs: l.FirstTokenMs,
ImageCount: l.ImageCount,
ImageSize: l.ImageSize,
ImageInputSize: l.ImageInputSize,
ImageOutputSize: l.ImageOutputSize,
ImageOutputTokens: l.ImageOutputTokens,
ImageOutputCost: l.ImageOutputCost,
ImageSizeSource: l.ImageSizeSource,
ImageSizeBreakdown: l.ImageSizeBreakdown,
MediaType: l.MediaType,
UserAgent: l.UserAgent,
IPAddress: l.IPAddress,
CacheTTLOverridden: l.CacheTTLOverridden,
BillingMode: l.BillingMode,
CreatedAt: l.CreatedAt,
User: UserFromServiceShallow(l.User),
APIKey: APIKeyFromService(l.APIKey),
Group: GroupFromServiceShallow(l.Group),
Subscription: UserSubscriptionFromService(l.Subscription),
}
}
+10 -7
View File
@@ -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"`
@@ -480,13 +482,14 @@ type UsageLog struct {
CacheCreation5mTokens int `json:"cache_creation_5m_tokens"`
CacheCreation1hTokens int `json:"cache_creation_1h_tokens"`
InputCost float64 `json:"input_cost"`
OutputCost float64 `json:"output_cost"`
CacheCreationCost float64 `json:"cache_creation_cost"`
CacheReadCost float64 `json:"cache_read_cost"`
TotalCost float64 `json:"total_cost"`
ActualCost float64 `json:"actual_cost"`
RateMultiplier float64 `json:"rate_multiplier"`
InputCost float64 `json:"input_cost"`
OutputCost float64 `json:"output_cost"`
CacheCreationCost float64 `json:"cache_creation_cost"`
CacheReadCost float64 `json:"cache_read_cost"`
TotalCost float64 `json:"total_cost"`
ActualCost float64 `json:"actual_cost"`
RateMultiplier float64 `json:"rate_multiplier"`
LongContextBillingApplied bool `json:"long_context_billing_applied"`
BillingType int8 `json:"billing_type"`
RequestType string `json:"request_type"`
+15 -3
View File
@@ -18,11 +18,14 @@ 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"
EndpointImagesEdits = "/v1/images/edits"
EndpointVideosGenerations = "/v1/videos/generations"
EndpointVideosEdits = "/v1/videos/edits"
EndpointVideosExtensions = "/v1/videos/extensions"
EndpointVideos = "/v1/videos"
EndpointGeminiModels = "/v1beta/models"
)
@@ -75,6 +78,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):
@@ -85,6 +90,10 @@ func NormalizeInboundEndpoint(path string) string {
return EndpointImagesEdits
case strings.Contains(path, EndpointVideosGenerations) || strings.Contains(path, "/videos/generations"):
return EndpointVideosGenerations
case strings.Contains(path, EndpointVideosEdits) || strings.Contains(path, "/videos/edits"):
return EndpointVideosEdits
case strings.Contains(path, EndpointVideosExtensions) || strings.Contains(path, "/videos/extensions"):
return EndpointVideosExtensions
case strings.Contains(path, EndpointVideos) || strings.Contains(path, "/videos/"):
return EndpointVideos
case strings.Contains(path, EndpointResponsesCompact) || isResponsesCompactAliasPath(path):
@@ -155,8 +164,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 +179,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 == EndpointVideosEdits || inbound == EndpointVideosExtensions || inbound == EndpointVideos {
return inbound
}
// OpenAI forwards everything to the Responses API.
+59
View File
@@ -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
// ──────────────────────────────────────────────────────────
+7 -4
View File
@@ -29,7 +29,8 @@ const (
)
const (
// maxSameAccountRetries 同账号重试次数上限(针对 RetryableOnSameAccount 错误)
// maxSameAccountRetries 同账号重试次数默认上限(针对 RetryableOnSameAccount 错误)。
// 生产调用方通常传入账号级配置 account.GetPoolModeRetryCount(),该常量仅作兜底/测试默认值。
maxSameAccountRetries = 3
// sameAccountRetryDelay 同账号重试间隔
sameAccountRetryDelay = 500 * time.Millisecond
@@ -67,6 +68,7 @@ func (s *FailoverState) HandleFailoverError(
gatewayService TempUnscheduler,
accountID int64,
platform string,
retryLimit int,
failoverErr *service.UpstreamFailoverError,
) FailoverAction {
s.LastFailoverErr = failoverErr
@@ -76,14 +78,15 @@ func (s *FailoverState) HandleFailoverError(
s.ForceCacheBilling = true
}
// 同账号重试:对 RetryableOnSameAccount 的临时性错误,先在同一账号上重试
if failoverErr.RetryableOnSameAccount && s.SameAccountRetryCount[accountID] < maxSameAccountRetries {
// 同账号重试:对 RetryableOnSameAccount 的临时性错误,先在同一账号上重试。
// 重试次数上限 retryLimit 由调用方传入(账号级 pool_mode_retry_count 配置)。
if failoverErr.RetryableOnSameAccount && s.SameAccountRetryCount[accountID] < retryLimit {
s.SameAccountRetryCount[accountID]++
logger.FromContext(ctx).Warn("gateway.failover_same_account_retry",
zap.Int64("account_id", accountID),
zap.Int("upstream_status", failoverErr.StatusCode),
zap.Int("same_account_retry_count", s.SameAccountRetryCount[accountID]),
zap.Int("same_account_retry_max", maxSameAccountRetries),
zap.Int("same_account_retry_max", retryLimit),
)
if !sleepWithContext(ctx, sameAccountRetryDelay) {
return FailoverCanceled
+87 -50
View File
@@ -133,7 +133,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, false, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SwitchCount)
@@ -150,7 +150,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
err := newTestFailoverErr(500, false, false)
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformAntigravity, err)
action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
@@ -166,7 +166,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
err := newTestFailoverErr(500, false, false)
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 200, service.PlatformAntigravity, err)
action := fs.HandleFailoverError(context.Background(), mock, 200, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
@@ -181,19 +181,19 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
// 第一次切换:0→1
err1 := newTestFailoverErr(500, false, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err1)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SwitchCount)
// 第二次切换:1→2
err2 := newTestFailoverErr(502, false, false)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", err2)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 2, fs.SwitchCount)
// 第三次已耗尽:SwitchCount(2) >= MaxSwitches(2)
err3 := newTestFailoverErr(503, false, false)
action = fs.HandleFailoverError(context.Background(), mock, 300, "openai", err3)
action = fs.HandleFailoverError(context.Background(), mock, 300, "openai", maxSameAccountRetries, err3)
require.Equal(t, FailoverExhausted, action)
require.Equal(t, 2, fs.SwitchCount, "耗尽时不应继续递增")
@@ -212,7 +212,7 @@ func TestHandleFailoverError_BasicSwitch(t *testing.T) {
fs := NewFailoverState(0, false)
err := newTestFailoverErr(500, false, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverExhausted, action)
require.Equal(t, 0, fs.SwitchCount)
require.Contains(t, fs.FailedAccountIDs, int64(100))
@@ -229,7 +229,7 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
fs := NewFailoverState(3, true) // hasBoundSession=true
err := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.True(t, fs.ForceCacheBilling)
})
@@ -238,7 +238,7 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, false, true) // ForceCacheBilling=true
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.True(t, fs.ForceCacheBilling)
})
@@ -247,7 +247,7 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.False(t, fs.ForceCacheBilling)
})
@@ -257,12 +257,12 @@ func TestHandleFailoverError_CacheBilling(t *testing.T) {
// 第一次:ForceCacheBilling=true → 设置
err1 := newTestFailoverErr(500, false, true)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err1)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
require.True(t, fs.ForceCacheBilling)
// 第二次:ForceCacheBilling=false → 仍然保持 true
err2 := newTestFailoverErr(502, false, false)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", err2)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
require.True(t, fs.ForceCacheBilling, "ForceCacheBilling 一旦设置不应被重置")
})
}
@@ -278,7 +278,7 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
err := newTestFailoverErr(400, true, false)
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
@@ -297,7 +297,7 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
err := newTestFailoverErr(400, true, false)
for i := 1; i <= maxSameAccountRetries; i++ {
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, i, fs.SameAccountRetryCount[100])
}
@@ -311,12 +311,12 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
err := newTestFailoverErr(400, true, false)
for i := 0; i < maxSameAccountRetries; i++ {
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
}
require.Equal(t, maxSameAccountRetries, fs.SameAccountRetryCount[100])
// 第 maxSameAccountRetries+1 次:重试耗尽,应切换账号
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SwitchCount)
require.Contains(t, fs.FailedAccountIDs, int64(100))
@@ -333,12 +333,12 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
err := newTestFailoverErr(400, true, false)
// 账号 100 第一次重试
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[100])
// 账号 200 第一次重试(独立计数)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", err)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[200])
require.Equal(t, 1, fs.SameAccountRetryCount[100], "账号 100 的计数不应受影响")
@@ -351,17 +351,54 @@ func TestHandleFailoverError_SameAccountRetry(t *testing.T) {
// 耗尽账号 100 的重试
for i := 0; i < maxSameAccountRetries; i++ {
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
}
// 第 maxSameAccountRetries+1 次: 重试耗尽 → 切换
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
// 再次遇到账号 100,计数仍为 maxSameAccountRetries,条件不满足 → 直接切换
action = fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action = fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Len(t, mock.calls, 2, "第二次耗尽也应调用 TempUnschedule")
})
t.Run("尊重账号级retryLimit_配置1次只重试1次", func(t *testing.T) {
// 回归测试:Anthropic 等路径此前硬编码同账号重试 3 次,忽略账号
// pool_mode_retry_count 配置。此处验证传入 retryLimit=1 时只重试 1 次即切换。
mock := &mockTempUnscheduler{}
fs := NewFailoverState(5, false)
err := newTestFailoverErr(403, true, false)
const retryLimit = 1
// 第 1 次:同账号重试
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryLimit, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[100])
require.Equal(t, 0, fs.SwitchCount, "首次重试不应切换账号")
require.Empty(t, mock.calls, "未耗尽前不应 TempUnschedule")
// 第 2 次:已达上限 1 → 不再同账号重试,直接切换 + TempUnschedule
action = fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryLimit, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[100], "重试计数不应超过 retryLimit")
require.Equal(t, 1, fs.SwitchCount, "重试耗尽应切换账号")
require.Contains(t, fs.FailedAccountIDs, int64(100))
require.Len(t, mock.calls, 1, "重试耗尽应触发 TempUnschedule")
})
t.Run("retryLimit为0时立即切换不重试", func(t *testing.T) {
// pool_mode_retry_count=0 表示关闭同账号重试(如 GPT Image 账号)。
mock := &mockTempUnscheduler{}
fs := NewFailoverState(5, false)
err := newTestFailoverErr(403, true, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", 0, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 0, fs.SameAccountRetryCount[100], "retryLimit=0 不应发生同账号重试")
require.Equal(t, 1, fs.SwitchCount, "应立即切换账号")
require.Len(t, mock.calls, 1, "应立即 TempUnschedule")
})
}
// ---------------------------------------------------------------------------
@@ -374,7 +411,7 @@ func TestHandleFailoverError_TempUnschedule(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, false, false) // RetryableOnSameAccount=false
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Empty(t, mock.calls)
})
@@ -384,10 +421,10 @@ func TestHandleFailoverError_TempUnschedule(t *testing.T) {
err := newTestFailoverErr(502, true, false)
for i := 0; i < maxSameAccountRetries; i++ {
fs.HandleFailoverError(context.Background(), mock, 42, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 42, "openai", maxSameAccountRetries, err)
}
// 再次触发时才会执行 TempUnschedule + 切换
fs.HandleFailoverError(context.Background(), mock, 42, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 42, "openai", maxSameAccountRetries, err)
require.Len(t, mock.calls, 1)
require.Equal(t, int64(42), mock.calls[0].accountID)
@@ -410,7 +447,7 @@ func TestHandleFailoverError_ContextCanceled(t *testing.T) {
cancel() // 立即取消
start := time.Now()
action := fs.HandleFailoverError(ctx, mock, 100, "openai", err)
action := fs.HandleFailoverError(ctx, mock, 100, "openai", maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverCanceled, action)
@@ -429,7 +466,7 @@ func TestHandleFailoverError_ContextCanceled(t *testing.T) {
cancel() // 立即取消
start := time.Now()
action := fs.HandleFailoverError(ctx, mock, 100, service.PlatformAntigravity, err)
action := fs.HandleFailoverError(ctx, mock, 100, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverCanceled, action)
@@ -446,10 +483,10 @@ func TestHandleFailoverError_FailedAccountIDs(t *testing.T) {
mock := &mockTempUnscheduler{}
fs := NewFailoverState(3, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(500, false, false))
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
require.Contains(t, fs.FailedAccountIDs, int64(100))
fs.HandleFailoverError(context.Background(), mock, 200, "openai", newTestFailoverErr(502, false, false))
fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, newTestFailoverErr(502, false, false))
require.Contains(t, fs.FailedAccountIDs, int64(200))
require.Len(t, fs.FailedAccountIDs, 2)
})
@@ -458,7 +495,7 @@ func TestHandleFailoverError_FailedAccountIDs(t *testing.T) {
mock := &mockTempUnscheduler{}
fs := NewFailoverState(0, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(500, false, false))
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
require.Equal(t, FailoverExhausted, action)
require.Contains(t, fs.FailedAccountIDs, int64(100))
})
@@ -467,7 +504,7 @@ func TestHandleFailoverError_FailedAccountIDs(t *testing.T) {
mock := &mockTempUnscheduler{}
fs := NewFailoverState(3, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(400, true, false))
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(400, true, false))
require.Equal(t, FailoverContinue, action)
require.NotContains(t, fs.FailedAccountIDs, int64(100))
})
@@ -476,8 +513,8 @@ func TestHandleFailoverError_FailedAccountIDs(t *testing.T) {
mock := &mockTempUnscheduler{}
fs := NewFailoverState(5, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(500, false, false))
fs.HandleFailoverError(context.Background(), mock, 100, "openai", newTestFailoverErr(500, false, false))
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, newTestFailoverErr(500, false, false))
require.Len(t, fs.FailedAccountIDs, 1, "map 天然去重")
})
}
@@ -492,11 +529,11 @@ func TestHandleFailoverError_LastFailoverErr(t *testing.T) {
fs := NewFailoverState(3, false)
err1 := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err1)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
require.Equal(t, err1, fs.LastFailoverErr)
err2 := newTestFailoverErr(502, false, false)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", err2)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
require.Equal(t, err2, fs.LastFailoverErr)
})
@@ -505,7 +542,7 @@ func TestHandleFailoverError_LastFailoverErr(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(400, true, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, err, fs.LastFailoverErr)
})
}
@@ -522,30 +559,30 @@ func TestHandleFailoverError_IntegrationScenario(t *testing.T) {
// 1. 账号 100 遇到可重试错误,同账号重试 maxSameAccountRetries 次
retryErr := newTestFailoverErr(400, true, false)
for i := 0; i < maxSameAccountRetries; i++ {
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryErr)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, retryErr)
require.Equal(t, FailoverContinue, action)
}
require.True(t, fs.ForceCacheBilling, "hasBoundSession=true 应设置 ForceCacheBilling")
// 2. 账号 100 超过重试上限 → TempUnschedule + 切换
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", retryErr)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, retryErr)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SwitchCount)
require.Len(t, mock.calls, 1)
// 3. 账号 200 遇到不可重试错误 → 直接切换
switchErr := newTestFailoverErr(500, false, false)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", switchErr)
action = fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, switchErr)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 2, fs.SwitchCount)
// 4. 账号 300 遇到不可重试错误 → 再切换
action = fs.HandleFailoverError(context.Background(), mock, 300, "openai", switchErr)
action = fs.HandleFailoverError(context.Background(), mock, 300, "openai", maxSameAccountRetries, switchErr)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 3, fs.SwitchCount)
// 5. 账号 400 → 已耗尽 (SwitchCount=3 >= MaxSwitches=3)
action = fs.HandleFailoverError(context.Background(), mock, 400, "openai", switchErr)
action = fs.HandleFailoverError(context.Background(), mock, 400, "openai", maxSameAccountRetries, switchErr)
require.Equal(t, FailoverExhausted, action)
// 最终状态验证
@@ -563,21 +600,21 @@ func TestHandleFailoverError_IntegrationScenario(t *testing.T) {
// 第一次切换:delay = 0s
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformAntigravity, err)
action := fs.HandleFailoverError(context.Background(), mock, 100, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
require.Less(t, elapsed, 200*time.Millisecond, "第一次切换延迟为 0")
// 第二次切换:delay = 1s
start = time.Now()
action = fs.HandleFailoverError(context.Background(), mock, 200, service.PlatformAntigravity, err)
action = fs.HandleFailoverError(context.Background(), mock, 200, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed = time.Since(start)
require.Equal(t, FailoverContinue, action)
require.GreaterOrEqual(t, elapsed, 800*time.Millisecond, "第二次切换延迟约 1s")
// 第三次:耗尽(无延迟,因为在检查延迟之前就返回了)
start = time.Now()
action = fs.HandleFailoverError(context.Background(), mock, 300, service.PlatformAntigravity, err)
action = fs.HandleFailoverError(context.Background(), mock, 300, service.PlatformAntigravity, maxSameAccountRetries, err)
elapsed = time.Since(start)
require.Equal(t, FailoverExhausted, action)
require.Less(t, elapsed, 200*time.Millisecond, "耗尽时不应有延迟")
@@ -589,17 +626,17 @@ func TestHandleFailoverError_IntegrationScenario(t *testing.T) {
// 第一次:ForceCacheBilling=false
err1 := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", err1)
fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err1)
require.False(t, fs.ForceCacheBilling)
// 第二次:ForceCacheBilling=true(Antigravity 粘性会话切换)
err2 := newTestFailoverErr(500, false, true)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", err2)
fs.HandleFailoverError(context.Background(), mock, 200, "openai", maxSameAccountRetries, err2)
require.True(t, fs.ForceCacheBilling, "错误标志应触发 ForceCacheBilling")
// 第三次:ForceCacheBilling=false,但状态仍保持 true
err3 := newTestFailoverErr(500, false, false)
fs.HandleFailoverError(context.Background(), mock, 300, "openai", err3)
fs.HandleFailoverError(context.Background(), mock, 300, "openai", maxSameAccountRetries, err3)
require.True(t, fs.ForceCacheBilling, "不应重置")
})
}
@@ -614,7 +651,7 @@ func TestHandleFailoverError_EdgeCases(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(0, false, false)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
})
@@ -623,7 +660,7 @@ func TestHandleFailoverError_EdgeCases(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, true, false)
action := fs.HandleFailoverError(context.Background(), mock, 0, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, 0, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[0])
})
@@ -633,7 +670,7 @@ func TestHandleFailoverError_EdgeCases(t *testing.T) {
fs := NewFailoverState(3, false)
err := newTestFailoverErr(500, true, false)
action := fs.HandleFailoverError(context.Background(), mock, -1, "openai", err)
action := fs.HandleFailoverError(context.Background(), mock, -1, "openai", maxSameAccountRetries, err)
require.Equal(t, FailoverContinue, action)
require.Equal(t, 1, fs.SameAccountRetryCount[-1])
})
@@ -645,7 +682,7 @@ func TestHandleFailoverError_EdgeCases(t *testing.T) {
err := newTestFailoverErr(500, false, false)
start := time.Now()
action := fs.HandleFailoverError(context.Background(), mock, 100, "", err)
action := fs.HandleFailoverError(context.Background(), mock, 100, "", maxSameAccountRetries, err)
elapsed := time.Since(start)
require.Equal(t, FailoverContinue, action)
+10 -9
View File
@@ -448,7 +448,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
h.handleFailoverExhausted(c, failoverErr, service.PlatformGemini, true)
return
}
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, failoverErr)
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch action {
case FailoverContinue:
continue
@@ -868,7 +868,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
h.handleFailoverExhausted(c, failoverErr, account.Platform, true)
return
}
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, failoverErr)
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch action {
case FailoverContinue:
continue
@@ -1454,13 +1454,14 @@ func (h *GatewayHandler) usageUnrestricted(c *gin.Context, ctx context.Context,
remaining := h.calculateSubscriptionRemaining(apiKey.Group, subscription)
resp["remaining"] = remaining
resp["subscription"] = gin.H{
"daily_usage_usd": subscription.DailyUsageUSD,
"weekly_usage_usd": subscription.WeeklyUsageUSD,
"monthly_usage_usd": subscription.MonthlyUsageUSD,
"daily_limit_usd": apiKey.Group.DailyLimitUSD,
"weekly_limit_usd": apiKey.Group.WeeklyLimitUSD,
"monthly_limit_usd": apiKey.Group.MonthlyLimitUSD,
"expires_at": subscription.ExpiresAt,
"daily_usage_usd": subscription.DailyUsageUSD,
"weekly_usage_usd": subscription.WeeklyUsageUSD,
"monthly_usage_usd": subscription.MonthlyUsageUSD,
"daily_limit_usd": apiKey.Group.DailyLimitUSD,
"weekly_limit_usd": apiKey.Group.WeeklyLimitUSD,
"monthly_limit_usd": apiKey.Group.MonthlyLimitUSD,
"weekly_window_start": subscription.WeeklyWindowStart,
"expires_at": subscription.ExpiresAt,
}
}
@@ -254,7 +254,7 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
h.handleCCFailoverExhausted(c, failoverErr, true)
return
}
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, failoverErr)
action := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch action {
case FailoverContinue:
continue
@@ -233,7 +233,7 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
h.handleResponsesFailoverExhausted(c, failoverErr, true)
return
}
action := fs.HandleFailoverError(requestCtx, h.gatewayService, account.ID, account.Platform, failoverErr)
action := fs.HandleFailoverError(requestCtx, h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch action {
case FailoverContinue:
continue
@@ -0,0 +1,51 @@
package handler
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestUsageUnrestrictedIncludesWeeklyWindowStart(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/v1/usage", nil)
weeklyWindowStart := time.Date(2026, time.July, 13, 0, 30, 0, 0, time.FixedZone("UTC+8", 8*60*60))
c.Set(string(middleware.ContextKeySubscription), &service.UserSubscription{
WeeklyWindowStart: &weeklyWindowStart,
})
handler := &GatewayHandler{}
handler.usageUnrestricted(
c,
context.Background(),
&service.APIKey{Group: &service.Group{
Name: "Weekly plan",
SubscriptionType: service.SubscriptionTypeSubscription,
}},
middleware.AuthSubject{},
nil,
nil,
nil,
)
require.Equal(t, http.StatusOK, recorder.Code)
var response struct {
Subscription struct {
WeeklyWindowStart *time.Time `json:"weekly_window_start"`
} `json:"subscription"`
}
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &response))
require.NotNil(t, response.Subscription.WeeklyWindowStart)
require.True(t, weeklyWindowStart.Equal(*response.Subscription.WeeklyWindowStart))
}
@@ -220,6 +220,15 @@ func (h *ConcurrencyHelper) TryAcquireUserSlotForAPIKey(ctx context.Context, use
return h.withAPIKeySlot(ctx, apiKeyID, releaseFunc), true, nil
}
// AcquireOpenAIWSIngressLease bounds the whole client WebSocket lifecycle,
// independently from per-turn user and account slots.
func (h *ConcurrencyHelper) AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int) (*service.OpenAIWSIngressLease, bool, error) {
if h == nil || h.concurrencyService == nil {
return nil, false, fmt.Errorf("concurrency service is unavailable")
}
return h.concurrencyService.AcquireOpenAIWSIngressLease(ctx, apiKeyID, maxConnections)
}
// TryAcquireAccountSlot 尝试立即获取账号并发槽位。
// 返回值: (releaseFunc, acquired, error)
func (h *ConcurrencyHelper) TryAcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int) (func(), bool, error) {
@@ -11,10 +11,13 @@ import (
)
type concurrencyCacheMock struct {
acquireUserSlotFn func(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error)
acquireAccountSlotFn func(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error)
releaseUserCalled int32
releaseAccountCalled int32
acquireUserSlotFn func(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error)
acquireAccountSlotFn func(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error)
acquireIngressLeaseFn func(ctx context.Context, apiKeyID int64, maxConnections int, leaseID string) (bool, error)
releaseIngressLeaseFn func(ctx context.Context, apiKeyID int64, leaseID string) error
releaseUserCalled int32
releaseAccountCalled int32
releaseIngressCalled int32
}
func (m *concurrencyCacheMock) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
@@ -97,6 +100,25 @@ func (m *concurrencyCacheMock) CleanupStaleProcessSlots(ctx context.Context, act
return nil
}
func (m *concurrencyCacheMock) AcquireOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, maxConnections int, leaseID string) (bool, error) {
if m.acquireIngressLeaseFn != nil {
return m.acquireIngressLeaseFn(ctx, apiKeyID, maxConnections, leaseID)
}
return false, nil
}
func (m *concurrencyCacheMock) RefreshOpenAIWSIngressLease(context.Context, int64, string) (bool, error) {
return true, nil
}
func (m *concurrencyCacheMock) ReleaseOpenAIWSIngressLease(ctx context.Context, apiKeyID int64, leaseID string) error {
atomic.AddInt32(&m.releaseIngressCalled, 1)
if m.releaseIngressLeaseFn != nil {
return m.releaseIngressLeaseFn(ctx, apiKeyID, leaseID)
}
return nil
}
func TestConcurrencyHelper_TryAcquireUserSlot(t *testing.T) {
cache := &concurrencyCacheMock{
acquireUserSlotFn: func(ctx context.Context, userID int64, maxConcurrency int, requestID string) (bool, error) {
@@ -482,7 +482,7 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
if err != nil {
var failoverErr *service.UpstreamFailoverError
if errors.As(err, &failoverErr) {
failoverAction := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, failoverErr)
failoverAction := fs.HandleFailoverError(c.Request.Context(), h.gatewayService, account.ID, account.Platform, account.GetPoolModeRetryCount(), failoverErr)
switch failoverAction {
case FailoverContinue:
continue
+11 -1
View File
@@ -31,6 +31,16 @@ func (h *OpenAIGatewayHandler) GrokVideoGeneration(c *gin.Context) {
h.handleGrokMedia(c, service.GrokMediaEndpointVideosGenerations, "")
}
// GrokVideoEdit handles asynchronous xAI video edits through Grok groups.
func (h *OpenAIGatewayHandler) GrokVideoEdit(c *gin.Context) {
h.handleGrokMedia(c, service.GrokMediaEndpointVideosEdits, "")
}
// GrokVideoExtension handles asynchronous xAI video extensions through Grok groups.
func (h *OpenAIGatewayHandler) GrokVideoExtension(c *gin.Context) {
h.handleGrokMedia(c, service.GrokMediaEndpointVideosExtensions, "")
}
// GrokVideoStatus handles xAI video status retrieval through Grok groups.
func (h *OpenAIGatewayHandler) GrokVideoStatus(c *gin.Context) {
h.handleGrokMedia(c, service.GrokMediaEndpointVideoStatus, c.Param("request_id"))
@@ -298,7 +308,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil)
if endpoint == service.GrokMediaEndpointVideosGenerations && strings.TrimSpace(result.ResponseID) != "" {
if endpoint.IsGenerationRequest() && strings.TrimSpace(result.ResponseID) != "" {
if err := h.gatewayService.BindGrokMediaVideoRequestAccount(requestCtx, apiKey.GroupID, result.ResponseID, account.ID); err != nil {
reqLog.Warn("grok_media.bind_video_request_account_failed",
zap.Int64("account_id", account.ID),
@@ -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)
}
@@ -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}}
@@ -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))
}
})
}
@@ -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)
}
@@ -15,11 +15,13 @@ import (
// Codex CLI and the Codex desktop app refresh their model picker from
// GET {base_url}/models?client_version=... (custom provider mode) or
// GET /backend-api/codex/models (chatgpt_base_url mode). Both routes land
// here. The manifest is proxied verbatim from the ChatGPT backend with a
// schedulable OAuth account's credentials, so clients pointed at the gateway
// see the account's real, always-current model entitlements instead of a
// frozen local cache.
// here. The manifest is proxied verbatim from the selected account's ChatGPT
// backend or custom API key upstream. API key manifests use a short-lived,
// asynchronously revalidated cache to tolerate canceled client requests.
func (h *OpenAIGatewayHandler) CodexModels(c *gin.Context) {
if c.Request.Context().Err() != nil {
return
}
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
if !ok || apiKey.Group == nil {
h.errorResponse(c, http.StatusUnauthorized, "invalid_request_error", "API key group is required")
@@ -30,24 +32,54 @@ func (h *OpenAIGatewayHandler) CodexModels(c *gin.Context) {
return
}
account, err := h.gatewayService.SelectAccountForModel(c.Request.Context(), apiKey.GroupID, "", "")
if err != nil {
h.errorResponse(c, http.StatusServiceUnavailable, "upstream_error", "No available OpenAI accounts")
return
maxAccountSwitches := h.maxAccountSwitches
if maxAccountSwitches <= 0 {
maxAccountSwitches = 3
}
failedAccountIDs := make(map[int64]struct{})
switchCount := 0
var lastUpstreamErr error
manifest, err := h.gatewayService.FetchCodexModelsManifest(c.Request.Context(), account, c.Query("client_version"), c.GetHeader("If-None-Match"))
if err != nil {
h.errorResponse(c, infraerrors.Code(err), "upstream_error", infraerrors.Message(err))
return
}
for {
account, err := h.gatewayService.SelectAccountForModelWithExclusions(c.Request.Context(), apiKey.GroupID, "", "", failedAccountIDs)
if err != nil {
if c.Request.Context().Err() != nil {
return
}
if lastUpstreamErr != nil {
h.errorResponse(c, infraerrors.Code(lastUpstreamErr), "upstream_error", infraerrors.Message(lastUpstreamErr))
return
}
h.errorResponse(c, http.StatusServiceUnavailable, "upstream_error", "No available OpenAI accounts")
return
}
if manifest.ETag != "" {
c.Header("ETag", manifest.ETag)
}
if manifest.NotModified {
c.Status(http.StatusNotModified)
manifest, err := h.gatewayService.FetchCodexModelsManifest(c.Request.Context(), account, c.Query("client_version"), c.GetHeader("If-None-Match"))
if err != nil {
if c.Request.Context().Err() != nil {
return
}
if service.IsRetryableCodexModelsManifestError(err) && switchCount < maxAccountSwitches {
failedAccountIDs[account.ID] = struct{}{}
switchCount++
lastUpstreamErr = err
continue
}
h.errorResponse(c, infraerrors.Code(err), "upstream_error", infraerrors.Message(err))
return
}
if c.Request.Context().Err() != nil {
return
}
if manifest.ETag != "" {
c.Header("ETag", manifest.ETag)
}
if manifest.NotModified {
c.Status(http.StatusNotModified)
return
}
c.Data(http.StatusOK, "application/json", manifest.Body)
return
}
c.Data(http.StatusOK, "application/json", manifest.Body)
}
@@ -0,0 +1,288 @@
package handler
import (
"context"
"errors"
"fmt"
"io"
"net"
"net/http"
"net/http/httptest"
"strings"
"sync"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
type codexModelsFailoverAccountRepo struct {
service.AccountRepository
accounts []service.Account
}
func (r codexModelsFailoverAccountRepo) GetByID(_ context.Context, id int64) (*service.Account, error) {
for i := range r.accounts {
if r.accounts[i].ID == id {
account := r.accounts[i]
return &account, nil
}
}
return nil, service.ErrNoAvailableAccounts
}
func (r codexModelsFailoverAccountRepo) ListSchedulableByPlatform(_ context.Context, platform string) ([]service.Account, error) {
accounts := make([]service.Account, 0, len(r.accounts))
for _, account := range r.accounts {
if account.Platform == platform {
accounts = append(accounts, account)
}
}
return accounts, nil
}
type codexModelsFailoverHTTPUpstream struct {
service.HTTPUpstream
mu sync.Mutex
accountIDs []int64
firstErr error
firstStatus int
statuses map[int64]int
}
func (u *codexModelsFailoverHTTPUpstream) Do(_ *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
u.mu.Lock()
u.accountIDs = append(u.accountIDs, accountID)
u.mu.Unlock()
status, hasStatus := u.statuses[accountID]
if accountID == 1 || hasStatus {
if u.firstErr != nil {
return nil, u.firstErr
}
if !hasStatus {
status = u.firstStatus
}
if status == 0 {
status = http.StatusServiceUnavailable
}
return &http.Response{
StatusCode: status,
Status: http.StatusText(status),
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(
`{"error":{"message":"No available OpenAI accounts","type":"upstream_error"}}`,
)),
}, nil
}
return &http.Response{
StatusCode: http.StatusOK,
Status: "200 OK",
Header: make(http.Header),
Body: io.NopCloser(strings.NewReader(`{"models":[{"slug":"gpt-5.6-sol"}]}`)),
}, nil
}
func (u *codexModelsFailoverHTTPUpstream) calls() []int64 {
u.mu.Lock()
defer u.mu.Unlock()
return append([]int64(nil), u.accountIDs...)
}
func TestCodexModelsCanceledRequestDoesNotWriteResponse(t *testing.T) {
gin.SetMode(gin.TestMode)
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
ctx, cancel := context.WithCancel(context.Background())
cancel()
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil).WithContext(ctx)
h := &OpenAIGatewayHandler{}
h.CodexModels(c)
if c.Writer.Written() {
t.Fatalf("canceled request wrote an HTTP response: status=%d body=%q", recorder.Code, recorder.Body.String())
}
}
func TestCodexModelsFailsOverFromRetryableUpstreamStatus(t *testing.T) {
retryableStatuses := []int{
http.StatusTooManyRequests,
http.StatusInternalServerError,
http.StatusBadGateway,
http.StatusServiceUnavailable,
http.StatusGatewayTimeout,
}
for _, status := range retryableStatuses {
t.Run(http.StatusText(status), func(t *testing.T) {
handler, upstream, groupID := newCodexModelsFailoverTestHandler(status)
recorder := performCodexModelsRequest(t, handler, groupID)
if got, want := upstream.calls(), []int64{1, 2}; !equalInt64Slices(got, want) {
t.Fatalf("upstream account calls: got %v, want %v", got, want)
}
if recorder.Code != http.StatusOK {
t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusOK, recorder.Body.String())
}
if got, want := recorder.Body.String(), `{"models":[{"slug":"gpt-5.6-sol"}]}`; got != want {
t.Fatalf("body: got %q, want %q", got, want)
}
})
}
}
func TestCodexModelsFailsOverFromUpstreamTransportError(t *testing.T) {
handler, upstream, groupID := newCodexModelsFailoverTestHandler(http.StatusServiceUnavailable)
upstream.firstErr = &net.OpError{
Op: "read",
Net: "tcp",
Err: errors.New("connection reset"),
}
recorder := performCodexModelsRequest(t, handler, groupID)
if got, want := upstream.calls(), []int64{1, 2}; !equalInt64Slices(got, want) {
t.Fatalf("upstream account calls: got %v, want %v", got, want)
}
if recorder.Code != http.StatusOK {
t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusOK, recorder.Body.String())
}
}
func TestCodexModelsDoesNotFailOverFromPermanentUpstreamStatus(t *testing.T) {
statuses := []int{
http.StatusBadRequest,
http.StatusUnauthorized,
http.StatusForbidden,
http.StatusNotFound,
600,
}
for _, status := range statuses {
t.Run(fmt.Sprintf("status_%d", status), func(t *testing.T) {
handler, upstream, groupID := newCodexModelsFailoverTestHandler(status)
recorder := performCodexModelsRequest(t, handler, groupID)
if got, want := upstream.calls(), []int64{1}; !equalInt64Slices(got, want) {
t.Fatalf("upstream account calls: got %v, want %v", got, want)
}
if recorder.Code != http.StatusBadGateway {
t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusBadGateway, recorder.Body.String())
}
})
}
}
func TestCodexModelsDoesNotFailOverFromUpstreamConfigurationError(t *testing.T) {
handler, upstream, groupID := newCodexModelsFailoverTestHandler(http.StatusServiceUnavailable)
upstream.firstErr = errors.New("invalid proxy URL")
recorder := performCodexModelsRequest(t, handler, groupID)
if got, want := upstream.calls(), []int64{1}; !equalInt64Slices(got, want) {
t.Fatalf("upstream account calls: got %v, want %v", got, want)
}
if recorder.Code != http.StatusBadGateway {
t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusBadGateway, recorder.Body.String())
}
}
func TestCodexModelsReturnsLastUpstreamErrorWhenAccountsAreExhausted(t *testing.T) {
handler, upstream, groupID := newCodexModelsFailoverTestHandler(http.StatusServiceUnavailable)
upstream.statuses = map[int64]int{
1: http.StatusServiceUnavailable,
2: http.StatusGatewayTimeout,
}
recorder := performCodexModelsRequest(t, handler, groupID)
if got, want := upstream.calls(), []int64{1, 2}; !equalInt64Slices(got, want) {
t.Fatalf("upstream account calls: got %v, want %v", got, want)
}
if recorder.Code != http.StatusBadGateway {
t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusBadGateway, recorder.Body.String())
}
if body := recorder.Body.String(); !strings.Contains(body, "upstream error 504") {
t.Fatalf("body does not preserve the last upstream error: %s", body)
}
}
func TestCodexModelsHonorsAccountSwitchLimit(t *testing.T) {
handler, upstream, groupID := newCodexModelsFailoverTestHandlerWithAccountCount(http.StatusServiceUnavailable, 4, 2)
upstream.statuses = map[int64]int{
1: http.StatusServiceUnavailable,
2: http.StatusBadGateway,
3: http.StatusGatewayTimeout,
4: http.StatusInternalServerError,
}
recorder := performCodexModelsRequest(t, handler, groupID)
if got, want := upstream.calls(), []int64{1, 2, 3}; !equalInt64Slices(got, want) {
t.Fatalf("upstream account calls: got %v, want %v", got, want)
}
if recorder.Code != http.StatusBadGateway {
t.Fatalf("status: got %d, want %d; body=%s", recorder.Code, http.StatusBadGateway, recorder.Body.String())
}
if body := recorder.Body.String(); !strings.Contains(body, "upstream error 504") {
t.Fatalf("body does not preserve the limit-ending upstream error: %s", body)
}
}
func newCodexModelsFailoverTestHandler(firstStatus int) (*OpenAIGatewayHandler, *codexModelsFailoverHTTPUpstream, int64) {
return newCodexModelsFailoverTestHandlerWithAccountCount(firstStatus, 2, 3)
}
func newCodexModelsFailoverTestHandlerWithAccountCount(firstStatus, accountCount, maxSwitches int) (*OpenAIGatewayHandler, *codexModelsFailoverHTTPUpstream, int64) {
gin.SetMode(gin.TestMode)
groupID := int64(42)
accounts := make([]service.Account, 0, accountCount)
for i := 1; i <= accountCount; i++ {
accounts = append(accounts, service.Account{
ID: int64(i),
Name: fmt.Sprintf("upstream-%d", i),
Platform: service.PlatformOpenAI,
Type: service.AccountTypeAPIKey,
Status: service.StatusActive,
Schedulable: true,
Priority: i - 1,
Concurrency: 1,
Credentials: map[string]any{
"api_key": fmt.Sprintf("sk-%d", i),
"base_url": fmt.Sprintf("https://upstream-%d.example/v1", i),
},
})
}
upstream := &codexModelsFailoverHTTPUpstream{firstStatus: firstStatus}
cfg := &config.Config{RunMode: config.RunModeSimple}
gatewayService := service.NewOpenAIGatewayService(
codexModelsFailoverAccountRepo{accounts: accounts},
nil, nil, nil, nil, nil, nil, cfg, nil, nil, nil, nil, nil,
upstream,
nil, nil, nil, nil, nil, nil, nil, nil,
)
return &OpenAIGatewayHandler{gatewayService: gatewayService, maxAccountSwitches: maxSwitches}, upstream, groupID
}
func performCodexModelsRequest(t *testing.T, handler *OpenAIGatewayHandler, groupID int64) *httptest.ResponseRecorder {
t.Helper()
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models?client_version=0.144.0", nil)
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
GroupID: &groupID,
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI},
})
handler.CodexModels(c)
return recorder
}
func equalInt64Slices(got, want []int64) bool {
if len(got) != len(want) {
return false
}
for i := range got {
if got[i] != want[i] {
return false
}
}
return true
}
@@ -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) {
@@ -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)
}
@@ -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
@@ -1287,10 +1299,36 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
wsConn.SetReadLimit(service.ResolveOpenAIWSClientReadLimitBytes(h.cfg))
ctx := c.Request.Context()
maxIngressConnections := 0
if h.cfg != nil {
maxIngressConnections = h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey
}
ingressLease, ingressLeaseAcquired, ingressLeaseErr := h.concurrencyHelper.AcquireOpenAIWSIngressLease(ctx, apiKey.ID, maxIngressConnections)
if ingressLeaseErr != nil {
reqLog.Error("openai.websocket_ingress_lease_acquire_failed", zap.Error(ingressLeaseErr))
closeOpenAIClientWS(wsConn, coderws.StatusInternalError, "failed to reserve websocket ingress capacity")
return
}
if !ingressLeaseAcquired {
reqLog.Info("openai.websocket_ingress_capacity_rejected", zap.Int("max_ingress_connections_per_api_key", maxIngressConnections))
closeOpenAIClientWS(wsConn, coderws.StatusTryAgainLater, "too many open websocket connections, please retry later")
return
}
if ingressLease != nil {
defer ingressLease.Release()
ctx = ingressLease.Context()
c.Request = c.Request.WithContext(ctx)
}
readCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
msgType, firstMessage, err := wsConn.Read(readCtx)
cancel()
if err != nil {
if errors.Is(context.Cause(ctx), service.ErrOpenAIWSIngressLeaseLost) {
reqLog.Warn("openai.websocket_ingress_lease_lost_before_first_message", zap.Error(err))
closeOpenAIClientWS(wsConn, coderws.StatusTryAgainLater, "websocket ingress capacity lease lost; please reconnect")
return
}
closeStatus, closeReason := summarizeWSCloseErrorForLog(err)
reqLog.Warn("openai.websocket_read_first_message_failed",
zap.Error(err),
@@ -1444,7 +1482,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 +1639,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) {
@@ -1680,6 +1718,25 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
continue
}
if errors.Is(context.Cause(ctx), service.ErrOpenAIWSIngressLeaseLost) {
reqLog.Warn("openai.websocket_ingress_lease_lost",
zap.Int64("account_id", account.ID),
zap.Error(err),
)
closeOpenAIClientWS(wsConn, coderws.StatusTryAgainLater, "websocket ingress capacity lease lost; please reconnect")
return
}
var closeErr *service.OpenAIWSClientCloseError
if errors.As(err, &closeErr) && closeErr.StatusCode() == coderws.StatusNormalClosure {
reqLog.Info("openai.websocket_ingress_closed_normally",
zap.Int64("account_id", account.ID),
zap.String("reason", closeErr.Reason()),
)
closeOpenAIClientWS(wsConn, closeErr.StatusCode(), closeErr.Reason())
return
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
closeStatus, closeReason := summarizeWSCloseErrorForLog(err)
reqLog.Warn("openai.websocket_proxy_failed",
@@ -1688,7 +1745,6 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
zap.String("close_status", closeStatus),
zap.String("close_reason", closeReason),
)
var closeErr *service.OpenAIWSClientCloseError
if errors.As(err, &closeErr) {
closeOpenAIClientWS(wsConn, closeErr.StatusCode(), closeErr.Reason())
return
@@ -2035,7 +2091,8 @@ func openAIForwardErrorAlreadyCommunicated(c *gin.Context, writerSizeBeforeForwa
}
// 与快照同口径:排除 compact 心跳字节,避免"仅心跳写出"被误判为
// 响应已写出(#3887)。
if service.OpenAICompactKeepaliveAdjustedWrittenSize(c) == writerSizeBeforeForward {
if service.OpenAICompactKeepaliveAdjustedWrittenSize(c) == writerSizeBeforeForward ||
service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) == writerSizeBeforeForward {
return false
}
@@ -2441,7 +2498,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 {
@@ -8,6 +8,7 @@ import (
"net/http/httptest"
"strings"
"sync"
"sync/atomic"
"testing"
"time"
@@ -414,11 +415,13 @@ func TestResolveOpenAIMessagesDispatchMappedModel(t *testing.T) {
SonnetMappedModel: "gpt-5.2",
ExactModelMappings: map[string]string{
"claude-sonnet-4-5-20250929": "gpt-5.4-mini-high",
"claude-fable-5": "gpt-5.6-sol",
},
},
},
}
require.Equal(t, "gpt-5.4-mini", resolveOpenAIMessagesDispatchMappedModel(apiKey, "claude-sonnet-4-5-20250929"))
require.Equal(t, "gpt-5.6-sol", resolveOpenAIMessagesDispatchMappedModel(apiKey, "claude-fable-5"))
})
t.Run("uses_family_default_when_no_override", func(t *testing.T) {
@@ -711,6 +714,68 @@ func TestOpenAIResponsesWebSocket_InvalidUpgradeDoesNotSetTransport(t *testing.T
require.Equal(t, service.OpenAIClientTransportUnknown, service.GetOpenAIClientTransport(c))
}
func TestOpenAIResponsesWebSocket_IngressCapacityRejected(t *testing.T) {
gin.SetMode(gin.TestMode)
cache := &concurrencyCacheMock{
acquireIngressLeaseFn: func(context.Context, int64, int, string) (bool, error) {
return false, nil
},
}
h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache)
h.cfg = &config.Config{}
h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = 1
wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1})
defer wsServer.Close()
dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil)
cancelDial()
require.NoError(t, err)
defer func() { _ = clientConn.CloseNow() }()
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
_, _, err = clientConn.Read(readCtx)
cancelRead()
var closeErr coderws.CloseError
require.ErrorAs(t, err, &closeErr)
require.Equal(t, coderws.StatusTryAgainLater, closeErr.Code)
}
func TestOpenAIResponsesWebSocket_IngressLeaseReleasedOnEarlyReturn(t *testing.T) {
gin.SetMode(gin.TestMode)
cache := &concurrencyCacheMock{
acquireIngressLeaseFn: func(context.Context, int64, int, string) (bool, error) {
return true, nil
},
}
h := newOpenAIHandlerForPreviousResponseIDValidation(t, cache)
h.cfg = &config.Config{}
h.cfg.Gateway.OpenAIWS.MaxIngressConnectionsPerAPIKey = 1
wsServer := newOpenAIWSHandlerTestServer(t, h, middleware.AuthSubject{UserID: 1, Concurrency: 1})
defer wsServer.Close()
dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
clientConn, _, err := coderws.Dial(dialCtx, "ws"+strings.TrimPrefix(wsServer.URL, "http")+"/openai/v1/responses", nil)
cancelDial()
require.NoError(t, err)
defer func() { _ = clientConn.CloseNow() }()
writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
err = clientConn.Write(writeCtx, coderws.MessageBinary, []byte("not a response.create frame"))
cancelWrite()
require.NoError(t, err)
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
_, _, err = clientConn.Read(readCtx)
cancelRead()
var closeErr coderws.CloseError
require.ErrorAs(t, err, &closeErr)
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
require.Eventually(t, func() bool {
return atomic.LoadInt32(&cache.releaseIngressCalled) == 1
}, time.Second, 10*time.Millisecond)
}
func TestOpenAIResponsesWebSocket_RejectsMessageIDAsPreviousResponseID(t *testing.T) {
gin.SetMode(gin.TestMode)
+16 -2
View File
@@ -142,6 +142,9 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
failedAccountIDs := make(map[int64]struct{})
sameAccountRetryCount := make(map[int64]int)
var lastFailoverErr *service.UpstreamFailoverError
stopJSONKeepalive := func() {}
jsonKeepaliveStarted := false
defer func() { stopJSONKeepalive() }()
for {
reqLog.Debug("openai.images.account_selecting", zap.Int("excluded_account_count", len(failedAccountIDs)))
@@ -210,8 +213,12 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
}
service.SetOpsLatencyMs(c, service.OpsRoutingLatencyMsKey, time.Since(routingStart).Milliseconds())
if !parsed.Stream && !jsonKeepaliveStarted {
stopJSONKeepalive = service.StartOpenAIImagesJSONKeepalive(c, h.openAIImagesJSONKeepaliveInterval())
jsonKeepaliveStarted = true
}
forwardStart := time.Now()
writerSizeBeforeForward := c.Writer.Size()
writerSizeBeforeForward := service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c)
result, err := func() (*service.OpenAIForwardResult, error) {
defer func() {
if accountReleaseFunc != nil {
@@ -258,7 +265,7 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
var failoverErr *service.UpstreamFailoverError
if errors.As(err, &failoverErr) {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, false, nil)
if c.Writer.Size() != writerSizeBeforeForward {
if service.OpenAIImagesJSONKeepaliveAdjustedWrittenSize(c) != writerSizeBeforeForward {
reqLog.Warn("openai.images.upstream_failover_skipped_after_flush",
zap.Int64("account_id", account.ID),
zap.Int("upstream_status", failoverErr.StatusCode),
@@ -383,6 +390,13 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
}
}
func (h *OpenAIGatewayHandler) openAIImagesJSONKeepaliveInterval() time.Duration {
if h.cfg == nil || h.cfg.Gateway.ImageNonstreamKeepaliveInterval <= 0 {
return 0
}
return time.Duration(h.cfg.Gateway.ImageNonstreamKeepaliveInterval) * time.Second
}
func isMultipartImagesContentType(contentType string) bool {
return strings.HasPrefix(strings.ToLower(strings.TrimSpace(contentType)), "multipart/form-data")
}
@@ -0,0 +1,90 @@
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) {
w := &opsCaptureWriter{}
w.ResponseWriter = nil
assert.NotPanics(t, func() {
assert.Equal(t, 0, w.Status())
})
assert.NotPanics(t, func() {
assert.Equal(t, -1, w.Size())
})
assert.NotPanics(t, func() {
assert.False(t, w.Written())
})
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() {
h := w.Header()
assert.NotNil(t, h)
})
assert.NotPanics(t, func() {
w.WriteHeader(200)
})
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() {
p := w.Pusher()
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)
}
@@ -1,11 +1,14 @@
package handler
import (
"bufio"
"bytes"
"context"
"encoding/json"
"errors"
"log"
"net"
"net/http"
"runtime"
"runtime/debug"
"strconv"
@@ -498,7 +501,82 @@ func releaseOpsCaptureWriter(w *opsCaptureWriter) {
opsCaptureWriterPool.Put(w)
}
func (w *opsCaptureWriter) Status() int {
if w.ResponseWriter == nil {
return 0
}
return w.ResponseWriter.Status()
}
func (w *opsCaptureWriter) Size() int {
if w.ResponseWriter == nil {
return -1
}
return w.ResponseWriter.Size()
}
func (w *opsCaptureWriter) Written() bool {
if w.ResponseWriter == nil {
return false
}
return w.ResponseWriter.Written()
}
func (w *opsCaptureWriter) Header() http.Header {
if w.ResponseWriter == nil {
return http.Header{}
}
return w.ResponseWriter.Header()
}
func (w *opsCaptureWriter) WriteHeader(code int) {
if w.ResponseWriter == nil {
return
}
w.ResponseWriter.WriteHeader(code)
}
func (w *opsCaptureWriter) WriteHeaderNow() {
if w.ResponseWriter == nil {
return
}
w.ResponseWriter.WriteHeaderNow()
}
func (w *opsCaptureWriter) Flush() {
if w.ResponseWriter == nil {
return
}
w.ResponseWriter.Flush()
}
func (w *opsCaptureWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
if w.ResponseWriter == nil {
return nil, nil, errors.New("response writer released")
}
return w.ResponseWriter.Hijack()
}
func (w *opsCaptureWriter) CloseNotify() <-chan bool {
if w.ResponseWriter == nil {
ch := make(chan bool)
close(ch)
return ch
}
return w.ResponseWriter.CloseNotify()
}
func (w *opsCaptureWriter) Pusher() http.Pusher {
if w.ResponseWriter == nil {
return nil
}
return w.ResponseWriter.Pusher()
}
func (w *opsCaptureWriter) Write(b []byte) (int, error) {
if w.ResponseWriter == nil {
return 0, nil
}
if w.Status() >= 400 && w.limit > 0 && w.buf.Len() < w.limit {
remaining := w.limit - w.buf.Len()
if len(b) > remaining {
@@ -511,6 +589,9 @@ func (w *opsCaptureWriter) Write(b []byte) (int, error) {
}
func (w *opsCaptureWriter) WriteString(s string) (int, error) {
if w.ResponseWriter == nil {
return 0, nil
}
if w.Status() >= 400 && w.limit > 0 && w.buf.Len() < w.limit {
remaining := w.limit - w.buf.Len()
if len(s) > remaining {
+1 -15
View File
@@ -9,7 +9,6 @@ import (
dbent "github.com/Wei-Shaw/sub2api/ent"
"github.com/Wei-Shaw/sub2api/internal/payment"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
@@ -19,15 +18,13 @@ import (
// PaymentHandler handles user-facing payment requests.
type PaymentHandler struct {
channelService *service.ChannelService
paymentService *service.PaymentService
configService *service.PaymentConfigService
}
// NewPaymentHandler creates a new PaymentHandler.
func NewPaymentHandler(paymentService *service.PaymentService, configService *service.PaymentConfigService, channelService *service.ChannelService) *PaymentHandler {
func NewPaymentHandler(paymentService *service.PaymentService, configService *service.PaymentConfigService) *PaymentHandler {
return &PaymentHandler{
channelService: channelService,
paymentService: paymentService,
configService: configService,
}
@@ -91,17 +88,6 @@ func (h *PaymentHandler) GetPlans(c *gin.Context) {
response.Success(c, result)
}
// GetChannels returns enabled payment channels.
// GET /api/v1/payment/channels
func (h *PaymentHandler) GetChannels(c *gin.Context) {
channels, _, err := h.channelService.List(c.Request.Context(), pagination.PaginationParams{Page: 1, PageSize: 1000}, "active", "")
if err != nil {
response.ErrorFrom(c, err)
return
}
response.Success(c, channels)
}
// GetCheckoutInfo returns all data the payment page needs in a single call:
// payment methods with limits, subscription plans, and configuration.
// GET /api/v1/payment/checkout-info
@@ -119,7 +119,7 @@ func TestVerifyOrderPublicReturnsLegacyOrderState(t *testing.T) {
require.NoError(t, err)
paymentSvc := service.NewPaymentService(client, payment.NewRegistry(), nil, nil, nil, nil, nil, nil, nil)
h := NewPaymentHandler(paymentSvc, nil, nil)
h := NewPaymentHandler(paymentSvc, nil)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
@@ -219,7 +219,7 @@ func TestResolveOrderPublicByResumeTokenReturnsFrontendContractFields(t *testing
configSvc := service.NewPaymentConfigService(client, nil, []byte("0123456789abcdef0123456789abcdef"))
paymentSvc := service.NewPaymentService(client, payment.NewRegistry(), nil, nil, nil, configSvc, nil, nil, nil)
h := NewPaymentHandler(paymentSvc, nil, nil)
h := NewPaymentHandler(paymentSvc, nil)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
@@ -307,7 +307,7 @@ func TestResolveOrderPublicByResumeTokenReturnsBadRequestForMismatchedToken(t *t
configSvc := service.NewPaymentConfigService(client, nil, []byte("0123456789abcdef0123456789abcdef"))
paymentSvc := service.NewPaymentService(client, payment.NewRegistry(), nil, nil, nil, configSvc, nil, nil, nil)
h := NewPaymentHandler(paymentSvc, nil, nil)
h := NewPaymentHandler(paymentSvc, nil)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
@@ -347,7 +347,7 @@ func TestVerifyOrderPublicRejectsBlankOutTradeNo(t *testing.T) {
t.Cleanup(func() { _ = client.Close() })
paymentSvc := service.NewPaymentService(client, payment.NewRegistry(), nil, nil, nil, nil, nil, nil, nil)
h := NewPaymentHandler(paymentSvc, nil, nil)
h := NewPaymentHandler(paymentSvc, nil)
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
+1 -1
View File
@@ -164,7 +164,7 @@ var ProviderSet = wire.NewSet(
admin.NewDashboardHandler,
admin.NewUserHandler,
admin.NewGroupHandler,
admin.NewAccountHandler,
admin.ProvideAccountHandler,
admin.NewAnnouncementHandler,
admin.NewDataManagementHandler,
admin.NewBackupHandler,
+9 -9
View File
@@ -17,6 +17,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
)
// ForbiddenError 表示上游返回 403 Forbidden
@@ -279,7 +280,6 @@ func NewClient(proxyURL string) (*Client, error) {
}
client.Transport = transport
}
return &Client{
httpClient: client,
}, nil
@@ -341,7 +341,7 @@ func (c *Client) ExchangeCode(ctx context.Context, code, codeVerifier string) (*
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("token 交换请求失败: %w", err)
}
@@ -383,7 +383,7 @@ func (c *Client) RefreshToken(ctx context.Context, refreshToken string) (*TokenR
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("token 刷新请求失败: %w", err)
}
@@ -414,7 +414,7 @@ func (c *Client) GetUserInfo(ctx context.Context, accessToken string) (*UserInfo
}
req.Header.Set("Authorization", "Bearer "+accessToken)
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("用户信息请求失败: %w", err)
}
@@ -465,7 +465,7 @@ func (c *Client) LoadCodeAssist(ctx context.Context, accessToken string) (*LoadC
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
lastErr = fmt.Errorf("loadCodeAssist 请求失败: %w", err)
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
@@ -544,7 +544,7 @@ func (c *Client) OnboardUser(ctx context.Context, accessToken, tierID string) (s
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
lastErr = fmt.Errorf("onboardUser 请求失败: %w", err)
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
@@ -683,7 +683,7 @@ func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectI
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
resp, err := fetchClient.Do(req)
resp, err := servertiming.Do(fetchClient, req)
if err != nil {
lastErr = fmt.Errorf("fetchAvailableModels 请求失败: %w", err)
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
@@ -842,7 +842,7 @@ func (c *Client) SetUserSettings(ctx context.Context, accessToken string) (*SetU
req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1")
req.Host = "daily-cloudcode-pa.googleapis.com"
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("setUserSettings 请求失败: %w", err)
}
@@ -885,7 +885,7 @@ func (c *Client) FetchUserInfo(ctx context.Context, accessToken, projectID strin
req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1")
req.Host = "daily-cloudcode-pa.googleapis.com"
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("fetchUserInfo 请求失败: %w", err)
}
@@ -718,7 +718,7 @@ func TestStreamingToolCallDoneWithoutDeltaEmitsArguments(t *testing.T) {
assert.Equal(t, "content_block_stop", events[1].Type)
}
func TestStreamingReadToolDropsEmptyPages(t *testing.T) {
func TestStreamingReadToolStreamsDeltas(t *testing.T) {
state := NewResponsesEventToAnthropicState()
ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
@@ -739,18 +739,17 @@ func TestStreamingReadToolDropsEmptyPages(t *testing.T) {
OutputIndex: 0,
Delta: `{"file_path":"/tmp/demo.py","limit":2000,"offset":0,"pages":""}`,
}, state)
assert.Len(t, events, 0)
require.Len(t, events, 1, "Read tool deltas must be streamed like any other tool")
assert.Equal(t, "content_block_delta", events[0].Type)
assert.Equal(t, "input_json_delta", events[0].Delta.Type)
events = ResponsesEventToAnthropicEvents(&ResponsesStreamEvent{
Type: "response.function_call_arguments.done",
OutputIndex: 0,
Arguments: `{"file_path":"/tmp/demo.py","limit":2000,"offset":0,"pages":""}`,
}, state)
require.Len(t, events, 2)
assert.Equal(t, "content_block_delta", events[0].Type)
assert.Equal(t, "input_json_delta", events[0].Delta.Type)
assert.JSONEq(t, `{"file_path":"/tmp/demo.py","limit":2000,"offset":0}`, events[0].Delta.PartialJSON)
assert.Equal(t, "content_block_stop", events[1].Type)
require.Len(t, events, 1, "after streaming deltas, .done should just close the block")
assert.Equal(t, "content_block_stop", events[0].Type)
}
func TestStreamingReasoning(t *testing.T) {
@@ -102,9 +102,10 @@ func AnthropicToResponsesResponse(resp *AnthropicResponse) *ResponsesResponse {
resp.Usage.CacheReadInputTokens +
resp.Usage.CacheCreationInputTokens
out.Usage = &ResponsesUsage{
InputTokens: totalInputTokens,
OutputTokens: resp.Usage.OutputTokens,
TotalTokens: totalInputTokens + resp.Usage.OutputTokens,
InputTokens: totalInputTokens,
OutputTokens: resp.Usage.OutputTokens,
TotalTokens: totalInputTokens + resp.Usage.OutputTokens,
CacheCreationInputTokens: resp.Usage.CacheCreationInputTokens,
}
if resp.Usage.CacheReadInputTokens > 0 {
out.Usage.InputTokensDetails = &ResponsesInputTokensDetails{
@@ -163,6 +164,8 @@ type AnthropicEventToResponsesState struct {
OutputTokens int
CacheReadInputTokens int
CacheCreationInputTokens int
StopReason string
}
// NewAnthropicEventToResponsesState returns an initialised stream state.
@@ -404,7 +407,6 @@ func anthToResHandleContentBlockStop(evt *AnthropicStreamEvent, state *Anthropic
}
func anthToResHandleMessageDelta(evt *AnthropicStreamEvent, state *AnthropicEventToResponsesState) []ResponsesStreamEvent {
// Update usage
if evt.Usage != nil {
state.OutputTokens = evt.Usage.OutputTokens
if evt.Usage.InputTokens > 0 {
@@ -417,6 +419,9 @@ func anthToResHandleMessageDelta(evt *AnthropicStreamEvent, state *AnthropicEven
state.CacheCreationInputTokens = evt.Usage.CacheCreationInputTokens
}
}
if evt.Delta != nil && evt.Delta.StopReason != "" {
state.StopReason = evt.Delta.StopReason
}
return nil
}
@@ -427,15 +432,15 @@ func anthToResHandleMessageStop(state *AnthropicEventToResponsesState) []Respons
}
var events []ResponsesStreamEvent
// Close any open item
events = append(events, closeCurrentResponsesItem(state)...)
// Determine status
status := "completed"
var incompleteDetails *ResponsesIncompleteDetails
if state.StopReason == "max_tokens" {
status = "incomplete"
incompleteDetails = &ResponsesIncompleteDetails{Reason: "max_output_tokens"}
}
// Emit response.completed
events = append(events, makeResponsesCompletedEvent(state, status, incompleteDetails))
state.CompletedSent = true
return events
@@ -497,9 +502,10 @@ func makeResponsesCompletedEvent(
// back to match OpenAI Responses semantics where input_tokens is the total.
totalInputTokens := state.InputTokens + state.CacheReadInputTokens + state.CacheCreationInputTokens
usage := &ResponsesUsage{
InputTokens: totalInputTokens,
OutputTokens: state.OutputTokens,
TotalTokens: totalInputTokens + state.OutputTokens,
InputTokens: totalInputTokens,
OutputTokens: state.OutputTokens,
TotalTokens: totalInputTokens + state.OutputTokens,
CacheCreationInputTokens: state.CacheCreationInputTokens,
}
if state.CacheReadInputTokens > 0 {
usage.InputTokensDetails = &ResponsesInputTokensDetails{
@@ -507,15 +513,20 @@ func makeResponsesCompletedEvent(
}
}
eventType := "response.completed"
if status == "incomplete" {
eventType = "response.incomplete"
}
return ResponsesStreamEvent{
Type: "response.completed",
Type: eventType,
SequenceNumber: seq,
Response: &ResponsesResponse{
ID: state.ResponseID,
Object: "response",
Model: state.Model,
Status: status,
Output: []ResponsesOutput{}, // Simplified; full output tracking would add complexity
Output: []ResponsesOutput{},
Usage: usage,
IncompleteDetails: incompleteDetails,
},
@@ -1,6 +1,8 @@
package apicompat
import (
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"strings"
@@ -33,11 +35,30 @@ func ResponsesToChatCompletionsRequest(req *ResponsesRequest) (*ChatCompletionsR
if req.Reasoning != nil {
out.ReasoningEffort = req.Reasoning.Effort
}
if len(req.Tools) > 0 {
out.Tools = responsesToolsToChatTools(req.Tools)
effectiveTools, err := EffectiveResponsesTools(req)
if err != nil {
return nil, err
}
if len(req.ToolChoice) > 0 {
out.ToolChoice = responsesToolChoiceToChatToolChoice(req.ToolChoice)
if len(effectiveTools) > 0 {
tools, err := responsesToolsToChatTools(effectiveTools)
if err != nil {
return nil, err
}
out.Tools = tools
}
// tools 全部被丢弃(如仅含 web_search/image_generation 等服务端工具)时不再转发
// tool_choice:上游会拒绝 "'tool_choice' is only allowed when 'tools' are specified"。
// 指向被丢弃工具的选择项同理(见 responsesToolChoiceToChatToolChoice)。
if len(out.Tools) > 0 && len(req.ToolChoice) > 0 {
declared := make(map[string]bool, len(out.Tools))
for _, tool := range out.Tools {
if tool.Function != nil {
declared[tool.Function.Name] = true
}
}
if tc := responsesToolChoiceToChatToolChoice(req.ToolChoice, declared); len(tc) > 0 {
out.ToolChoice = tc
}
}
if req.Text != nil {
out.ResponseFormat = responsesTextFormatToChatResponseFormat(req.Text.Format)
@@ -46,6 +67,110 @@ func ResponsesToChatCompletionsRequest(req *ResponsesRequest) (*ChatCompletionsR
return out, nil
}
// EffectiveResponsesTools returns every client-executable tool declared by a
// Responses request. Newer Codex clients place their runtime tools in an
// input item shaped as {"type":"additional_tools","tools":[...]} instead of
// the top-level tools field. Chat-only upstreams must receive both forms.
func EffectiveResponsesTools(req *ResponsesRequest) ([]ResponsesTool, error) {
if req == nil {
return nil, nil
}
tools := append([]ResponsesTool(nil), req.Tools...)
inputRaw := bytesTrimSpace(req.Input)
if len(inputRaw) == 0 || string(inputRaw) == "null" || inputRaw[0] != '[' {
return tools, nil
}
var items []json.RawMessage
if err := json.Unmarshal(inputRaw, &items); err != nil {
return nil, fmt.Errorf("parse responses input for additional tools: %w", err)
}
for _, raw := range items {
raw = bytesTrimSpace(raw)
if len(raw) == 0 || raw[0] != '{' {
continue
}
var item struct {
Type string `json:"type"`
Tools []ResponsesTool `json:"tools"`
}
if err := json.Unmarshal(raw, &item); err != nil {
return nil, fmt.Errorf("parse responses additional tools item: %w", err)
}
if item.Type == "additional_tools" {
tools = append(tools, item.Tools...)
}
}
return tools, nil
}
// CustomToolNames 收集 Responses 请求中 custom/freeform 工具的名字。chat 桥回程时
// 需要据此把模型对这些工具的调用还原为 custom_tool_call 项(codex 只按该类型路由)。
func CustomToolNames(tools []ResponsesTool) map[string]bool {
var out map[string]bool
for _, tool := range tools {
if tool.Type == "custom" && tool.Name != "" {
if out == nil {
out = make(map[string]bool)
}
out[tool.Name] = true
}
}
return out
}
// NamespacedToolName 记录 namespace 子工具的原始归属(命名空间 + 裸子工具名)。
type NamespacedToolName struct {
Namespace string
Name string
}
// NamespaceToolNames 收集 Responses 请求中 namespace 子工具的摊平名 →(namespace,
// 子工具名)映射。chat 桥回程时需据此把模型对摊平工具的调用还原为带 namespace 字段
// 的 function_call 项:codex 按 namespace+name 路由,平铺名会被判为 unsupported
// call;摊平名超长时带截断哈希(见 flattenNamespaceToolName),无法按字符串切分还原。
// 摊平名撞名的请求已在转换阶段被显式拒绝(见 namespaceChildrenToChatTools),
// 此处映射不存在歧义。
func NamespaceToolNames(tools []ResponsesTool) map[string]NamespacedToolName {
var out map[string]NamespacedToolName
for _, tool := range tools {
if tool.Type != "namespace" || tool.Name == "" {
continue
}
children := tool.Tools
if len(children) == 0 {
children = tool.Children
}
for _, child := range children {
if child.Type != "function" || child.Name == "" {
continue
}
if out == nil {
out = make(map[string]NamespacedToolName)
}
out[flattenNamespaceToolName(tool.Name, child.Name)] = NamespacedToolName{
Namespace: tool.Name,
Name: child.Name,
}
}
}
return out
}
// HasToolSearchTool 判断 Responses 请求是否声明了 tool_search 服务端工具。chat 桥
// 回程时需据此把模型对代理工具的调用还原为 tool_search_call 项:codex 只在该项类型
// 且 execution=client 时执行 tool search,同名 function_call 会因 payload 不匹配
// 触发 fatal 中止整个 turn。
func HasToolSearchTool(tools []ResponsesTool) bool {
for _, tool := range tools {
if tool.Type == "tool_search" {
return true
}
}
return false
}
// responsesInputToChatMessages converts a Responses request's instructions +
// input[] into Chat Completions messages. It is a three-stage pipeline:
//
@@ -133,33 +258,68 @@ func buildChatMessagesFromItems(messages []ChatMessage, rawItems []json.RawMessa
if strings.TrimSpace(arguments) == "" {
arguments = "{}"
}
name := rawString(item["name"])
// namespace 子工具的历史调用带 namespace 字段,需与请求方向的摊平
// 命名(namespaceChildrenToChatTools)保持一致。
if ns := rawString(item["namespace"]); ns != "" {
name = flattenNamespaceToolName(ns, name)
}
toolCall := ChatToolCall{
ID: rawString(item["call_id"]),
Type: "function",
Function: ChatFunctionCall{
Name: name,
Arguments: arguments,
},
}
messages = appendAssistantToolCall(messages, toolCall, pendingReasoning)
pendingReasoning = ""
continue
case "tool_search_call":
// tool_search 调用的 arguments 是 JSON 对象(如 {"query": ...}),
// 原文即为降级 function 调用的 arguments 字符串。
arguments := strings.TrimSpace(string(bytesTrimSpace(item["arguments"])))
if s := rawString(item["arguments"]); s != "" {
arguments = s
}
if arguments == "" || arguments == "null" {
arguments = "{}"
}
toolCall := ChatToolCall{
ID: rawString(item["call_id"]),
Type: "function",
Function: ChatFunctionCall{
Name: toolSearchProxyName,
Arguments: arguments,
},
}
messages = appendAssistantToolCall(messages, toolCall, pendingReasoning)
pendingReasoning = ""
continue
case "custom_tool_call":
// custom/freeform 工具的历史调用:input 自由文本包进降级 function 工具
// 的 {"input": ...} 参数,与请求方向的工具降级(customToolInputSchema)
// 保持一致,模型才能把历史与当前工具定义对上。
arguments, _ := json.Marshal(map[string]string{"input": rawString(item["input"])})
toolCall := ChatToolCall{
ID: rawString(item["call_id"]),
Type: "function",
Function: ChatFunctionCall{
Name: rawString(item["name"]),
Arguments: arguments,
Arguments: string(arguments),
},
}
// Parallel tool calls arrive as consecutive function_call items and
// must share one assistant message; the matching tool replies then
// follow it. Merge into the immediately preceding assistant message.
if n := len(messages); n > 0 && messages[n-1].Role == "assistant" {
messages[n-1].ToolCalls = append(messages[n-1].ToolCalls, toolCall)
if messages[n-1].ReasoningContent == "" {
messages[n-1].ReasoningContent = pendingReasoning
}
} else {
messages = append(messages, ChatMessage{
Role: "assistant",
ToolCalls: []ChatToolCall{toolCall},
ReasoningContent: pendingReasoning,
})
}
messages = appendAssistantToolCall(messages, toolCall, pendingReasoning)
pendingReasoning = ""
continue
case "function_call_output":
content, _ := json.Marshal(rawString(item["output"]))
case "function_call_output", "custom_tool_call_output", "tool_search_output":
outputRaw := bytesTrimSpace(item["output"])
outputText := rawString(outputRaw)
if outputText == "" && len(outputRaw) > 0 && string(outputRaw) != "null" && string(outputRaw) != `""` {
// 对象/数组形式的输出(如 tool_search 的结果列表)整体字符串化。
outputText = string(outputRaw)
}
content, _ := json.Marshal(outputText)
messages = append(messages, ChatMessage{
Role: "tool",
ToolCallID: rawString(item["call_id"]),
@@ -184,9 +344,9 @@ func buildChatMessagesFromItems(messages []ChatMessage, rawItems []json.RawMessa
// Only genuine message items become chat messages. Codex emits other
// Responses item types with no Chat equivalent (web_search_call,
// local_shell_call, custom tool calls, file_search_call, ...). Converting
// them via the generic path would insert a spurious message between an
// assistant tool_calls message and its tool reply, which DeepSeek rejects
// local_shell_call, file_search_call, ...). Converting them via the
// generic path would insert a spurious message between an assistant
// tool_calls message and its tool reply, which DeepSeek rejects
// ("insufficient tool messages following tool_calls message"). Skip them.
if itemType != "" && itemType != "message" {
pendingReasoning = ""
@@ -213,6 +373,25 @@ func buildChatMessagesFromItems(messages []ChatMessage, rawItems []json.RawMessa
return messages, nil
}
// appendAssistantToolCall merges a tool call into the chat message list.
// Parallel tool calls arrive as consecutive *_call items and must share one
// assistant message; the matching tool replies then follow it. Merge into the
// immediately preceding assistant message.
func appendAssistantToolCall(messages []ChatMessage, toolCall ChatToolCall, pendingReasoning string) []ChatMessage {
if n := len(messages); n > 0 && messages[n-1].Role == "assistant" {
messages[n-1].ToolCalls = append(messages[n-1].ToolCalls, toolCall)
if messages[n-1].ReasoningContent == "" {
messages[n-1].ReasoningContent = pendingReasoning
}
return messages
}
return append(messages, ChatMessage{
Role: "assistant",
ToolCalls: []ChatToolCall{toolCall},
ReasoningContent: pendingReasoning,
})
}
// normalizeChatMessages is the single place that enforces the tool-call
// invariant the DeepSeek / OpenAI Chat Completions schema requires: an assistant
// message with tool_calls must be immediately followed by one tool message per
@@ -427,39 +606,187 @@ func chatContentFromSingleResponsesPart(partType string, part map[string]json.Ra
}
}
func responsesToolsToChatTools(tools []ResponsesTool) []ChatTool {
// customToolInputSchema 是 custom/freeform 工具降级为 function 工具时的参数 schema。
// chat 协议无法表达 custom 工具的自由文本输入(及其 grammar 约束),退化为单一
// input 字符串参数;回程时再从 arguments 的 input 字段还原(见
// extractCustomToolCallInput)。
const customToolInputSchema = `{"type":"object","properties":{"input":{"type":"string","description":"The raw input for this tool, passed through verbatim."}},"required":["input"]}`
func responsesToolsToChatTools(tools []ResponsesTool) ([]ChatTool, error) {
// 顶层 function/custom 工具名集合:namespace 子工具摊平后与其撞名时,chat
// 上游无法按 namespace 区分调用归属。这类请求在原生 Responses 上游是合法的
// (按 namespace+name 路由),歧义由摊平转换制造且无法消除,必须显式拒绝,
// 不能静默降级(重复声明发给上游、回程还原到错误工具)。
topLevel := make(map[string]bool)
for _, tool := range tools {
if (tool.Type == "function" || tool.Type == "custom") && tool.Name != "" {
topLevel[tool.Name] = true
}
}
flatOwner := make(map[string]NamespacedToolName)
toolSearchDeclared := false
out := make([]ChatTool, 0, len(tools))
for _, tool := range tools {
if tool.Type != "function" {
switch tool.Type {
case "function":
out = append(out, ChatTool{
Type: "function",
Function: &ChatFunction{
Name: tool.Name,
Description: tool.Description,
Parameters: tool.Parameters,
Strict: tool.Strict,
},
})
case "custom":
// codex 0.14x 的核心执行工具 exec 即为 custom 类型;丢弃它会让模型
// 无法执行任何命令,必须降级为 function 工具透传。
out = append(out, ChatTool{
Type: "function",
Function: &ChatFunction{
Name: tool.Name,
Description: tool.Description,
Parameters: json.RawMessage(customToolInputSchema),
},
})
case "tool_search":
// 代理不能改名(codex 的模型侧按 tool_search 这个名字调用),与客户端
// 声明的同名工具无法区分——回程会把普通工具的调用劫持成 tool_search_call,
// 必须显式拒绝;重复声明 type=tool_search 去重即可。
if topLevel[toolSearchProxyName] {
return nil, fmt.Errorf("built-in tool_search conflicts with a declared tool named %q; this upstream cannot disambiguate them, rename the tool", toolSearchProxyName)
}
if toolSearchDeclared {
continue
}
toolSearchDeclared = true
out = append(out, toolSearchProxyChatTool())
case "namespace":
flattened, err := namespaceChildrenToChatTools(tool, topLevel, flatOwner)
if err != nil {
return nil, err
}
out = append(out, flattened...)
}
// 其余类型(web_search、image_generation 等服务端工具)在 chat 上游没有
// 对应能力,维持丢弃。
}
return out, nil
}
// toolSearchProxyName 是 tool_search 服务端工具降级后的 function 工具名。模型对
// 它的调用以同名 function_call 原样回传,由 codex 端路由。
const toolSearchProxyName = "tool_search"
const toolSearchProxySchema = `{"type":"object","properties":{"query":{"type":"string","description":"Search query for tools or connectors to load."},"limit":{"type":"integer","description":"Maximum number of tool groups to return."}},"required":["query"]}`
func toolSearchProxyChatTool() ChatTool {
return ChatTool{
Type: "function",
Function: &ChatFunction{
Name: toolSearchProxyName,
Description: "Search and load Codex tools, plugins, connectors, and MCP namespaces for the current task.",
Parameters: json.RawMessage(toolSearchProxySchema),
},
}
}
// namespaceChildrenToChatTools 将 namespace 工具的子 function 工具摊平为顶层
// function 工具,名字加 "<namespace>__" 前缀。摊平名与顶层工具或其他 namespace
// 撞名时返回错误(歧义不可消除,显式拒绝);同一 (namespace, 子工具) 的重复声明
// 去重后不算冲突。
func namespaceChildrenToChatTools(tool ResponsesTool, topLevel map[string]bool, flatOwner map[string]NamespacedToolName) ([]ChatTool, error) {
if tool.Name == "" {
return nil, nil
}
children := tool.Tools
if len(children) == 0 {
children = tool.Children
}
var out []ChatTool
for _, child := range children {
if child.Type != "function" || child.Name == "" {
continue
}
flat := flattenNamespaceToolName(tool.Name, child.Name)
entry := NamespacedToolName{Namespace: tool.Name, Name: child.Name}
if topLevel[flat] {
return nil, fmt.Errorf("namespace tool %q/%q flattens to %q which conflicts with a top-level tool of the same name; this upstream cannot disambiguate them, rename one of the tools", tool.Name, child.Name, flat)
}
if prev, ok := flatOwner[flat]; ok {
if prev == entry {
continue
}
return nil, fmt.Errorf("namespace tools %q/%q and %q/%q both flatten to %q; this upstream cannot disambiguate them, rename one of the tools", prev.Namespace, prev.Name, tool.Name, child.Name, flat)
}
flatOwner[flat] = entry
out = append(out, ChatTool{
Type: "function",
Function: &ChatFunction{
Name: tool.Name,
Description: tool.Description,
Parameters: tool.Parameters,
Strict: tool.Strict,
Name: flat,
Description: child.Description,
Parameters: child.Parameters,
Strict: child.Strict,
},
})
}
return out
return out, nil
}
func responsesToolChoiceToChatToolChoice(raw json.RawMessage) json.RawMessage {
// chatToolNameMaxLen 是 Chat Completions function 工具名的通用长度上限。
const chatToolNameMaxLen = 64
// flattenNamespaceToolName 生成 namespace 子工具的摊平名;超长时截断并追加
// sha256 短哈希保证唯一性。
func flattenNamespaceToolName(namespace, name string) string {
full := namespace + "__" + name
if len(full) <= chatToolNameMaxLen {
return full
}
sum := sha256.Sum256([]byte(full))
suffix := "__" + hex.EncodeToString(sum[:4])
prefixLen := chatToolNameMaxLen - len(suffix)
var prefix strings.Builder
for _, ch := range full {
if prefix.Len()+len(string(ch)) > prefixLen {
break
}
_, _ = prefix.WriteRune(ch)
}
return prefix.String() + suffix
}
// responsesToolChoiceToChatToolChoice 把 Responses 的 tool_choice 转为 chat 形态。
// declared 是转换后实际声明的 chat 工具名集合:具名选择项仅在目标工具幸存时转发,
// 服务端工具(web_search 等)的选择项随工具本身丢弃——指向未声明工具的 tool_choice
// 会被 chat 上游 400 拒绝。返回 nil 表示丢弃 tool_choice。
func responsesToolChoiceToChatToolChoice(raw json.RawMessage, declared map[string]bool) json.RawMessage {
var choice map[string]json.RawMessage
if err := json.Unmarshal(raw, &choice); err != nil {
// "auto"/"none"/"required" 等字符串形式原样转发。
return raw
}
if rawString(choice["type"]) != "function" {
return raw
var name string
switch rawString(choice["type"]) {
case "tool_search":
// tool_search 未被丢弃而是降级为同名 function 代理(见
// responsesToolsToChatTools),强制选择它同样降级为 function 选择,
// 静默丢弃会把强制搜索退化为自动选择。
name = toolSearchProxyName
case "function", "custom":
// custom 工具已降级为 function 工具,指向它的 tool_choice 同样按 function 转换。
name = rawString(choice["name"])
if name == "" {
name = rawNestedString(choice["function"], "name")
}
if name == "" {
return raw
}
default:
return nil
}
name := rawString(choice["name"])
if name == "" {
name = rawNestedString(choice["function"], "name")
}
if name == "" {
return raw
if !declared[name] {
return nil
}
out, err := json.Marshal(map[string]any{
"type": "function",
@@ -473,9 +800,38 @@ func responsesToolChoiceToChatToolChoice(raw json.RawMessage) json.RawMessage {
return out
}
// extractCustomToolCallInput 从降级 function 调用的 arguments 中还原 custom 工具的
// 自由文本输入:优先取 {"input": "..."} 的 input 字段;模型未按 schema 输出时原样
// 回传,交由客户端校验、模型重试。
func extractCustomToolCallInput(arguments string) string {
trimmed := strings.TrimSpace(arguments)
if trimmed == "" {
return ""
}
var obj map[string]json.RawMessage
if err := json.Unmarshal([]byte(trimmed), &obj); err != nil {
return trimmed
}
if raw, ok := obj["input"]; ok {
var s string
if err := json.Unmarshal(raw, &s); err == nil {
return s
}
return trimmed
}
if len(obj) == 0 {
return ""
}
return trimmed
}
// ChatCompletionsResponseToResponses converts a non-streaming Chat Completions
// response into a Responses API response.
func ChatCompletionsResponseToResponses(resp *ChatCompletionsResponse, model string) *ResponsesResponse {
// response into a Responses API response. customTools 是客户端请求中 custom 工具
// 的名字集合(见 CustomToolNames),命中的调用会还原为 custom_tool_call 项;
// toolSearch 表示客户端声明了 tool_search 工具(见 HasToolSearchTool),代理工具
// 的调用会还原为 tool_search_call 项;namespaceTools 是 namespace 子工具的摊平名
// 映射(见 NamespaceToolNames),命中的调用还原为带 namespace 字段的 function_call 项。
func ChatCompletionsResponseToResponses(resp *ChatCompletionsResponse, model string, customTools map[string]bool, toolSearch bool, namespaceTools map[string]NamespacedToolName) *ResponsesResponse {
id := ""
if resp != nil {
id = resp.ID
@@ -500,7 +856,7 @@ func ChatCompletionsResponseToResponses(resp *ChatCompletionsResponse, model str
if len(resp.Choices) > 0 {
choice := resp.Choices[0]
out.Output = chatMessageToResponsesOutput(choice.Message)
out.Output = chatMessageToResponsesOutput(choice.Message, customTools, toolSearch, namespaceTools)
if choice.FinishReason == "length" {
out.Status = "incomplete"
out.IncompleteDetails = &ResponsesIncompleteDetails{Reason: "max_output_tokens"}
@@ -515,7 +871,7 @@ func ChatCompletionsResponseToResponses(resp *ChatCompletionsResponse, model str
return out
}
func chatMessageToResponsesOutput(message ChatMessage) []ResponsesOutput {
func chatMessageToResponsesOutput(message ChatMessage, customTools map[string]bool, toolSearch bool, namespaceTools map[string]NamespacedToolName) []ResponsesOutput {
var outputs []ResponsesOutput
if message.ReasoningContent != "" {
outputs = append(outputs, ResponsesOutput{
@@ -550,6 +906,39 @@ func chatMessageToResponsesOutput(message ChatMessage) []ResponsesOutput {
if strings.TrimSpace(arguments) == "" {
arguments = "{}"
}
if customTools[toolCall.Function.Name] {
outputs = append(outputs, ResponsesOutput{
Type: "custom_tool_call",
ID: generateItemID(),
CallID: toolCall.ID,
Name: toolCall.Function.Name,
Input: extractCustomToolCallInput(arguments),
Status: "completed",
})
continue
}
if toolSearch && toolCall.Function.Name == toolSearchProxyName {
outputs = append(outputs, ResponsesOutput{
Type: "tool_search_call",
ID: generateItemID(),
CallID: toolCall.ID,
Arguments: arguments,
Status: "completed",
})
continue
}
if ns, ok := namespaceTools[toolCall.Function.Name]; ok {
outputs = append(outputs, ResponsesOutput{
Type: "function_call",
ID: generateItemID(),
CallID: toolCall.ID,
Name: ns.Name,
Namespace: ns.Namespace,
Arguments: arguments,
Status: "completed",
})
continue
}
outputs = append(outputs, ResponsesOutput{
Type: "function_call",
ID: generateItemID(),
@@ -563,6 +952,21 @@ func chatMessageToResponsesOutput(message ChatMessage) []ResponsesOutput {
return outputs
}
// toolSearchCallArgumentsJSON 把降级 function 调用累积的 arguments 字符串还原为
// tool_search_call 线上要求的 JSON 对象;模型未按 schema 输出(非法 JSON)时按
// 字符串值兜底,交由 codex 解析报错后让模型重试。
func toolSearchCallArgumentsJSON(arguments string) json.RawMessage {
trimmed := strings.TrimSpace(arguments)
if trimmed == "" {
return json.RawMessage(`{}`)
}
if json.Valid([]byte(trimmed)) {
return json.RawMessage(trimmed)
}
fallback, _ := json.Marshal(arguments)
return fallback
}
func emptyResponsesMessageOutput() ResponsesOutput {
return ResponsesOutput{
Type: "message",
@@ -662,6 +1066,35 @@ type ChatCompletionsToResponsesStreamState struct {
ToolItemIDs map[int]string
ToolOutputIndex map[int]int
// CustomTools 是客户端请求中 custom/freeform 工具的名字集合(见
// CustomToolNames)。命中的调用按 custom_tool_call 生命周期下发,codex 才能
// 路由回它注册的 custom 工具。
CustomTools map[string]bool
// ToolSearchDeclared 表示客户端请求声明了 tool_search 工具(见
// HasToolSearchTool)。命中的代理调用按 tool_search_call 项还原,codex 只按
// 该项类型(且 execution=client)执行 tool search。
ToolSearchDeclared bool
// NamespaceTools 是 namespace 子工具的摊平名 → 原始归属映射(见
// NamespaceToolNames)。命中的调用还原为带 namespace 字段的 function_call 项,
// codex 按 namespace+name 路由。
NamespaceTools map[string]NamespacedToolName
// toolIsCustom 记录每个工具调用宣告时的类型判定,保证 added/done 事件的
// 项类型一致。
toolIsCustom map[int]bool
// toolIsToolSearch 记录工具调用是否判定为 tool_search 代理调用。
toolIsToolSearch map[int]bool
// toolNamespace 记录工具调用宣告时命中的 namespace 归属(见 NamespaceTools)。
toolNamespace map[int]NamespacedToolName
// toolAnnounced 记录 output_item.added 是否已发出。存在 custom 工具且名字
// 尚未到达时延迟宣告,待名字可判定类型后再补发(见 announceChatToolItem)。
toolAnnounced map[int]bool
FinishReason string
Usage *ResponsesUsage
}
@@ -669,12 +1102,16 @@ type ChatCompletionsToResponsesStreamState struct {
// NewChatCompletionsToResponsesStreamState returns an initialized stream state.
func NewChatCompletionsToResponsesStreamState(model string) *ChatCompletionsToResponsesStreamState {
return &ChatCompletionsToResponsesStreamState{
ResponseID: generateResponsesID(),
Model: model,
Created: time.Now().Unix(),
ToolCalls: make(map[int]*ChatToolCall),
ToolItemIDs: make(map[int]string),
ToolOutputIndex: make(map[int]int),
ResponseID: generateResponsesID(),
Model: model,
Created: time.Now().Unix(),
ToolCalls: make(map[int]*ChatToolCall),
ToolItemIDs: make(map[int]string),
ToolOutputIndex: make(map[int]int),
toolIsCustom: make(map[int]bool),
toolIsToolSearch: make(map[int]bool),
toolNamespace: make(map[int]NamespacedToolName),
toolAnnounced: make(map[int]bool),
}
}
@@ -758,19 +1195,8 @@ func ChatCompletionsChunkToResponsesEvents(
copyCall.Function.Arguments = ""
state.ToolCalls[idx] = &copyCall
stored = &copyCall
itemID := generateItemID()
state.ToolItemIDs[idx] = itemID
state.ToolItemIDs[idx] = generateItemID()
state.ToolOutputIndex[idx] = state.allocOutputIndex()
events = append(events, chatToResponsesEvent(state, "response.output_item.added", &ResponsesStreamEvent{
OutputIndex: state.ToolOutputIndex[idx],
Item: &ResponsesOutput{
Type: "function_call",
ID: itemID,
CallID: stored.ID,
Name: stored.Function.Name,
Status: "in_progress",
},
}))
} else {
if toolCall.ID != "" {
stored.ID = toolCall.ID
@@ -779,15 +1205,22 @@ func ChatCompletionsChunkToResponsesEvents(
stored.Function.Name = toolCall.Function.Name
}
}
events = append(events, announceChatToolItem(state, idx, stored, false)...)
if toolCall.Function.Arguments != "" {
stored.Function.Arguments += toolCall.Function.Arguments
events = append(events, chatToResponsesEvent(state, "response.function_call_arguments.delta", &ResponsesStreamEvent{
OutputIndex: state.ToolOutputIndex[idx],
ItemID: state.ToolItemIDs[idx],
Delta: toolCall.Function.Arguments,
CallID: stored.ID,
Name: stored.Function.Name,
}))
// 未宣告(名字未到)时仅累积,宣告时统一补发;custom 调用的
// arguments 是包裹 input 的 JSON 片段,无法增量还原为自由文本
// 输入,缓冲整份 arguments 收尾时一次性下发(见 closeChatToolItems);
// tool_search 调用同样收尾时随 output_item.done 全量下发。
if state.toolAnnounced[idx] && !state.toolIsCustom[idx] && !state.toolIsToolSearch[idx] {
events = append(events, chatToResponsesEvent(state, "response.function_call_arguments.delta", &ResponsesStreamEvent{
OutputIndex: state.ToolOutputIndex[idx],
ItemID: state.ToolItemIDs[idx],
Delta: toolCall.Function.Arguments,
CallID: stored.ID,
Name: stored.Function.Name,
}))
}
}
}
if choice.FinishReason != nil && *choice.FinishReason != "" {
@@ -998,6 +1431,64 @@ func ensureChatToResponsesTextPart(state *ChatCompletionsToResponsesStreamState)
})}
}
// announceChatToolItem 在类型可判定时发出工具调用的 output_item.added。custom
// 工具的判定依赖名字:名字未到且请求里存在 custom 工具时延迟宣告,避免 added/done
// 的项类型不一致;force 用于流收尾,名字始终未到时按 function_call 兜底。
func announceChatToolItem(
state *ChatCompletionsToResponsesStreamState,
idx int,
stored *ChatToolCall,
force bool,
) []ResponsesStreamEvent {
if state.toolAnnounced[idx] {
return nil
}
if !force && stored.Function.Name == "" && (len(state.CustomTools) > 0 || state.ToolSearchDeclared || len(state.NamespaceTools) > 0) {
return nil
}
state.toolAnnounced[idx] = true
isCustom := state.CustomTools[stored.Function.Name]
isToolSearch := !isCustom && state.ToolSearchDeclared && stored.Function.Name == toolSearchProxyName
state.toolIsCustom[idx] = isCustom
state.toolIsToolSearch[idx] = isToolSearch
itemType := "function_call"
if isCustom {
itemType = "custom_tool_call"
}
if isToolSearch {
itemType = "tool_search_call"
}
// namespace 子工具的调用仍按 function_call 生命周期下发,但 added/done 项要
// 还原为裸子工具名 + namespace 字段(codex 按 namespace+name 路由)。
itemName, itemNamespace := stored.Function.Name, ""
if ns, ok := state.NamespaceTools[stored.Function.Name]; ok && !isCustom && !isToolSearch {
state.toolNamespace[idx] = ns
itemName, itemNamespace = ns.Name, ns.Namespace
}
events := []ResponsesStreamEvent{chatToResponsesEvent(state, "response.output_item.added", &ResponsesStreamEvent{
OutputIndex: state.ToolOutputIndex[idx],
Item: &ResponsesOutput{
Type: itemType,
ID: state.ToolItemIDs[idx],
CallID: stored.ID,
Name: itemName,
Namespace: itemNamespace,
Status: "in_progress",
},
})}
// 迟到宣告时补发已累积的参数增量(custom/tool_search 的输入收尾统一下发,不补发)。
if !isCustom && !isToolSearch && stored.Function.Arguments != "" {
events = append(events, chatToResponsesEvent(state, "response.function_call_arguments.delta", &ResponsesStreamEvent{
OutputIndex: state.ToolOutputIndex[idx],
ItemID: state.ToolItemIDs[idx],
Delta: stored.Function.Arguments,
CallID: stored.ID,
Name: stored.Function.Name,
}))
}
return events
}
// closeChatToolItems emits function_call_arguments.done + output_item.done for
// every tool call opened during the stream, carrying the full call_id/name/
// arguments so codex can deserialize and execute the call. Mirrors cc-switch's
@@ -1016,17 +1507,72 @@ func closeChatToolItems(state *ChatCompletionsToResponsesStreamState) []Response
if !opened {
continue
}
// 名字始终未到导致尚未宣告的调用,收尾前按最终名字兜底宣告。
events = append(events, announceChatToolItem(state, i, toolCall, true)...)
arguments := toolCall.Function.Arguments
if strings.TrimSpace(arguments) == "" {
arguments = "{}"
}
outputIndex := state.ToolOutputIndex[i]
if state.toolIsCustom[i] {
// custom 调用按 custom_tool_call 生命周期收尾:input 在此处一次性下发
// (流中不产出增量,见 ChatCompletionsChunkToResponsesEvents)。
input := extractCustomToolCallInput(arguments)
if input != "" {
events = append(events, chatToResponsesEvent(state, "response.custom_tool_call_input.delta", &ResponsesStreamEvent{
OutputIndex: outputIndex,
ItemID: itemID,
Delta: input,
}))
}
events = append(events,
chatToResponsesEvent(state, "response.custom_tool_call_input.done", &ResponsesStreamEvent{
OutputIndex: outputIndex,
ItemID: itemID,
CallID: toolCall.ID,
Name: toolCall.Function.Name,
Input: input,
}),
chatToResponsesEvent(state, "response.output_item.done", &ResponsesStreamEvent{
OutputIndex: outputIndex,
Item: &ResponsesOutput{
Type: "custom_tool_call",
ID: itemID,
CallID: toolCall.ID,
Name: toolCall.Function.Name,
Input: input,
Status: "completed",
},
}),
)
continue
}
if state.toolIsToolSearch[i] {
// tool_search 调用按 tool_search_call 项收尾:codex 从 output_item.done
// 物化该调用(无参数增量事件),arguments 全量随项下发。
events = append(events, chatToResponsesEvent(state, "response.output_item.done", &ResponsesStreamEvent{
OutputIndex: outputIndex,
Item: &ResponsesOutput{
Type: "tool_search_call",
ID: itemID,
CallID: toolCall.ID,
Arguments: arguments,
Status: "completed",
},
}))
continue
}
// namespace 子工具调用在宣告时已记录归属,收尾项同样带还原名与 namespace。
name, namespace := toolCall.Function.Name, ""
if ns, ok := state.toolNamespace[i]; ok {
name, namespace = ns.Name, ns.Namespace
}
events = append(events,
chatToResponsesEvent(state, "response.function_call_arguments.done", &ResponsesStreamEvent{
OutputIndex: outputIndex,
ItemID: itemID,
CallID: toolCall.ID,
Name: toolCall.Function.Name,
Name: name,
Arguments: arguments,
}),
chatToResponsesEvent(state, "response.output_item.done", &ResponsesStreamEvent{
@@ -1035,7 +1581,8 @@ func closeChatToolItems(state *ChatCompletionsToResponsesStreamState) []Response
Type: "function_call",
ID: itemID,
CallID: toolCall.ID,
Name: toolCall.Function.Name,
Name: name,
Namespace: namespace,
Arguments: arguments,
Status: "completed",
},
@@ -1078,11 +1625,37 @@ func (state *ChatCompletionsToResponsesStreamState) chatOutput() []ResponsesOutp
if strings.TrimSpace(arguments) == "" {
arguments = "{}"
}
if state.toolIsCustom[i] {
outputs = append(outputs, ResponsesOutput{
Type: "custom_tool_call",
ID: generateItemID(),
CallID: toolCall.ID,
Name: toolCall.Function.Name,
Input: extractCustomToolCallInput(arguments),
Status: "completed",
})
continue
}
if state.toolIsToolSearch[i] {
outputs = append(outputs, ResponsesOutput{
Type: "tool_search_call",
ID: generateItemID(),
CallID: toolCall.ID,
Arguments: arguments,
Status: "completed",
})
continue
}
name, namespace := toolCall.Function.Name, ""
if ns, ok := state.toolNamespace[i]; ok {
name, namespace = ns.Name, ns.Namespace
}
outputs = append(outputs, ResponsesOutput{
Type: "function_call",
ID: generateItemID(),
CallID: toolCall.ID,
Name: toolCall.Function.Name,
Name: name,
Namespace: namespace,
Arguments: arguments,
Status: "completed",
})
@@ -0,0 +1,923 @@
package apicompat
// custom/freeform 工具(如 Codex 0.14x 的 exec)在 responses→chat 桥上的双向转换。
// 背景:Codex 的核心命令执行工具 exec 是 type=custom(输入为自由文本),此前被
// responsesToolsToChatTools 丢弃,导致模型工具列表中没有 exec、无法执行任何命令。
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestResponsesToChatCompletionsRequest_CustomToolBecomesFunctionTool(t *testing.T) {
req := &ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"run dir"`),
Tools: []ResponsesTool{
{Type: "custom", Name: "exec", Description: "Run JavaScript code"},
{Type: "function", Name: "wait", Parameters: json.RawMessage(`{"type":"object","properties":{}}`)},
},
}
out, err := ResponsesToChatCompletionsRequest(req)
require.NoError(t, err)
require.Len(t, out.Tools, 2)
assert.Equal(t, "function", out.Tools[0].Type)
assert.Equal(t, "exec", out.Tools[0].Function.Name)
assert.Equal(t, "Run JavaScript code", out.Tools[0].Function.Description)
assert.JSONEq(t, customToolInputSchema, string(out.Tools[0].Function.Parameters))
assert.Equal(t, "wait", out.Tools[1].Function.Name)
}
func TestResponsesToChatCompletionsRequest_AdditionalToolsItem(t *testing.T) {
req := &ResponsesRequest{
Model: "gpt-test",
Input: json.RawMessage(`[
{"type":"additional_tools","role":"developer","tools":[
{"type":"custom","name":"exec","description":"Run PowerShell","format":{"type":"text"}},
{"type":"function","name":"wait","parameters":{"type":"object","properties":{}}},
{"type":"namespace","name":"collaboration","tools":[
{"type":"function","name":"send_message","parameters":{"type":"object","properties":{}}}
]}
]},
{"type":"message","role":"user","content":[{"type":"input_text","text":"run Get-Location"}]}
]`),
ToolChoice: json.RawMessage(`"auto"`),
}
effective, err := EffectiveResponsesTools(req)
require.NoError(t, err)
require.Len(t, effective, 3)
assert.True(t, CustomToolNames(effective)["exec"])
assert.Equal(t, NamespacedToolName{Namespace: "collaboration", Name: "send_message"}, NamespaceToolNames(effective)["collaboration__send_message"])
out, err := ResponsesToChatCompletionsRequest(req)
require.NoError(t, err)
require.Len(t, out.Tools, 3)
assert.Equal(t, "exec", out.Tools[0].Function.Name)
assert.Equal(t, "wait", out.Tools[1].Function.Name)
assert.Equal(t, "collaboration__send_message", out.Tools[2].Function.Name)
assert.JSONEq(t, `"auto"`, string(out.ToolChoice))
require.Len(t, out.Messages, 1, "additional_tools must not become a chat message")
assert.Equal(t, "user", out.Messages[0].Role)
}
func TestEffectiveResponsesTools_SkipsStringInputItems(t *testing.T) {
req := &ResponsesRequest{
Input: json.RawMessage(`["plain input",{"type":"additional_tools","tools":[{"type":"custom","name":"exec"}]}]`),
}
tools, err := EffectiveResponsesTools(req)
require.NoError(t, err)
require.Len(t, tools, 1)
assert.Equal(t, "exec", tools[0].Name)
}
func TestResponsesToChatCompletionsRequest_DropsToolChoiceWhenNoConvertibleTools(t *testing.T) {
req := &ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{
{Type: "web_search"},
{Type: "image_generation"},
},
ToolChoice: json.RawMessage(`"auto"`),
}
out, err := ResponsesToChatCompletionsRequest(req)
require.NoError(t, err)
assert.Empty(t, out.Tools)
assert.Empty(t, out.ToolChoice, "tools 为空时转发 tool_choice 会被上游 400 拒绝")
}
func TestResponsesToChatCompletionsRequest_CustomToolChoiceMapsToFunctionChoice(t *testing.T) {
req := &ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"run dir"`),
Tools: []ResponsesTool{{Type: "custom", Name: "exec"}},
ToolChoice: json.RawMessage(`{"type":"custom","name":"exec"}`),
}
out, err := ResponsesToChatCompletionsRequest(req)
require.NoError(t, err)
assert.JSONEq(t, `{"type":"function","function":{"name":"exec"}}`, string(out.ToolChoice))
}
func TestResponsesInputToChatMessages_CustomToolCallHistory(t *testing.T) {
input := json.RawMessage(`[
{"role":"user","content":"list files"},
{"type":"custom_tool_call","call_id":"call_1","name":"exec","input":"dir"},
{"type":"custom_tool_call_output","call_id":"call_1","output":"main.go"}
]`)
messages, err := responsesInputToChatMessages("", input)
require.NoError(t, err)
require.Len(t, messages, 3)
assert.Equal(t, []string{"user", "assistant", "tool"}, chatMessageRoles(messages))
require.Len(t, messages[1].ToolCalls, 1)
toolCall := messages[1].ToolCalls[0]
assert.Equal(t, "call_1", toolCall.ID)
assert.Equal(t, "exec", toolCall.Function.Name)
assert.JSONEq(t, `{"input":"dir"}`, toolCall.Function.Arguments)
assert.Equal(t, "call_1", messages[2].ToolCallID)
assert.JSONEq(t, `"main.go"`, string(messages[2].Content))
}
func TestChatCompletionsResponseToResponses_CustomToolCallOutputItem(t *testing.T) {
resp := &ChatCompletionsResponse{
ID: "cc-1",
Choices: []ChatChoice{{
Message: ChatMessage{
Role: "assistant",
ToolCalls: []ChatToolCall{
{ID: "call_1", Function: ChatFunctionCall{Name: "exec", Arguments: `{"input": "dir"}`}},
{ID: "call_2", Function: ChatFunctionCall{Name: "wait", Arguments: `{"cell_id": 3}`}},
},
},
}},
}
out := ChatCompletionsResponseToResponses(resp, "glm-5.2", map[string]bool{"exec": true}, false, nil)
require.Len(t, out.Output, 2)
assert.Equal(t, "custom_tool_call", out.Output[0].Type)
assert.Equal(t, "call_1", out.Output[0].CallID)
assert.Equal(t, "exec", out.Output[0].Name)
assert.Equal(t, "dir", out.Output[0].Input)
assert.Empty(t, out.Output[0].Arguments)
assert.Equal(t, "function_call", out.Output[1].Type)
assert.Equal(t, "wait", out.Output[1].Name)
assert.Equal(t, `{"cell_id": 3}`, out.Output[1].Arguments)
}
func TestExtractCustomToolCallInput_FallsBackToRawArguments(t *testing.T) {
assert.Equal(t, "dir", extractCustomToolCallInput(`{"input": "dir"}`))
assert.Equal(t, "console.log(1)", extractCustomToolCallInput(`console.log(1)`))
assert.Equal(t, `{"other": "x"}`, extractCustomToolCallInput(`{"other": "x"}`))
assert.Equal(t, "", extractCustomToolCallInput(`{}`))
assert.Equal(t, "", extractCustomToolCallInput(""))
}
func TestChatCompletionsChunkToResponsesEvents_CustomToolCallStream(t *testing.T) {
state := NewChatCompletionsToResponsesStreamState("glm-5.2")
state.CustomTools = map[string]bool{"exec": true}
idx := 0
chunk := &ChatCompletionsChunk{
ID: "cc-1",
Choices: []ChatChunkChoice{{
Delta: ChatDelta{
ToolCalls: []ChatToolCall{{
Index: &idx,
ID: "call_1",
Function: ChatFunctionCall{Name: "exec", Arguments: `{"input": "dir"}`},
}},
},
}},
}
events := ChatCompletionsChunkToResponsesEvents(chunk, state)
events = append(events, FinalizeChatCompletionsResponsesStream(state)...)
var added, inputDone, itemDone *ResponsesStreamEvent
for i := range events {
evt := &events[i]
switch evt.Type {
case "response.output_item.added":
if evt.Item != nil && evt.Item.Type != "message" && evt.Item.Type != "reasoning" {
added = evt
}
case "response.custom_tool_call_input.done":
inputDone = evt
case "response.output_item.done":
if evt.Item != nil && evt.Item.Type == "custom_tool_call" {
itemDone = evt
}
case "response.function_call_arguments.delta", "response.function_call_arguments.done":
t.Fatalf("custom 工具调用不应产出 function_call 参数事件: %s", evt.Type)
}
}
require.NotNil(t, added, "缺少 custom_tool_call 的 output_item.added")
assert.Equal(t, "custom_tool_call", added.Item.Type)
assert.Equal(t, "exec", added.Item.Name)
require.NotNil(t, inputDone, "缺少 response.custom_tool_call_input.done")
assert.Equal(t, "dir", inputDone.Input)
assert.Equal(t, "call_1", inputDone.CallID)
require.NotNil(t, itemDone, "缺少 custom_tool_call 的 output_item.done")
assert.Equal(t, "call_1", itemDone.Item.CallID)
assert.Equal(t, "exec", itemDone.Item.Name)
assert.Equal(t, "dir", itemDone.Item.Input)
assert.Empty(t, itemDone.Item.Arguments)
// response.completed 的 output 数组同样携带 custom_tool_call 项。
final := events[len(events)-1]
require.Equal(t, "response.completed", final.Type)
require.NotNil(t, final.Response)
foundCustom := false
for _, item := range final.Response.Output {
if item.Type == "custom_tool_call" {
foundCustom = true
assert.Equal(t, "exec", item.Name)
assert.Equal(t, "dir", item.Input)
}
}
assert.True(t, foundCustom, "response.completed 缺少 custom_tool_call 输出项")
}
func TestResponsesToChatCompletionsRequest_ToolSearchToolBecomesProxyFunction(t *testing.T) {
req := &ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{{Type: "tool_search"}},
}
out, err := ResponsesToChatCompletionsRequest(req)
require.NoError(t, err)
require.Len(t, out.Tools, 1)
assert.Equal(t, "function", out.Tools[0].Type)
assert.Equal(t, "tool_search", out.Tools[0].Function.Name)
assert.Contains(t, string(out.Tools[0].Function.Parameters), `"query"`)
}
// codex 只在 ResponseItem 为 tool_search_call 变体且 execution=client 时执行
// tool search;同名 function_call 会命中 ToolSearchHandler 后因 payload 不匹配
// 触发 FunctionCallError::Fatal,直接中止整个 turn,因此回程必须还原项类型。
func TestChatCompletionsResponseToResponses_ToolSearchCallOutputItem(t *testing.T) {
resp := &ChatCompletionsResponse{
ID: "cc-1",
Choices: []ChatChoice{{
Message: ChatMessage{
Role: "assistant",
ToolCalls: []ChatToolCall{
{ID: "call_s", Function: ChatFunctionCall{Name: "tool_search", Arguments: `{"query":"gmail","limit":2}`}},
},
},
}},
}
out := ChatCompletionsResponseToResponses(resp, "glm-5.2", nil, true, nil)
require.Len(t, out.Output, 1)
item := out.Output[0]
assert.Equal(t, "tool_search_call", item.Type)
assert.Equal(t, "call_s", item.CallID)
// 线上形态:execution 必须为 "client"(codex 的必填字段,非 client 被忽略),
// arguments 必须是 JSON 对象而非字符串(codex 按对象解析 query/limit)。
b, err := json.Marshal(item)
require.NoError(t, err)
var m map[string]any
require.NoError(t, json.Unmarshal(b, &m))
assert.Equal(t, "client", m["execution"])
args, ok := m["arguments"].(map[string]any)
require.True(t, ok, "arguments 必须序列化为 JSON 对象")
assert.Equal(t, "gmail", args["query"])
}
func TestChatCompletionsResponseToResponses_ToolSearchNotDeclaredKeepsFunctionCall(t *testing.T) {
resp := &ChatCompletionsResponse{
Choices: []ChatChoice{{
Message: ChatMessage{
Role: "assistant",
ToolCalls: []ChatToolCall{
{ID: "call_s", Function: ChatFunctionCall{Name: "tool_search", Arguments: `{"query":"gmail"}`}},
},
},
}},
}
// 客户端未声明 type=tool_search 时,同名普通 function 工具不受影响。
out := ChatCompletionsResponseToResponses(resp, "glm-5.2", nil, false, nil)
require.Len(t, out.Output, 1)
assert.Equal(t, "function_call", out.Output[0].Type)
}
func TestChatCompletionsChunkToResponsesEvents_ToolSearchCallStream(t *testing.T) {
state := NewChatCompletionsToResponsesStreamState("glm-5.2")
state.ToolSearchDeclared = true
idx := 0
chunk := &ChatCompletionsChunk{
ID: "cc-1",
Choices: []ChatChunkChoice{{
Delta: ChatDelta{
ToolCalls: []ChatToolCall{{
Index: &idx,
ID: "call_s",
Function: ChatFunctionCall{Name: "tool_search", Arguments: `{"query":"gmail"}`},
}},
},
}},
}
events := ChatCompletionsChunkToResponsesEvents(chunk, state)
events = append(events, FinalizeChatCompletionsResponsesStream(state)...)
var added, itemDone *ResponsesStreamEvent
for i := range events {
evt := &events[i]
switch evt.Type {
case "response.output_item.added":
if evt.Item != nil && evt.Item.Type != "message" && evt.Item.Type != "reasoning" {
added = evt
}
case "response.output_item.done":
if evt.Item != nil && evt.Item.Type == "tool_search_call" {
itemDone = evt
}
case "response.function_call_arguments.delta", "response.function_call_arguments.done",
"response.custom_tool_call_input.delta", "response.custom_tool_call_input.done":
t.Fatalf("tool_search 调用不应产出 %s", evt.Type)
}
}
require.NotNil(t, added, "缺少 tool_search_call 的 output_item.added")
assert.Equal(t, "tool_search_call", added.Item.Type)
require.NotNil(t, itemDone, "缺少 tool_search_call 的 output_item.done")
assert.Equal(t, "call_s", itemDone.Item.CallID)
// SSE 线上形态经 responsesItemWire 白名单重组,必须单独断言。
sse, err := ResponsesEventToSSE(*itemDone)
require.NoError(t, err)
assert.Contains(t, sse, `"execution":"client"`)
assert.Contains(t, sse, `"arguments":{"query":"gmail"}`)
assert.Contains(t, sse, `"call_id":"call_s"`)
// response.completed 的 output 数组同样携带 tool_search_call 项。
final := events[len(events)-1]
require.Equal(t, "response.completed", final.Type)
require.NotNil(t, final.Response)
found := false
for _, item := range final.Response.Output {
if item.Type == "tool_search_call" {
found = true
assert.Equal(t, "call_s", item.CallID)
}
}
assert.True(t, found, "response.completed 缺少 tool_search_call 输出项")
}
func TestHasToolSearchTool(t *testing.T) {
assert.True(t, HasToolSearchTool([]ResponsesTool{{Type: "tool_search"}}))
assert.False(t, HasToolSearchTool([]ResponsesTool{{Type: "function", Name: "tool_search"}}))
assert.False(t, HasToolSearchTool(nil))
}
func TestResponsesToChatCompletionsRequest_NamespaceToolFlattensChildren(t *testing.T) {
req := &ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{{
Type: "namespace",
Name: "gmail",
Tools: []ResponsesTool{
{Type: "function", Name: "send", Description: "Send mail", Parameters: json.RawMessage(`{"type":"object","properties":{}}`)},
{Type: "custom", Name: "ignored_child"},
},
}},
}
out, err := ResponsesToChatCompletionsRequest(req)
require.NoError(t, err)
require.Len(t, out.Tools, 1, "namespace 子工具中仅 function 类型被摊平")
assert.Equal(t, "gmail__send", out.Tools[0].Function.Name)
assert.Equal(t, "Send mail", out.Tools[0].Function.Description)
}
func TestResponsesToolsParsing_StringToolBecomesCustom(t *testing.T) {
var req ResponsesRequest
require.NoError(t, json.Unmarshal([]byte(`{"model":"glm-5.2","input":"hi","tools":["exec",{"type":"function","name":"wait"}]}`), &req))
require.Len(t, req.Tools, 2)
assert.Equal(t, "custom", req.Tools[0].Type)
assert.Equal(t, "exec", req.Tools[0].Name)
assert.Equal(t, "function", req.Tools[1].Type)
assert.True(t, CustomToolNames(req.Tools)["exec"])
}
func TestFlattenNamespaceToolName_CapsAt64WithHashSuffix(t *testing.T) {
assert.Equal(t, "gmail__send", flattenNamespaceToolName("gmail", "send"))
long := flattenNamespaceToolName("very_long_namespace_prefix_for_testing_purposes", "and_a_rather_long_tool_name_too")
assert.LessOrEqual(t, len(long), 64)
assert.Contains(t, long, "__")
// 同输入结果稳定
assert.Equal(t, long, flattenNamespaceToolName("very_long_namespace_prefix_for_testing_purposes", "and_a_rather_long_tool_name_too"))
}
func TestResponsesInputToChatMessages_ToolSearchCallHistory(t *testing.T) {
input := json.RawMessage(`[
{"role":"user","content":"find tools"},
{"type":"tool_search_call","call_id":"call_s","arguments":{"query":"gmail"}},
{"type":"tool_search_output","call_id":"call_s","output":{"groups":["gmail"]}}
]`)
messages, err := responsesInputToChatMessages("", input)
require.NoError(t, err)
require.Len(t, messages, 3)
require.Len(t, messages[1].ToolCalls, 1)
assert.Equal(t, "tool_search", messages[1].ToolCalls[0].Function.Name)
assert.JSONEq(t, `{"query":"gmail"}`, messages[1].ToolCalls[0].Function.Arguments)
assert.Equal(t, "tool", messages[2].Role)
assert.Equal(t, "call_s", messages[2].ToolCallID)
assert.JSONEq(t, `"{\"groups\":[\"gmail\"]}"`, string(messages[2].Content))
}
func TestResponsesInputToChatMessages_NamespacedFunctionCallHistory(t *testing.T) {
input := json.RawMessage(`[
{"type":"function_call","call_id":"call_n","name":"send","namespace":"gmail","arguments":"{\"to\":\"a\"}"},
{"type":"function_call_output","call_id":"call_n","output":"ok"}
]`)
messages, err := responsesInputToChatMessages("", input)
require.NoError(t, err)
require.Len(t, messages, 2)
require.Len(t, messages[0].ToolCalls, 1)
assert.Equal(t, "gmail__send", messages[0].ToolCalls[0].Function.Name)
}
func TestChatCompletionsChunkToResponsesEvents_CustomToolNameArrivesLate(t *testing.T) {
state := NewChatCompletionsToResponsesStreamState("glm-5.2")
state.CustomTools = map[string]bool{"exec": true}
idx := 0
chunk1 := &ChatCompletionsChunk{Choices: []ChatChunkChoice{{Delta: ChatDelta{
ToolCalls: []ChatToolCall{{Index: &idx, ID: "call_1", Function: ChatFunctionCall{Arguments: `{"inp`}}},
}}}}
chunk2 := &ChatCompletionsChunk{Choices: []ChatChunkChoice{{Delta: ChatDelta{
ToolCalls: []ChatToolCall{{Index: &idx, Function: ChatFunctionCall{Name: "exec", Arguments: `ut": "dir"}`}}},
}}}}
var events []ResponsesStreamEvent
events = append(events, ChatCompletionsChunkToResponsesEvents(chunk1, state)...)
events = append(events, ChatCompletionsChunkToResponsesEvents(chunk2, state)...)
events = append(events, FinalizeChatCompletionsResponsesStream(state)...)
addedCount := 0
for _, evt := range events {
switch evt.Type {
case "response.output_item.added":
if evt.Item != nil && evt.Item.Type != "reasoning" && evt.Item.Type != "message" {
addedCount++
assert.Equal(t, "custom_tool_call", evt.Item.Type, "迟到的名字命中 custom 工具时按 custom_tool_call 宣告")
assert.Equal(t, "exec", evt.Item.Name)
}
case "response.function_call_arguments.delta", "response.function_call_arguments.done":
t.Fatalf("custom 调用不应产出 function 参数事件: %s", evt.Type)
case "response.custom_tool_call_input.done":
assert.Equal(t, "dir", evt.Input)
}
}
assert.Equal(t, 1, addedCount, "工具调用只宣告一次")
}
func TestChatCompletionsChunkToResponsesEvents_FunctionToolNameArrivesLate(t *testing.T) {
state := NewChatCompletionsToResponsesStreamState("glm-5.2")
state.CustomTools = map[string]bool{"exec": true}
idx := 0
chunk1 := &ChatCompletionsChunk{Choices: []ChatChunkChoice{{Delta: ChatDelta{
ToolCalls: []ChatToolCall{{Index: &idx, ID: "call_9", Function: ChatFunctionCall{Arguments: `{"cell`}}},
}}}}
chunk2 := &ChatCompletionsChunk{Choices: []ChatChunkChoice{{Delta: ChatDelta{
ToolCalls: []ChatToolCall{{Index: &idx, Function: ChatFunctionCall{Name: "wait", Arguments: `_id": 3}`}}},
}}}}
var events []ResponsesStreamEvent
events = append(events, ChatCompletionsChunkToResponsesEvents(chunk1, state)...)
events = append(events, ChatCompletionsChunkToResponsesEvents(chunk2, state)...)
events = append(events, FinalizeChatCompletionsResponsesStream(state)...)
deltas := ""
argsDone := ""
for _, evt := range events {
switch evt.Type {
case "response.function_call_arguments.delta":
deltas += evt.Delta
case "response.function_call_arguments.done":
argsDone = evt.Arguments
case "response.custom_tool_call_input.done":
t.Fatal("function 调用不应产出 custom 事件")
}
}
assert.Equal(t, `{"cell_id": 3}`, deltas, "宣告前累积的参数需在宣告时补发")
assert.Equal(t, `{"cell_id": 3}`, argsDone)
}
// 序列化层(MarshalJSON → responsesItemWire)单独走白名单重组,事件结构体上的字段
// 齐全不代表落到 SSE 线上的 JSON 齐全,必须在 wire 层再断言一次。
func TestResponsesEventToSSE_CustomToolCallItemCarriesAllFields(t *testing.T) {
evt := ResponsesStreamEvent{
Type: "response.output_item.done",
OutputIndex: 1,
Item: &ResponsesOutput{
Type: "custom_tool_call",
ID: "item_1",
CallID: "call_1",
Name: "exec",
Input: "dir",
Status: "completed",
},
}
sse, err := ResponsesEventToSSE(evt)
require.NoError(t, err)
assert.Contains(t, sse, `"call_id":"call_1"`)
assert.Contains(t, sse, `"name":"exec"`)
assert.Contains(t, sse, `"input":"dir"`)
assert.Contains(t, sse, `"type":"custom_tool_call"`)
}
func TestNamespaceToolNames_MapsFlattenedNames(t *testing.T) {
tools := []ResponsesTool{
{Type: "namespace", Name: "gmail", Tools: []ResponsesTool{
{Type: "function", Name: "send"},
{Type: "custom", Name: "skip_me"},
}},
{Type: "namespace", Name: "crm", Children: []ResponsesTool{
{Type: "function", Name: "query"},
}},
{Type: "function", Name: "wait"},
}
m := NamespaceToolNames(tools)
require.Len(t, m, 2)
assert.Equal(t, NamespacedToolName{Namespace: "gmail", Name: "send"}, m["gmail__send"])
assert.Equal(t, NamespacedToolName{Namespace: "crm", Name: "query"}, m["crm__query"])
// 摊平名超长时截断加哈希,无法按字符串切分还原,必须经映射反查。
longNS := "very_long_namespace_prefix_for_testing_purposes"
longChild := "and_a_rather_long_tool_name_too"
m2 := NamespaceToolNames([]ResponsesTool{{
Type: "namespace", Name: longNS,
Tools: []ResponsesTool{{Type: "function", Name: longChild}},
}})
assert.Equal(t, NamespacedToolName{Namespace: longNS, Name: longChild},
m2[flattenNamespaceToolName(longNS, longChild)])
assert.Nil(t, NamespaceToolNames(nil))
}
// 内置 tool_search 降级后的代理 function 与客户端声明的同名工具无法区分:回程会把
// 普通工具的调用劫持成 tool_search_call,必须显式拒绝(代理不能改名,codex 的模型
// 侧按 tool_search 这个名字调用)。
func TestResponsesToChatCompletionsRequest_RejectsToolSearchNameConflict(t *testing.T) {
// 与顶层 function 工具同名。
_, err := ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{
{Type: "tool_search"},
{Type: "function", Name: "tool_search"},
},
})
require.Error(t, err, "与内置 tool_search 代理撞名的 function 工具必须拒绝")
assert.Contains(t, err.Error(), "tool_search")
// 与顶层 custom 工具同名。
_, err = ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{
{Type: "custom", Name: "tool_search"},
{Type: "tool_search"},
},
})
require.Error(t, err, "与内置 tool_search 代理撞名的 custom 工具必须拒绝")
// 重复声明 type=tool_search 去重后只产出一个代理,不拒绝。
out, err := ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{{Type: "tool_search"}, {Type: "tool_search"}},
})
require.NoError(t, err)
require.Len(t, out.Tools, 1)
assert.Equal(t, "tool_search", out.Tools[0].Function.Name)
}
// tool_choice 指向被转换丢弃的工具(如 web_search)或不存在的名字时不能原样转发,
// chat 上游会因选择项指向未声明工具而 400;字符串形式与指向幸存工具的选择保持转发。
func TestResponsesToChatCompletionsRequest_DropsToolChoiceForDroppedTool(t *testing.T) {
// 强制选择被丢弃的 web_search:工具没了,选择项也必须丢。
out, err := ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{
{Type: "function", Name: "wait", Parameters: json.RawMessage(`{"type":"object","properties":{}}`)},
{Type: "web_search"},
},
ToolChoice: json.RawMessage(`{"type":"web_search"}`),
})
require.NoError(t, err)
require.Len(t, out.Tools, 1)
assert.Empty(t, out.ToolChoice, "指向被丢弃服务端工具的 tool_choice 必须丢弃")
// 具名选择指向不存在的工具名。
out, err = ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{{Type: "function", Name: "wait"}},
ToolChoice: json.RawMessage(`{"type":"function","name":"missing"}`),
})
require.NoError(t, err)
assert.Empty(t, out.ToolChoice, "指向不存在工具名的 tool_choice 必须丢弃")
// 字符串形式与指向幸存工具的选择保持原有转发行为。
out, err = ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{{Type: "function", Name: "wait"}},
ToolChoice: json.RawMessage(`"auto"`),
})
require.NoError(t, err)
assert.JSONEq(t, `"auto"`, string(out.ToolChoice))
out, err = ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{{Type: "function", Name: "wait"}},
ToolChoice: json.RawMessage(`{"type":"function","name":"wait"}`),
})
require.NoError(t, err)
assert.JSONEq(t, `{"type":"function","function":{"name":"wait"}}`, string(out.ToolChoice))
}
// tool_search 工具没有被丢弃而是降级为同名 function 代理,强制选择它的 tool_choice
// 必须同步降级为指向代理的 function 选择,不能静默丢弃(丢弃会把强制搜索退化为
// 自动选择,模型可以不执行搜索)。
func TestResponsesToChatCompletionsRequest_ToolSearchToolChoiceMapsToProxy(t *testing.T) {
out, err := ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{{Type: "tool_search"}},
ToolChoice: json.RawMessage(`{"type":"tool_search"}`),
})
require.NoError(t, err)
assert.JSONEq(t, `{"type":"function","function":{"name":"tool_search"}}`, string(out.ToolChoice))
// 未声明 type=tool_search 时强制选择它没有可指向的代理,丢弃选择项。
out, err = ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{{Type: "function", Name: "wait"}},
ToolChoice: json.RawMessage(`{"type":"tool_search"}`),
})
require.NoError(t, err)
assert.Empty(t, out.ToolChoice)
}
// 客户端请求在原生 Responses API 上合法(namespace 子工具按 namespace+name 路由),
// 是摊平转换让名字产生歧义;歧义无法消除时必须显式拒绝整个请求(400),而不是
// 静默降级——否则重复声明发给上游、回程还原到错误工具,问题只能靠抓包定位。
func TestResponsesToChatCompletionsRequest_RejectsAmbiguousFlattenedNames(t *testing.T) {
// 摊平名与顶层 function 工具撞名。
_, err := ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{
{Type: "function", Name: "gmail__send"},
{Type: "namespace", Name: "gmail", Tools: []ResponsesTool{{Type: "function", Name: "send"}}},
},
})
require.Error(t, err, "与顶层工具撞名的摊平必须拒绝")
assert.Contains(t, err.Error(), "gmail__send")
// 不同 namespace 组合产生相同摊平名。
_, err = ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{
{Type: "namespace", Name: "a", Tools: []ResponsesTool{{Type: "function", Name: "b__c"}}},
{Type: "namespace", Name: "a__b", Tools: []ResponsesTool{{Type: "function", Name: "c"}}},
},
})
require.Error(t, err, "跨 namespace 撞名的摊平必须拒绝")
assert.Contains(t, err.Error(), "a__b__c")
}
// 完全相同的 (namespace, 子工具) 重复声明不构成歧义:去重后正常转换,不拒绝。
func TestResponsesToChatCompletionsRequest_DedupesIdenticalNamespaceChildren(t *testing.T) {
out, err := ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{
{Type: "namespace", Name: "gmail", Tools: []ResponsesTool{
{Type: "function", Name: "send"},
{Type: "function", Name: "send"},
}},
},
})
require.NoError(t, err)
require.Len(t, out.Tools, 1, "重复声明的同一子工具只声明一次")
assert.Equal(t, "gmail__send", out.Tools[0].Function.Name)
}
// codex 按 namespace+name 路由 namespace 子工具的调用:回程必须把摊平名还原为
// 裸子工具名并带独立 namespace 字段,平铺名的 function_call 会被 codex 判为
// unsupported call 拒绝执行。
func TestChatCompletionsResponseToResponses_NamespacedToolCallRestored(t *testing.T) {
resp := &ChatCompletionsResponse{
ID: "cc-1",
Choices: []ChatChoice{{
Message: ChatMessage{
Role: "assistant",
ToolCalls: []ChatToolCall{
{ID: "call_n", Function: ChatFunctionCall{Name: "mcp__svc__echo", Arguments: `{"text":"hi"}`}},
{ID: "call_9", Function: ChatFunctionCall{Name: "wait", Arguments: `{"cell_id": 3}`}},
},
},
}},
}
nsTools := map[string]NamespacedToolName{
"mcp__svc__echo": {Namespace: "mcp__svc", Name: "echo"},
}
out := ChatCompletionsResponseToResponses(resp, "glm-5.2", nil, false, nsTools)
require.Len(t, out.Output, 2)
item := out.Output[0]
assert.Equal(t, "function_call", item.Type)
assert.Equal(t, "echo", item.Name)
assert.Equal(t, "mcp__svc", item.Namespace)
assert.Equal(t, "call_n", item.CallID)
assert.Equal(t, `{"text":"hi"}`, item.Arguments)
// 非流式响应体走 ResponsesOutput.MarshalJSON,namespace 必须落到线上 JSON。
b, err := json.Marshal(item)
require.NoError(t, err)
assert.Contains(t, string(b), `"namespace":"mcp__svc"`)
assert.Contains(t, string(b), `"name":"echo"`)
// 未命中映射的普通 function 调用不受影响,且不携带 namespace 字段。
assert.Equal(t, "wait", out.Output[1].Name)
assert.Empty(t, out.Output[1].Namespace)
b2, err := json.Marshal(out.Output[1])
require.NoError(t, err)
assert.NotContains(t, string(b2), `"namespace"`)
}
func TestChatCompletionsChunkToResponsesEvents_NamespacedToolCallStream(t *testing.T) {
state := NewChatCompletionsToResponsesStreamState("glm-5.2")
state.NamespaceTools = map[string]NamespacedToolName{
"mcp__svc__echo": {Namespace: "mcp__svc", Name: "echo"},
}
idx := 0
chunk := &ChatCompletionsChunk{
ID: "cc-1",
Choices: []ChatChunkChoice{{
Delta: ChatDelta{
ToolCalls: []ChatToolCall{{
Index: &idx,
ID: "call_n",
Function: ChatFunctionCall{Name: "mcp__svc__echo", Arguments: `{"text":"hi"}`},
}},
},
}},
}
events := ChatCompletionsChunkToResponsesEvents(chunk, state)
events = append(events, FinalizeChatCompletionsResponsesStream(state)...)
var added, itemDone *ResponsesStreamEvent
for i := range events {
evt := &events[i]
switch evt.Type {
case "response.output_item.added":
if evt.Item != nil && evt.Item.Type != "message" && evt.Item.Type != "reasoning" {
added = evt
}
case "response.output_item.done":
if evt.Item != nil && evt.Item.Type == "function_call" {
itemDone = evt
}
case "response.custom_tool_call_input.delta", "response.custom_tool_call_input.done":
t.Fatalf("namespace 子工具调用不应产出 custom 事件: %s", evt.Type)
}
}
require.NotNil(t, added, "缺少 namespace 调用的 output_item.added")
assert.Equal(t, "function_call", added.Item.Type)
assert.Equal(t, "echo", added.Item.Name)
assert.Equal(t, "mcp__svc", added.Item.Namespace)
require.NotNil(t, itemDone, "缺少 namespace 调用的 output_item.done")
assert.Equal(t, "call_n", itemDone.Item.CallID)
assert.Equal(t, "echo", itemDone.Item.Name)
assert.Equal(t, "mcp__svc", itemDone.Item.Namespace)
assert.Equal(t, `{"text":"hi"}`, itemDone.Item.Arguments)
// SSE 线上形态经 responsesItemWire 白名单重组,必须单独断言 namespace 落线。
sse, err := ResponsesEventToSSE(*itemDone)
require.NoError(t, err)
assert.Contains(t, sse, `"namespace":"mcp__svc"`)
assert.Contains(t, sse, `"name":"echo"`)
assert.Contains(t, sse, `"call_id":"call_n"`)
// response.completed 的 output 数组同样携带还原后的 namespace 调用项。
final := events[len(events)-1]
require.Equal(t, "response.completed", final.Type)
require.NotNil(t, final.Response)
found := false
for _, item := range final.Response.Output {
if item.Type == "function_call" {
found = true
assert.Equal(t, "echo", item.Name)
assert.Equal(t, "mcp__svc", item.Namespace)
}
}
assert.True(t, found, "response.completed 缺少还原后的 namespace 调用项")
}
func TestChatCompletionsChunkToResponsesEvents_NamespacedToolNameArrivesLate(t *testing.T) {
state := NewChatCompletionsToResponsesStreamState("glm-5.2")
state.NamespaceTools = map[string]NamespacedToolName{
"mcp__svc__echo": {Namespace: "mcp__svc", Name: "echo"},
}
idx := 0
chunk1 := &ChatCompletionsChunk{Choices: []ChatChunkChoice{{Delta: ChatDelta{
ToolCalls: []ChatToolCall{{Index: &idx, ID: "call_n", Function: ChatFunctionCall{Arguments: `{"te`}}},
}}}}
chunk2 := &ChatCompletionsChunk{Choices: []ChatChunkChoice{{Delta: ChatDelta{
ToolCalls: []ChatToolCall{{Index: &idx, Function: ChatFunctionCall{Name: "mcp__svc__echo", Arguments: `xt":"hi"}`}}},
}}}}
var events []ResponsesStreamEvent
events = append(events, ChatCompletionsChunkToResponsesEvents(chunk1, state)...)
events = append(events, ChatCompletionsChunkToResponsesEvents(chunk2, state)...)
events = append(events, FinalizeChatCompletionsResponsesStream(state)...)
addedCount := 0
deltas := ""
for _, evt := range events {
switch evt.Type {
case "response.output_item.added":
if evt.Item != nil && evt.Item.Type != "reasoning" && evt.Item.Type != "message" {
addedCount++
assert.Equal(t, "echo", evt.Item.Name, "迟到的名字命中 namespace 映射时按还原名宣告")
assert.Equal(t, "mcp__svc", evt.Item.Namespace)
}
case "response.function_call_arguments.delta":
deltas += evt.Delta
}
}
assert.Equal(t, 1, addedCount, "工具调用只宣告一次")
assert.Equal(t, `{"text":"hi"}`, deltas, "宣告前累积的参数需在宣告时补发")
}
func TestChatCompletionsChunkToResponsesEvents_FunctionToolStreamUnaffected(t *testing.T) {
state := NewChatCompletionsToResponsesStreamState("glm-5.2")
state.CustomTools = map[string]bool{"exec": true}
idx := 0
chunk := &ChatCompletionsChunk{
Choices: []ChatChunkChoice{{
Delta: ChatDelta{
ToolCalls: []ChatToolCall{{
Index: &idx,
ID: "call_9",
Function: ChatFunctionCall{Name: "wait", Arguments: `{"cell_id": 3}`},
}},
},
}},
}
events := ChatCompletionsChunkToResponsesEvents(chunk, state)
events = append(events, FinalizeChatCompletionsResponsesStream(state)...)
sawArgsDelta := false
for _, evt := range events {
if evt.Type == "response.function_call_arguments.delta" {
sawArgsDelta = true
}
if evt.Type == "response.custom_tool_call_input.done" {
t.Fatal("function 工具不应产出 custom_tool_call 事件")
}
}
assert.True(t, sawArgsDelta, "function 工具应保持原有参数增量事件")
}
@@ -459,7 +459,7 @@ func TestChatCompletionsResponseToResponses_DeepSeekReasoningOnlyFallsBackToMess
}},
}
out := ChatCompletionsResponseToResponses(resp, "deepseek-reasoner")
out := ChatCompletionsResponseToResponses(resp, "deepseek-reasoner", nil, false, nil)
require.Len(t, out.Output, 2)
require.Equal(t, "reasoning", out.Output[0].Type)
@@ -493,7 +493,7 @@ func TestChatCompletionsResponseToResponses_DeepSeekReasoningToolCallDoesNotFall
}},
}
out := ChatCompletionsResponseToResponses(resp, "deepseek-reasoner")
out := ChatCompletionsResponseToResponses(resp, "deepseek-reasoner", nil, false, nil)
require.Len(t, out.Output, 2)
require.Equal(t, "reasoning", out.Output[0].Type)
@@ -0,0 +1,99 @@
package apicompat
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAnthropicUsageFromResponsesUsage_CacheCreation(t *testing.T) {
usage := &ResponsesUsage{
InputTokens: 20,
OutputTokens: 5,
CacheCreationInputTokens: 6,
InputTokensDetails: &ResponsesInputTokensDetails{
CachedTokens: 4,
},
}
got := anthropicUsageFromResponsesUsage(usage)
assert.Equal(t, 10, got.InputTokens, "input = total(20) - cache_read(4) - cache_creation(6)")
assert.Equal(t, 5, got.OutputTokens)
assert.Equal(t, 4, got.CacheReadInputTokens)
assert.Equal(t, 6, got.CacheCreationInputTokens, "cache creation must be preserved")
}
func TestAnthropicUsageFromResponsesUsage_NoCacheCreation(t *testing.T) {
usage := &ResponsesUsage{
InputTokens: 10,
OutputTokens: 5,
InputTokensDetails: &ResponsesInputTokensDetails{
CachedTokens: 3,
},
}
got := anthropicUsageFromResponsesUsage(usage)
assert.Equal(t, 7, got.InputTokens)
assert.Equal(t, 3, got.CacheReadInputTokens)
assert.Equal(t, 0, got.CacheCreationInputTokens)
}
func TestResponsesEventToAnthropicEvents_StreamingCacheCreation(t *testing.T) {
state := NewResponsesEventToAnthropicState()
state.MessageStartSent = true
completedEvt := &ResponsesStreamEvent{
Type: "response.completed",
Response: &ResponsesResponse{
Status: "completed",
Usage: &ResponsesUsage{
InputTokens: 20,
OutputTokens: 5,
CacheCreationInputTokens: 6,
InputTokensDetails: &ResponsesInputTokensDetails{
CachedTokens: 4,
},
},
},
}
events := ResponsesEventToAnthropicEvents(completedEvt, state)
var deltaEvt *AnthropicStreamEvent
for i := range events {
if events[i].Type == "message_delta" {
deltaEvt = &events[i]
break
}
}
require.NotNil(t, deltaEvt, "should have message_delta event")
require.NotNil(t, deltaEvt.Usage)
assert.Equal(t, 6, deltaEvt.Usage.CacheCreationInputTokens, "streaming cache_creation must be preserved")
assert.Equal(t, 10, deltaEvt.Usage.InputTokens, "input = 20 - 4(read) - 6(creation)")
assert.Equal(t, 4, deltaEvt.Usage.CacheReadInputTokens)
}
func TestAnthropicToResponsesResponse_CacheCreation(t *testing.T) {
resp := AnthropicResponse{
ID: "msg_test",
Type: "message",
Role: "assistant",
Model: "claude-opus-4-6",
Usage: AnthropicUsage{
InputTokens: 10,
OutputTokens: 5,
CacheReadInputTokens: 4,
CacheCreationInputTokens: 6,
},
StopReason: "end_turn",
}
out := AnthropicToResponsesResponse(&resp)
require.NotNil(t, out.Usage)
assert.Equal(t, 20, out.Usage.InputTokens, "total = input(10) + cache_read(4) + cache_creation(6)")
assert.Equal(t, 6, out.Usage.CacheCreationInputTokens, "cache creation must round-trip")
}
@@ -0,0 +1,212 @@
package apicompat
import (
"bytes"
"encoding/json"
"fmt"
"strings"
)
// ResponsesNamespaceName identifies a function child in a Responses namespace.
// It aliases the chat bridge mapping so both native and bridged paths share one
// namespace identity contract.
type ResponsesNamespaceName = NamespacedToolName
// FlattenResponsesNamespaces converts Codex private namespace declarations into
// public Responses function tools and rewrites namespace-qualified request calls.
func FlattenResponsesNamespaces(req map[string]any) (map[string]ResponsesNamespaceName, bool, error) {
return FlattenResponsesNamespacesExcept(req, nil)
}
// FlattenResponsesNamespacesExcept is FlattenResponsesNamespaces with a set of
// service-owned namespace names that must remain native in the request.
func FlattenResponsesNamespacesExcept(req map[string]any, preserved map[string]bool) (map[string]ResponsesNamespaceName, bool, error) {
if req == nil {
return nil, false, nil
}
tools, ok := req["tools"].([]any)
if !ok || len(tools) == 0 {
return nil, false, nil
}
topLevel := make(map[string]bool)
for _, raw := range tools {
tool, ok := raw.(map[string]any)
if !ok {
continue
}
typ := strings.TrimSpace(stringValue(tool["type"]))
name := strings.TrimSpace(stringValue(tool["name"]))
if (typ == "function" || typ == "custom") && name != "" {
topLevel[name] = true
}
}
names := make(map[string]ResponsesNamespaceName)
for _, raw := range tools {
tool, ok := raw.(map[string]any)
if !ok || strings.TrimSpace(stringValue(tool["type"])) != "namespace" {
continue
}
namespace := strings.TrimSpace(stringValue(tool["name"]))
if namespace == "" || preserved[namespace] {
continue
}
for _, rawChild := range namespaceChildren(tool) {
child, ok := rawChild.(map[string]any)
if !ok || strings.TrimSpace(stringValue(child["type"])) != "function" {
continue
}
name := strings.TrimSpace(stringValue(child["name"]))
if name == "" {
continue
}
flat := flattenNamespaceToolName(namespace, name)
entry := ResponsesNamespaceName{Namespace: namespace, Name: name}
if topLevel[flat] {
return nil, false, fmt.Errorf("namespace tool %q/%q flattens to %q which conflicts with a top-level tool of the same name; this upstream cannot disambiguate them, rename one of the tools", namespace, name, flat)
}
if prev, exists := names[flat]; exists && prev != entry {
return nil, false, fmt.Errorf("namespace tools %q/%q and %q/%q both flatten to %q; this upstream cannot disambiguate them, rename one of the tools", prev.Namespace, prev.Name, namespace, name, flat)
}
names[flat] = entry
}
}
if len(names) == 0 {
return nil, false, nil
}
flattened := make([]any, 0, len(tools)+len(names))
seen := make(map[string]bool)
for _, raw := range tools {
tool, ok := raw.(map[string]any)
if !ok || strings.TrimSpace(stringValue(tool["type"])) != "namespace" {
flattened = append(flattened, raw)
continue
}
namespace := strings.TrimSpace(stringValue(tool["name"]))
if preserved[namespace] {
flattened = append(flattened, raw)
continue
}
for _, rawChild := range namespaceChildren(tool) {
child, ok := rawChild.(map[string]any)
if !ok || strings.TrimSpace(stringValue(child["type"])) != "function" {
continue
}
name := strings.TrimSpace(stringValue(child["name"]))
flat := flattenNamespaceToolName(namespace, name)
if name == "" || seen[flat] {
continue
}
seen[flat] = true
flatChild := make(map[string]any, len(child))
for key, value := range child {
flatChild[key] = value
}
flatChild["name"] = flat
flattened = append(flattened, flatChild)
}
}
req["tools"] = flattened
rewriteNamespaceQualifiedCalls(req["input"], names)
if choice, ok := req["tool_choice"].(map[string]any); ok {
choiceNamespace := strings.TrimSpace(stringValue(choice["name"]))
if strings.TrimSpace(stringValue(choice["type"])) == "namespace" && !preserved[choiceNamespace] {
req["tool_choice"] = "auto"
} else {
rewriteNamespaceQualifiedCall(choice, names)
}
}
return names, true, nil
}
// RestoreResponsesNamespaceCalls restores flattened function calls in a JSON
// Responses payload to the namespace/name identity expected by Codex.
func RestoreResponsesNamespaceCalls(payload []byte, names map[string]ResponsesNamespaceName) ([]byte, bool, error) {
if len(payload) == 0 || len(names) == 0 {
return payload, false, nil
}
var value any
if err := json.Unmarshal(payload, &value); err != nil {
return payload, false, err
}
changed := restoreResponsesNamespaceValue(value, names)
if !changed {
return payload, false, nil
}
var rebuilt bytes.Buffer
encoder := json.NewEncoder(&rebuilt)
encoder.SetEscapeHTML(false)
if err := encoder.Encode(value); err != nil {
return payload, false, err
}
return bytes.TrimSuffix(rebuilt.Bytes(), []byte("\n")), true, nil
}
func namespaceChildren(tool map[string]any) []any {
if children, ok := tool["tools"].([]any); ok && len(children) > 0 {
return children
}
children, _ := tool["children"].([]any)
return children
}
func rewriteNamespaceQualifiedCalls(value any, names map[string]ResponsesNamespaceName) {
switch typed := value.(type) {
case []any:
for _, item := range typed {
rewriteNamespaceQualifiedCalls(item, names)
}
case map[string]any:
if strings.TrimSpace(stringValue(typed["type"])) == "function_call" {
rewriteNamespaceQualifiedCall(typed, names)
}
for _, child := range typed {
rewriteNamespaceQualifiedCalls(child, names)
}
}
}
func rewriteNamespaceQualifiedCall(item map[string]any, names map[string]ResponsesNamespaceName) bool {
namespace := strings.TrimSpace(stringValue(item["namespace"]))
name := strings.TrimSpace(stringValue(item["name"]))
if namespace == "" || name == "" {
return false
}
flat := flattenNamespaceToolName(namespace, name)
entry, ok := names[flat]
if !ok || entry.Namespace != namespace || entry.Name != name {
return false
}
item["name"] = flat
delete(item, "namespace")
return true
}
func restoreResponsesNamespaceValue(value any, names map[string]ResponsesNamespaceName) bool {
changed := false
switch typed := value.(type) {
case []any:
for _, item := range typed {
changed = restoreResponsesNamespaceValue(item, names) || changed
}
case map[string]any:
if strings.TrimSpace(stringValue(typed["type"])) == "function_call" {
if entry, ok := names[strings.TrimSpace(stringValue(typed["name"]))]; ok {
typed["name"] = entry.Name
typed["namespace"] = entry.Namespace
changed = true
}
}
for _, child := range typed {
changed = restoreResponsesNamespaceValue(child, names) || changed
}
}
return changed
}
func stringValue(value any) string {
text, _ := value.(string)
return text
}
@@ -0,0 +1,164 @@
package apicompat
import (
"testing"
"github.com/stretchr/testify/require"
)
func TestFlattenResponsesNamespaces_RewritesDeclarationHistoryAndChoice(t *testing.T) {
req := map[string]any{
"model": "gpt-5.5",
"tools": []any{
map[string]any{"type": "function", "name": "plain", "description": "keep"},
map[string]any{
"type": "namespace",
"name": "collaboration",
"tools": []any{
map[string]any{"type": "function", "name": "spawn_agent", "description": "spawn", "parameters": map[string]any{"type": "object"}},
},
},
},
"tool_choice": map[string]any{"type": "function", "name": "spawn_agent", "namespace": "collaboration"},
"input": []any{
map[string]any{"type": "function_call", "call_id": "call_1", "name": "spawn_agent", "namespace": "collaboration", "arguments": "{}"},
map[string]any{"type": "message", "role": "user", "content": "hi", "name": "spawn_agent", "namespace": "collaboration"},
},
}
names, changed, err := FlattenResponsesNamespaces(req)
require.NoError(t, err)
require.True(t, changed)
require.Equal(t, ResponsesNamespaceName{Namespace: "collaboration", Name: "spawn_agent"}, names["collaboration__spawn_agent"])
tools, ok := req["tools"].([]any)
require.True(t, ok)
require.Len(t, tools, 2)
plainTool, ok := tools[0].(map[string]any)
require.True(t, ok)
require.Equal(t, "plain", plainTool["name"])
flatTool, ok := tools[1].(map[string]any)
require.True(t, ok)
require.Equal(t, "collaboration__spawn_agent", flatTool["name"])
require.Equal(t, "spawn", flatTool["description"])
choice, ok := req["tool_choice"].(map[string]any)
require.True(t, ok)
require.Equal(t, "collaboration__spawn_agent", choice["name"])
require.NotContains(t, choice, "namespace")
input, ok := req["input"].([]any)
require.True(t, ok)
require.Len(t, input, 2)
call, ok := input[0].(map[string]any)
require.True(t, ok)
require.Equal(t, "collaboration__spawn_agent", call["name"])
require.NotContains(t, call, "namespace")
message, ok := input[1].(map[string]any)
require.True(t, ok)
require.Equal(t, "spawn_agent", message["name"])
require.Equal(t, "collaboration", message["namespace"])
require.Equal(t, "gpt-5.5", req["model"])
}
func TestFlattenResponsesNamespaces_RejectsFlatNameCollision(t *testing.T) {
req := map[string]any{"tools": []any{
map[string]any{"type": "function", "name": "collaboration__spawn_agent"},
map[string]any{"type": "namespace", "name": "collaboration", "tools": []any{
map[string]any{"type": "function", "name": "spawn_agent"},
}},
}}
_, _, err := FlattenResponsesNamespaces(req)
require.ErrorContains(t, err, "conflicts with a top-level tool")
}
func TestFlattenResponsesNamespaces_NamespaceGroupChoiceFallsBackToAuto(t *testing.T) {
req := map[string]any{
"tools": []any{map[string]any{
"type": "namespace", "name": "collaboration", "tools": []any{
map[string]any{"type": "function", "name": "spawn_agent"},
map[string]any{"type": "function", "name": "send_message"},
},
}},
"tool_choice": map[string]any{"type": "namespace", "name": "collaboration"},
}
_, changed, err := FlattenResponsesNamespaces(req)
require.NoError(t, err)
require.True(t, changed)
require.Equal(t, "auto", req["tool_choice"])
}
func TestFlattenResponsesNamespacesExcept_PreservesBuiltInNamespaceAndChoice(t *testing.T) {
req := map[string]any{
"tools": []any{
map[string]any{"type": "namespace", "name": "image_gen", "tools": []any{
map[string]any{"type": "function", "name": "imagegen"},
}},
map[string]any{"type": "namespace", "name": "collaboration", "tools": []any{
map[string]any{"type": "function", "name": "spawn_agent"},
}},
},
"tool_choice": map[string]any{"type": "namespace", "name": "image_gen"},
}
names, changed, err := FlattenResponsesNamespacesExcept(req, map[string]bool{"image_gen": true})
require.NoError(t, err)
require.True(t, changed)
require.Contains(t, names, "collaboration__spawn_agent")
tools, ok := req["tools"].([]any)
require.True(t, ok)
require.Len(t, tools, 2)
preservedTool, ok := tools[0].(map[string]any)
require.True(t, ok)
require.Equal(t, "namespace", preservedTool["type"])
require.Equal(t, "image_gen", preservedTool["name"])
flatTool, ok := tools[1].(map[string]any)
require.True(t, ok)
require.Equal(t, "function", flatTool["type"])
require.Equal(t, "collaboration__spawn_agent", flatTool["name"])
require.Equal(t, map[string]any{"type": "namespace", "name": "image_gen"}, req["tool_choice"])
}
func TestFlattenResponsesNamespaces_RejectsNamespaceCollision(t *testing.T) {
req := map[string]any{"tools": []any{
map[string]any{"type": "namespace", "name": "a", "tools": []any{
map[string]any{"type": "function", "name": "b__c"},
}},
map[string]any{"type": "namespace", "name": "a__b", "tools": []any{
map[string]any{"type": "function", "name": "c"},
}},
}}
_, _, err := FlattenResponsesNamespaces(req)
require.ErrorContains(t, err, "both flatten")
}
func TestRestoreResponsesNamespaceCalls_RewritesOnlyFunctionCalls(t *testing.T) {
payload := []byte(`{"type":"response.completed","response":{"output":[{"type":"function_call","name":"collaboration__spawn_agent","call_id":"call_1","arguments":"{}","extra":"keep"},{"type":"function_call","name":"plain","arguments":"{}"},{"type":"message","name":"collaboration__spawn_agent","content":"<tag>&value</tag>"}]}}`)
names := map[string]ResponsesNamespaceName{
"collaboration__spawn_agent": {Namespace: "collaboration", Name: "spawn_agent"},
}
got, changed, err := RestoreResponsesNamespaceCalls(payload, names)
require.NoError(t, err)
require.True(t, changed)
require.JSONEq(t, `{"type":"response.completed","response":{"output":[{"type":"function_call","name":"spawn_agent","namespace":"collaboration","call_id":"call_1","arguments":"{}","extra":"keep"},{"type":"function_call","name":"plain","arguments":"{}"},{"type":"message","name":"collaboration__spawn_agent","content":"<tag>&value</tag>"}]}}`, string(got))
require.Contains(t, string(got), "<tag>&value</tag>")
require.NotContains(t, string(got), `\u003c`)
}
func TestRestoreResponsesNamespaceCalls_RewritesLifecycleItems(t *testing.T) {
for _, eventType := range []string{"response.output_item.added", "response.output_item.done"} {
t.Run(eventType, func(t *testing.T) {
payload := []byte(`{"type":"` + eventType + `","item":{"type":"function_call","name":"collaboration__spawn_agent","arguments":"{}"}}`)
got, changed, err := RestoreResponsesNamespaceCalls(payload, map[string]ResponsesNamespaceName{
"collaboration__spawn_agent": {Namespace: "collaboration", Name: "spawn_agent"},
})
require.NoError(t, err)
require.True(t, changed)
require.JSONEq(t, `{"type":"`+eventType+`","item":{"type":"function_call","name":"spawn_agent","namespace":"collaboration","arguments":"{}"}}`, string(got))
})
}
}
@@ -86,6 +86,23 @@ func (e ResponsesStreamEvent) MarshalJSON() ([]byte, error) {
}
return json.Marshal(m)
case "response.custom_tool_call_input.delta", "response.custom_tool_call_input.done":
m := e.wireBase()
e.putItemID(m)
m["output_index"] = e.OutputIndex
if e.CallID != "" {
m["call_id"] = e.CallID
}
if e.Name != "" {
m["name"] = e.Name
}
if e.Type == "response.custom_tool_call_input.done" {
m["input"] = e.Input
} else {
m["delta"] = e.Delta
}
return json.Marshal(m)
default:
// response.created / completed / done / failed / incomplete and any
// event type not shaped above keep the default struct marshalling.
@@ -167,6 +184,23 @@ func responsesItemWire(item *ResponsesOutput) map[string]any {
m["call_id"] = item.CallID
m["name"] = item.Name
m["arguments"] = item.Arguments
// namespace 子工具的还原调用:codex 按 namespace+name 路由,缺少该字段
// 会被判为 unsupported call。
if item.Namespace != "" {
m["namespace"] = item.Namespace
}
case "custom_tool_call":
// custom/freeform 工具调用(如 codex 的 exec):input 为自由文本。缺少
// call_id/name 时 codex 无法路由该调用(表现为 unsupported call)。
m["call_id"] = item.CallID
m["name"] = item.Name
m["input"] = item.Input
case "tool_search_call":
// tool_search 调用还原项:execution 必须为 "client"(否则 codex 忽略该
// 调用),arguments 在线上是 JSON 对象而非字符串。
m["call_id"] = item.CallID
m["execution"] = "client"
m["arguments"] = toolSearchCallArgumentsJSON(item.Arguments)
}
return m
}
@@ -102,6 +102,26 @@ func TestWire_ArgumentsDonePresentEvenEmpty(t *testing.T) {
require.Equal(t, "", m["arguments"])
}
// TestWire_CustomToolCallInputIndexPresentAtZero guards the omitempty trap for
// custom_tool_call_input.delta/done: output_index must serialize even when 0
// (custom tool call as the first output item).
func TestWire_CustomToolCallInputIndexPresentAtZero(t *testing.T) {
d := marshalEvent(t, ResponsesStreamEvent{
Type: "response.custom_tool_call_input.delta", OutputIndex: 0, ItemID: "ct_1", Delta: "dir",
})
require.Contains(t, d, "output_index")
require.EqualValues(t, 0, d["output_index"])
require.Equal(t, "dir", d["delta"])
done := marshalEvent(t, ResponsesStreamEvent{
Type: "response.custom_tool_call_input.done", OutputIndex: 0, ItemID: "ct_1", CallID: "call_1", Name: "exec", Input: "dir",
})
require.Contains(t, done, "output_index")
require.EqualValues(t, 0, done["output_index"])
require.Equal(t, "dir", done["input"])
require.NotContains(t, done, "delta")
}
// TestWire_UnknownEventFallsBackToDefault ensures non-streamed event types keep
// default marshalling (the response object is preserved).
func TestWire_UnknownEventFallsBackToDefault(t *testing.T) {
@@ -111,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)
}
@@ -100,15 +100,16 @@ func anthropicUsageFromResponsesUsage(usage *ResponsesUsage) AnthropicUsage {
cachedTokens = usage.InputTokensDetails.CachedTokens
}
inputTokens := usage.InputTokens - cachedTokens
inputTokens := usage.InputTokens - cachedTokens - usage.CacheCreationInputTokens
if inputTokens < 0 {
inputTokens = 0
}
return AnthropicUsage{
InputTokens: inputTokens,
OutputTokens: usage.OutputTokens,
CacheReadInputTokens: cachedTokens,
InputTokens: inputTokens,
OutputTokens: usage.OutputTokens,
CacheReadInputTokens: cachedTokens,
CacheCreationInputTokens: usage.CacheCreationInputTokens,
}
}
@@ -181,9 +182,10 @@ type ResponsesEventToAnthropicState struct {
// OutputIndexToBlockIdx maps Responses output_index → Anthropic content block index.
OutputIndexToBlockIdx map[int]int
InputTokens int
OutputTokens int
CacheReadInputTokens int
InputTokens int
OutputTokens int
CacheReadInputTokens int
CacheCreationInputTokens int
ResponseID string
Model string
@@ -258,9 +260,10 @@ func FinalizeResponsesAnthropicStream(state *ResponsesEventToAnthropicState) []A
StopReason: stopReason,
},
Usage: &AnthropicUsage{
InputTokens: state.InputTokens,
OutputTokens: state.OutputTokens,
CacheReadInputTokens: state.CacheReadInputTokens,
InputTokens: state.InputTokens,
OutputTokens: state.OutputTokens,
CacheReadInputTokens: state.CacheReadInputTokens,
CacheCreationInputTokens: state.CacheCreationInputTokens,
},
},
AnthropicStreamEvent{Type: "message_stop"},
@@ -410,10 +413,6 @@ func resToAnthHandleFuncArgsDelta(evt *ResponsesStreamEvent, state *ResponsesEve
return nil
}
if state.CurrentBlockType == "tool_use" && state.CurrentToolName == "Read" {
state.CurrentToolArgs += evt.Delta
return nil
}
if state.CurrentBlockType == "tool_use" {
state.CurrentToolHadDelta = true
}
@@ -578,6 +577,7 @@ func resToAnthHandleCompleted(evt *ResponsesStreamEvent, state *ResponsesEventTo
state.InputTokens = usage.InputTokens
state.OutputTokens = usage.OutputTokens
state.CacheReadInputTokens = usage.CacheReadInputTokens
state.CacheCreationInputTokens = usage.CacheCreationInputTokens
}
if evt.Response != nil {
if evt.Response.Usage != nil {
@@ -585,6 +585,7 @@ func resToAnthHandleCompleted(evt *ResponsesStreamEvent, state *ResponsesEventTo
state.InputTokens = usage.InputTokens
state.OutputTokens = usage.OutputTokens
state.CacheReadInputTokens = usage.CacheReadInputTokens
state.CacheCreationInputTokens = usage.CacheCreationInputTokens
}
switch evt.Response.Status {
case "incomplete":
@@ -605,9 +606,10 @@ func resToAnthHandleCompleted(evt *ResponsesStreamEvent, state *ResponsesEventTo
StopReason: stopReason,
},
Usage: &AnthropicUsage{
InputTokens: state.InputTokens,
OutputTokens: state.OutputTokens,
CacheReadInputTokens: state.CacheReadInputTokens,
InputTokens: state.InputTokens,
OutputTokens: state.OutputTokens,
CacheReadInputTokens: state.CacheReadInputTokens,
CacheCreationInputTokens: state.CacheCreationInputTokens,
},
},
AnthropicStreamEvent{Type: "message_stop"},
@@ -0,0 +1,84 @@
package apicompat
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestResToAnthFuncArgsDelta_ReadToolStreamsDeltas(t *testing.T) {
state := NewResponsesEventToAnthropicState()
state.MessageStartSent = true
state.CurrentBlockType = "tool_use"
state.CurrentToolName = "Read"
state.OutputIndexToBlockIdx = map[int]int{0: 0}
evt := &ResponsesStreamEvent{
Type: "response.function_call_arguments.delta",
OutputIndex: 0,
Delta: `{"file_path":"/tmp/test.go"}`,
}
events := ResponsesEventToAnthropicEvents(evt, state)
require.Len(t, events, 1, "Read tool delta must produce content_block_delta")
assert.Equal(t, "content_block_delta", events[0].Type)
assert.Equal(t, "input_json_delta", events[0].Delta.Type)
assert.Equal(t, `{"file_path":"/tmp/test.go"}`, events[0].Delta.PartialJSON)
assert.True(t, state.CurrentToolHadDelta, "Read deltas should set CurrentToolHadDelta")
}
func TestResToAnthFuncArgsDelta_ReadToolWithoutDone(t *testing.T) {
state := NewResponsesEventToAnthropicState()
state.MessageStartSent = true
state.ContentBlockIndex = 0
state.ContentBlockOpen = true
state.CurrentBlockType = "tool_use"
state.CurrentToolName = "Read"
state.OutputIndexToBlockIdx = map[int]int{0: 0}
delta := &ResponsesStreamEvent{
Type: "response.function_call_arguments.delta",
OutputIndex: 0,
Delta: `{"file_path":"/tmp/test.go"}`,
}
events := ResponsesEventToAnthropicEvents(delta, state)
require.Len(t, events, 1, "delta should be streamed")
completed := &ResponsesStreamEvent{
Type: "response.completed",
Response: &ResponsesResponse{
Status: "completed",
},
}
events = ResponsesEventToAnthropicEvents(completed, state)
hasStop := false
for _, e := range events {
if e.Type == "content_block_stop" {
hasStop = true
}
}
assert.True(t, hasStop, "block should be closed even without .done event")
}
func TestResToAnthFuncArgsDelta_NonReadToolUnchanged(t *testing.T) {
state := NewResponsesEventToAnthropicState()
state.MessageStartSent = true
state.CurrentBlockType = "tool_use"
state.CurrentToolName = "Write"
state.OutputIndexToBlockIdx = map[int]int{0: 0}
evt := &ResponsesStreamEvent{
Type: "response.function_call_arguments.delta",
OutputIndex: 0,
Delta: `{"file_path":"/tmp/out.txt","content":"hello"}`,
}
events := ResponsesEventToAnthropicEvents(evt, state)
require.Len(t, events, 1)
assert.Equal(t, "content_block_delta", events[0].Type)
assert.True(t, state.CurrentToolHadDelta)
}
@@ -89,8 +89,13 @@ func ResponsesToChatCompletions(resp *ResponsesResponse, model string) *ChatComp
func responsesStatusToChatFinishReason(status string, details *ResponsesIncompleteDetails, toolCalls []ChatToolCall) string {
switch status {
case "incomplete":
if details != nil && details.Reason == "max_output_tokens" {
return "length"
if details != nil {
switch details.Reason {
case "max_output_tokens":
return "length"
case "content_filter":
return "content_filter"
}
}
return "stop"
case "completed":
@@ -299,8 +304,13 @@ func resToChatHandleCompleted(evt *ResponsesStreamEvent, state *ResponsesEventTo
switch evt.Response.Status {
case "incomplete":
if evt.Response.IncompleteDetails != nil && evt.Response.IncompleteDetails.Reason == "max_output_tokens" {
finishReason = "length"
if evt.Response.IncompleteDetails != nil {
switch evt.Response.IncompleteDetails.Reason {
case "max_output_tokens":
finishReason = "length"
case "content_filter":
finishReason = "content_filter"
}
}
case "completed":
if state.SawToolCall {
@@ -0,0 +1,122 @@
package apicompat
import (
"testing"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestAnthropicStreamingMaxTokens_MapsToIncomplete(t *testing.T) {
state := NewAnthropicEventToResponsesState()
AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_start",
Message: &AnthropicResponse{ID: "msg_test", Model: "claude-opus-4-6", Role: "assistant"},
}, state)
AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_delta",
Delta: &AnthropicDelta{
StopReason: "max_tokens",
},
Usage: &AnthropicUsage{OutputTokens: 4096},
}, state)
require.Equal(t, "max_tokens", state.StopReason)
events := AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_stop",
}, state)
var completed *ResponsesStreamEvent
for i := range events {
if events[i].Type == "response.completed" || events[i].Type == "response.incomplete" {
completed = &events[i]
break
}
}
require.NotNil(t, completed, "should have terminal event")
assert.Equal(t, "response.incomplete", completed.Type)
require.NotNil(t, completed.Response)
assert.Equal(t, "incomplete", completed.Response.Status)
require.NotNil(t, completed.Response.IncompleteDetails)
assert.Equal(t, "max_output_tokens", completed.Response.IncompleteDetails.Reason)
}
func TestAnthropicStreamingEndTurn_MapsToCompleted(t *testing.T) {
state := NewAnthropicEventToResponsesState()
AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_start",
Message: &AnthropicResponse{ID: "msg_test", Model: "claude-opus-4-6", Role: "assistant"},
}, state)
AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_delta",
Delta: &AnthropicDelta{StopReason: "end_turn"},
Usage: &AnthropicUsage{OutputTokens: 100},
}, state)
events := AnthropicEventToResponsesEvents(&AnthropicStreamEvent{
Type: "message_stop",
}, state)
var completed *ResponsesStreamEvent
for i := range events {
if events[i].Type == "response.completed" {
completed = &events[i]
break
}
}
require.NotNil(t, completed)
assert.Equal(t, "completed", completed.Response.Status)
assert.Nil(t, completed.Response.IncompleteDetails)
}
func TestResponsesToChatCompletions_ContentFilter(t *testing.T) {
resp := &ResponsesResponse{
ID: "resp_cf",
Status: "incomplete",
IncompleteDetails: &ResponsesIncompleteDetails{
Reason: "content_filter",
},
Output: []ResponsesOutput{{
Type: "message",
Content: []ResponsesContentPart{{Type: "output_text", Text: "partial"}},
}},
Usage: &ResponsesUsage{InputTokens: 10, OutputTokens: 5},
}
cc := ResponsesToChatCompletions(resp, "gpt-5.5")
require.Len(t, cc.Choices, 1)
assert.Equal(t, "content_filter", cc.Choices[0].FinishReason)
}
func TestResponsesToChatCompletionsStreaming_ContentFilter(t *testing.T) {
state := NewResponsesEventToChatState()
state.ID = "resp_cf"
state.Model = "gpt-5.5"
state.SentRole = true
events := ResponsesEventToChatChunks(&ResponsesStreamEvent{
Type: "response.completed",
Response: &ResponsesResponse{
ID: "resp_cf",
Status: "incomplete",
IncompleteDetails: &ResponsesIncompleteDetails{
Reason: "content_filter",
},
},
}, state)
hasContentFilter := false
for _, chunk := range events {
for _, choice := range chunk.Choices {
if choice.FinishReason != nil && *choice.FinishReason == "content_filter" {
hasContentFilter = true
}
}
}
assert.True(t, hasContentFilter, "streaming content_filter should map to finish_reason content_filter")
}
+101 -1
View File
@@ -249,11 +249,31 @@ type ResponsesContentPart struct {
// ResponsesTool describes a tool in the Responses API.
type ResponsesTool struct {
Type string `json:"type"` // "function" | "web_search" | "local_shell" etc.
Type string `json:"type"` // "function" | "custom" | "web_search" | "local_shell" etc.
Name string `json:"name,omitempty"`
Description string `json:"description,omitempty"`
Parameters json.RawMessage `json:"parameters,omitempty"`
Strict *bool `json:"strict,omitempty"`
// type=namespace 的子工具列表(tools 与 children 二选一,语义相同)。
Tools []ResponsesTool `json:"tools,omitempty"`
Children []ResponsesTool `json:"children,omitempty"`
}
// UnmarshalJSON 容忍字符串形式的工具声明:codex 会以 "name" 简写声明 custom 工具,
func (t *ResponsesTool) UnmarshalJSON(data []byte) error {
var name string
if err := json.Unmarshal(data, &name); err == nil {
*t = ResponsesTool{Type: "custom", Name: name}
return nil
}
type alias ResponsesTool
var a alias
if err := json.Unmarshal(data, &a); err != nil {
return err
}
*t = ResponsesTool(a)
return nil
}
// ResponsesResponse is the non-streaming response from POST /v1/responses.
@@ -301,11 +321,88 @@ type ResponsesOutput struct {
CallID string `json:"call_id,omitempty"`
Name string `json:"name,omitempty"`
Arguments string `json:"arguments,omitempty"`
// 来源为 namespace 子工具时的归属命名空间(codex 按 namespace+name 路由该调用)。
Namespace string `json:"namespace,omitempty"`
// type=custom_tool_call(custom/freeform 工具,input 为自由文本)
Input string `json:"input,omitempty"`
// type=web_search_call
Action *WebSearchAction `json:"action,omitempty"`
}
// MarshalJSON 处理 tool_search_call 项的线上形态(复用 CallID/Arguments 字段):
// execution 固定为 "client"(codex 的必填字段,非 client 的调用会被静默忽略),
// arguments 是 JSON 对象而非 function_call 语义下的字符串。其余类型走默认结构体
// 序列化,输出逐字节不变。
func (o ResponsesOutput) MarshalJSON() ([]byte, error) {
type responsesOutputAlias ResponsesOutput
if o.Type != "tool_search_call" {
return json.Marshal(responsesOutputAlias(o))
}
m := map[string]any{
"type": o.Type,
"id": o.ID,
"call_id": o.CallID,
"execution": "client",
"arguments": toolSearchCallArgumentsJSON(o.Arguments),
}
if o.Status != "" {
m["status"] = o.Status
}
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"
@@ -444,6 +541,9 @@ type ResponsesStreamEvent struct {
Name string `json:"name,omitempty"`
Arguments string `json:"arguments,omitempty"`
// response.custom_tool_call_input.done
Input string `json:"input,omitempty"`
// response.reasoning_summary_text.delta / done
// Reuses Text/Delta fields above, SummaryIndex identifies which summary part
SummaryIndex int `json:"summary_index,omitempty"`
+2
View File
@@ -25,6 +25,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
)
@@ -92,6 +93,7 @@ func buildClient(opts Options) (*http.Client, error) {
if opts.ValidateResolvedIP && !opts.AllowPrivateHosts {
rt = newValidatedTransport(transport)
}
rt = servertiming.WrapRoundTripper(rt)
return &http.Client{
Transport: rt,
Timeout: opts.Timeout,
@@ -38,7 +38,7 @@ func (p PaginationParams) Offset() int {
if p.Page < 1 {
p.Page = 1
}
return (p.Page - 1) * p.PageSize
return (p.Page - 1) * p.Limit()
}
// Limit 获取限制数
@@ -69,3 +69,30 @@ func TestPaginationParamsLimit(t *testing.T) {
})
}
}
func TestPaginationParamsOffsetUsesNormalizedLimit(t *testing.T) {
t.Parallel()
tests := []struct {
name string
page int
pageSize int
want int
}{
{name: "invalid page uses first page", page: 0, pageSize: 50, want: 0},
{name: "zero page size uses default", page: 2, pageSize: 0, want: 20},
{name: "negative page size uses default", page: 2, pageSize: -1, want: 20},
{name: "normal values", page: 3, pageSize: 50, want: 100},
{name: "page size beyond max is clamped", page: 2, pageSize: 1500, want: 1000},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
params := PaginationParams{Page: tt.page, PageSize: tt.pageSize}
if got := params.Offset(); got != tt.want {
t.Fatalf("Offset() for Page=%d, PageSize=%d = %d, want %d", tt.page, tt.pageSize, got, tt.want)
}
})
}
}
@@ -0,0 +1,348 @@
package servertiming
import (
"context"
"fmt"
"sort"
"strconv"
"strings"
"sync"
"time"
)
const (
HeaderName = "Server-Timing"
AdminUIHeader = "X-Admin-UI-Request"
MetricDatabase = "db"
MetricRedis = "redis"
dependencyPrefix = "dep_"
maxMetricNameLength = 48
maxIntervals = 2048
maxHeaderLength = 4096
)
type contextKey struct{}
type interval struct {
start time.Time
end time.Time
}
type metric struct {
count int64
intervals []interval
}
// Collector stores request-scoped timing samples. It is safe for concurrent use.
type Collector struct {
startedAt time.Time
mu sync.Mutex
metrics map[string]*metric
cacheStatus string
}
// New creates a collector whose total duration starts at startedAt.
func New(startedAt time.Time) *Collector {
if startedAt.IsZero() {
startedAt = time.Now()
}
return &Collector{
startedAt: startedAt,
metrics: make(map[string]*metric),
}
}
// WithCollector attaches a collector to a context.
func WithCollector(ctx context.Context, collector *Collector) context.Context {
if ctx == nil {
ctx = context.Background()
}
if collector == nil {
return ctx
}
return context.WithValue(ctx, contextKey{}, collector)
}
// FromContext returns the request timing collector, when one is active.
func FromContext(ctx context.Context) (*Collector, bool) {
if ctx == nil {
return nil, false
}
collector, ok := ctx.Value(contextKey{}).(*Collector)
return collector, ok && collector != nil
}
// Active reports whether timing collection is enabled for this request.
func Active(ctx context.Context) bool {
_, ok := FromContext(ctx)
return ok
}
// Record adds a completed interval and operation count to a metric.
func Record(ctx context.Context, name string, startedAt, endedAt time.Time, count int) {
collector, ok := FromContext(ctx)
if !ok {
return
}
collector.Record(name, startedAt, endedAt, count)
}
// RecordInterval adds timing without incrementing the operation count. It is
// useful when one logical operation has multiple blocking driver calls.
func RecordInterval(ctx context.Context, name string, startedAt, endedAt time.Time) {
collector, ok := FromContext(ctx)
if !ok {
return
}
collector.record(name, startedAt, endedAt, 0)
}
// Record adds a completed interval directly to the collector.
func (c *Collector) Record(name string, startedAt, endedAt time.Time, count int) {
if count <= 0 {
count = 1
}
c.record(name, startedAt, endedAt, count)
}
func (c *Collector) record(name string, startedAt, endedAt time.Time, count int) {
name = normalizeMetricName(name)
if c == nil || name == "" || startedAt.IsZero() || endedAt.Before(startedAt) {
return
}
if count < 0 {
count = 0
}
c.mu.Lock()
m := c.metrics[name]
if m == nil {
m = &metric{}
c.metrics[name] = m
}
m.count += int64(count)
if len(m.intervals) < maxIntervals {
m.intervals = append(m.intervals, interval{start: startedAt, end: endedAt})
}
c.mu.Unlock()
}
// Observe starts a metric span and returns an idempotent completion function.
func Observe(ctx context.Context, name string) func() {
collector, ok := FromContext(ctx)
name = normalizeMetricName(name)
if !ok || name == "" {
return func() {}
}
startedAt := time.Now()
var once sync.Once
return func() {
once.Do(func() {
collector.Record(name, startedAt, time.Now(), 1)
})
}
}
// ObserveDependency starts a named external dependency span.
func ObserveDependency(ctx context.Context, module string) func() {
return Observe(ctx, dependencyMetricName(module))
}
// RecordDependency records a completed external dependency interval.
func RecordDependency(ctx context.Context, module string, startedAt, endedAt time.Time) {
Record(ctx, dependencyMetricName(module), startedAt, endedAt, 1)
}
// SetCacheStatus records the response-cache outcome for the request.
func SetCacheStatus(ctx context.Context, status string) {
collector, ok := FromContext(ctx)
if !ok {
return
}
status = normalizeCacheStatus(status)
if status == "" {
return
}
collector.mu.Lock()
collector.cacheStatus = status
collector.mu.Unlock()
}
// HeaderValue renders a bounded, deterministic Server-Timing header.
func HeaderValue(ctx context.Context, endedAt time.Time, cacheStatus string) string {
collector, ok := FromContext(ctx)
if !ok {
return ""
}
return collector.HeaderValue(endedAt, cacheStatus)
}
// HeaderValue renders a bounded, deterministic Server-Timing header.
func (c *Collector) HeaderValue(endedAt time.Time, cacheStatus string) string {
if c == nil {
return ""
}
if endedAt.IsZero() {
endedAt = time.Now()
}
if endedAt.Before(c.startedAt) {
endedAt = c.startedAt
}
c.mu.Lock()
metrics := make(map[string]metric, len(c.metrics))
allIntervals := make([]interval, 0)
dependencyIntervals := make([]interval, 0)
var dependencyCount int64
for name, source := range c.metrics {
copied := metric{count: source.count, intervals: append([]interval(nil), source.intervals...)}
metrics[name] = copied
allIntervals = append(allIntervals, copied.intervals...)
if strings.HasPrefix(name, dependencyPrefix) {
dependencyIntervals = append(dependencyIntervals, copied.intervals...)
dependencyCount += copied.count
}
}
storedCacheStatus := c.cacheStatus
c.mu.Unlock()
total := endedAt.Sub(c.startedAt)
blocked := unionDuration(allIntervals, c.startedAt, endedAt)
app := total - blocked
if app < 0 {
app = 0
}
cacheStatus = normalizeCacheStatus(cacheStatus)
if cacheStatus == "" {
cacheStatus = normalizeCacheStatus(storedCacheStatus)
}
if cacheStatus == "" {
cacheStatus = "bypass"
}
database := metrics[MetricDatabase]
redisMetric := metrics[MetricRedis]
parts := []string{
"total;dur=" + formatDuration(total),
"app;dur=" + formatDuration(app),
fmt.Sprintf("db;dur=%s;desc=\"queries=%d\"", formatDuration(unionDuration(database.intervals, c.startedAt, endedAt)), database.count),
fmt.Sprintf("redis;dur=%s;desc=\"commands=%d\"", formatDuration(unionDuration(redisMetric.intervals, c.startedAt, endedAt)), redisMetric.count),
"cache;desc=\"" + cacheStatus + "\"",
fmt.Sprintf("deps;dur=%s;desc=\"calls=%d\"", formatDuration(unionDuration(dependencyIntervals, c.startedAt, endedAt)), dependencyCount),
}
dependencyNames := make([]string, 0)
for name := range metrics {
if strings.HasPrefix(name, dependencyPrefix) {
dependencyNames = append(dependencyNames, name)
}
}
sort.Strings(dependencyNames)
for _, name := range dependencyNames {
m := metrics[name]
part := fmt.Sprintf("%s;dur=%s;desc=\"calls=%d\"", name, formatDuration(unionDuration(m.intervals, c.startedAt, endedAt)), m.count)
candidate := strings.Join(append(parts, part), ", ")
if len(candidate) > maxHeaderLength {
break
}
parts = append(parts, part)
}
return strings.Join(parts, ", ")
}
func dependencyMetricName(module string) string {
module = normalizeMetricName(module)
module = strings.TrimPrefix(module, dependencyPrefix)
if module == "" {
module = "http"
}
return dependencyPrefix + module
}
func normalizeMetricName(name string) string {
name = strings.ToLower(strings.TrimSpace(name))
if name == "" {
return ""
}
var b strings.Builder
b.Grow(min(len(name), maxMetricNameLength))
for _, r := range name {
if b.Len() >= maxMetricNameLength {
break
}
switch {
case r >= 'a' && r <= 'z', r >= '0' && r <= '9':
_, _ = b.WriteRune(r)
case r == '_' || r == '-':
_ = b.WriteByte('_')
}
}
return strings.Trim(b.String(), "_")
}
func normalizeCacheStatus(status string) string {
switch strings.ToLower(strings.TrimSpace(status)) {
case "hit":
return "hit"
case "miss":
return "miss"
case "bypass":
return "bypass"
default:
return ""
}
}
func unionDuration(intervals []interval, lowerBound, upperBound time.Time) time.Duration {
if len(intervals) == 0 || !upperBound.After(lowerBound) {
return 0
}
normalized := make([]interval, 0, len(intervals))
for _, item := range intervals {
start := item.start
end := item.end
if start.Before(lowerBound) {
start = lowerBound
}
if end.After(upperBound) {
end = upperBound
}
if end.After(start) {
normalized = append(normalized, interval{start: start, end: end})
}
}
if len(normalized) == 0 {
return 0
}
sort.Slice(normalized, func(i, j int) bool {
return normalized[i].start.Before(normalized[j].start)
})
currentStart := normalized[0].start
currentEnd := normalized[0].end
var total time.Duration
for _, item := range normalized[1:] {
if !item.start.After(currentEnd) {
if item.end.After(currentEnd) {
currentEnd = item.end
}
continue
}
total += currentEnd.Sub(currentStart)
currentStart = item.start
currentEnd = item.end
}
total += currentEnd.Sub(currentStart)
return total
}
func formatDuration(value time.Duration) string {
if value < 0 {
value = 0
}
return strconv.FormatFloat(float64(value)/float64(time.Millisecond), 'f', 1, 64)
}
@@ -0,0 +1,129 @@
package servertiming
import (
"context"
"fmt"
"strings"
"sync"
"testing"
"time"
)
func TestCollectorHeaderValueAggregatesIntervals(t *testing.T) {
startedAt := time.Unix(100, 0)
collector := New(startedAt)
collector.Record(MetricDatabase, startedAt.Add(10*time.Millisecond), startedAt.Add(40*time.Millisecond), 2)
collector.Record(MetricRedis, startedAt.Add(30*time.Millisecond), startedAt.Add(50*time.Millisecond), 3)
collector.Record(dependencyMetricName("openai"), startedAt.Add(70*time.Millisecond), startedAt.Add(100*time.Millisecond), 1)
collector.Record(dependencyMetricName("github"), startedAt.Add(60*time.Millisecond), startedAt.Add(90*time.Millisecond), 1)
got := collector.HeaderValue(startedAt.Add(120*time.Millisecond), "miss")
want := `total;dur=120.0, app;dur=40.0, db;dur=30.0;desc="queries=2", redis;dur=20.0;desc="commands=3", cache;desc="miss", deps;dur=40.0;desc="calls=2", dep_github;dur=30.0;desc="calls=1", dep_openai;dur=30.0;desc="calls=1"`
if got != want {
t.Fatalf("HeaderValue() = %q, want %q", got, want)
}
}
func TestRecordIntervalDoesNotIncrementCount(t *testing.T) {
startedAt := time.Unix(200, 0)
collector := New(startedAt)
ctx := WithCollector(context.Background(), collector)
Record(ctx, MetricDatabase, startedAt.Add(10*time.Millisecond), startedAt.Add(20*time.Millisecond), 1)
RecordInterval(ctx, MetricDatabase, startedAt.Add(30*time.Millisecond), startedAt.Add(40*time.Millisecond))
header := HeaderValue(ctx, startedAt.Add(100*time.Millisecond), "hit")
if !strings.Contains(header, `db;dur=20.0;desc="queries=1"`) {
t.Fatalf("header %q does not contain one query with both blocking intervals", header)
}
if !strings.Contains(header, "app;dur=80.0") {
t.Fatalf("header %q does not subtract the interval union from app time", header)
}
}
func TestCollectorCacheStatusFallback(t *testing.T) {
startedAt := time.Unix(300, 0)
collector := New(startedAt)
ctx := WithCollector(context.Background(), collector)
SetCacheStatus(ctx, " HIT ")
if got := HeaderValue(ctx, startedAt.Add(time.Millisecond), "invalid"); !strings.Contains(got, `cache;desc="hit"`) {
t.Fatalf("HeaderValue() = %q, want stored cache hit", got)
}
other := New(startedAt)
if got := other.HeaderValue(startedAt.Add(time.Millisecond), "invalid"); !strings.Contains(got, `cache;desc="bypass"`) {
t.Fatalf("HeaderValue() = %q, want cache bypass", got)
}
}
func TestCollectorSanitizesDependencyMetric(t *testing.T) {
startedAt := time.Unix(400, 0)
collector := New(startedAt)
ctx := WithCollector(context.Background(), collector)
RecordDependency(ctx, "GitHub API\r\nInjected;dur=999", startedAt, startedAt.Add(time.Millisecond))
header := HeaderValue(ctx, startedAt.Add(2*time.Millisecond), "bypass")
if strings.ContainsAny(header, "\r\n") || strings.Contains(header, ";dur=999") {
t.Fatalf("unsafe metric content reached header: %q", header)
}
if !strings.Contains(header, "dep_githubapiinjecteddur999;dur=1.0") {
t.Fatalf("sanitized dependency metric missing from header: %q", header)
}
}
func TestCollectorBoundsHeaderLength(t *testing.T) {
startedAt := time.Unix(500, 0)
collector := New(startedAt)
for i := 0; i < 300; i++ {
collector.Record(
dependencyMetricName(fmt.Sprintf("module_%03d_with_a_deliberately_long_name", i)),
startedAt,
startedAt.Add(time.Millisecond),
1,
)
}
header := collector.HeaderValue(startedAt.Add(2*time.Millisecond), "bypass")
if len(header) > maxHeaderLength {
t.Fatalf("header length = %d, want <= %d", len(header), maxHeaderLength)
}
if !strings.Contains(header, "total;dur=2.0") || !strings.Contains(header, "deps;dur=1.0") {
t.Fatalf("bounded header lost fixed metrics: %q", header)
}
}
func TestCollectorConcurrentRecording(t *testing.T) {
startedAt := time.Now()
collector := New(startedAt)
ctx := WithCollector(context.Background(), collector)
const workers = 25
const recordsPerWorker = 100
var wg sync.WaitGroup
wg.Add(workers)
for i := 0; i < workers; i++ {
go func() {
defer wg.Done()
for j := 0; j < recordsPerWorker; j++ {
Record(ctx, MetricDatabase, startedAt, startedAt.Add(time.Microsecond), 1)
}
}()
}
wg.Wait()
header := HeaderValue(ctx, startedAt.Add(time.Millisecond), "bypass")
want := fmt.Sprintf(`queries=%d`, workers*recordsPerWorker)
if !strings.Contains(header, want) {
t.Fatalf("header %q does not contain %q", header, want)
}
}
func TestContextHelpersHandleMissingCollector(t *testing.T) {
if Active(context.Background()) {
t.Fatal("context without collector reported active")
}
if got := HeaderValue(context.Background(), time.Now(), "hit"); got != "" {
t.Fatalf("HeaderValue() = %q without collector, want empty", got)
}
}
Loaded 100 of 388 files, more files were not shown because too many files have changed in this diff. Show more