chore(openai): 同步上游并解决认证冲突
This commit is contained in:
388 files changed
+27587
-1591
No files matched your search
@@ -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:
|
||||
|
||||
@@ -116,6 +116,8 @@ backend/.installed
|
||||
# 其他
|
||||
# ===================
|
||||
tests
|
||||
!deploy/tests/
|
||||
!deploy/tests/**
|
||||
CLAUDE.md
|
||||
.claude
|
||||
scripts
|
||||
|
||||
@@ -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
@@ -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
@@ -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 @@
|
||||
0.1.150
|
||||
0.1.155
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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(", ")
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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).
|
||||
|
||||
@@ -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").
|
||||
|
||||
@@ -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
@@ -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))
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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=
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
// ──────────────────────────────────────────────────────────
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -164,7 +164,7 @@ var ProviderSet = wire.NewSet(
|
||||
admin.NewDashboardHandler,
|
||||
admin.NewUserHandler,
|
||||
admin.NewGroupHandler,
|
||||
admin.NewAccountHandler,
|
||||
admin.ProvideAccountHandler,
|
||||
admin.NewAnnouncementHandler,
|
||||
admin.NewDataManagementHandler,
|
||||
admin.NewBackupHandler,
|
||||
|
||||
@@ -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] = ©Call
|
||||
stored = ©Call
|
||||
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")
|
||||
}
|
||||
@@ -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"`
|
||||
|
||||
@@ -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
Reference in new issue
Block a user