mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
19149ca196 | ||
|
|
831812b39d | ||
|
|
d98ac48eb9 | ||
|
|
72b29a7d4e | ||
|
|
14608dc6d4 | ||
|
|
daa4989124 | ||
|
|
8a95a46a69 | ||
|
|
9d498c2474 | ||
|
|
a2f802d409 | ||
|
|
b01196a7a8 | ||
|
|
e14fb2b6ff | ||
|
|
dd7a2b22f0 | ||
|
|
4e8ea7d568 | ||
|
|
a3f2b8fd87 | ||
|
|
23cdd20606 | ||
|
|
a62b821b5f | ||
|
|
774ff5d8c8 | ||
|
|
d2667393b3 | ||
|
|
e002fbb349 | ||
|
|
005bc5c8d4 | ||
|
|
bc4bd118d8 | ||
|
|
f3925db11b | ||
|
|
f9a467a4c0 | ||
|
|
8ce9288a6a | ||
|
|
c1e702be5e | ||
|
|
bc6b69289c | ||
|
|
5a0492bf88 | ||
|
|
fdc9d92fdb | ||
|
|
8cdd372b7e | ||
|
|
080a52121a | ||
|
|
f0e0b7e6d8 | ||
|
|
4c845ca697 | ||
|
|
95d27f2dce | ||
|
|
f2e7a55703 | ||
|
|
1ed9f59599 | ||
|
|
b4ad7fd363 | ||
|
|
ef2a22be70 | ||
|
|
c831bb979f | ||
|
|
3edcbfd120 | ||
|
|
1db10dc559 | ||
|
|
3f6b5c7bd7 | ||
|
|
2946281a43 | ||
|
|
335edde9c8 | ||
|
|
1a855d3e0e | ||
|
|
b1a6b80267 | ||
|
|
c3f759d13a | ||
|
|
539bfc8bad | ||
|
|
8b75dd5576 | ||
|
|
bd0f2d6405 | ||
|
|
71afe24710 | ||
|
|
500f241692 | ||
|
|
b92bbf0299 | ||
|
|
2594950993 | ||
|
|
9386ce8a13 | ||
|
|
be74deae73 | ||
|
|
2264a33085 | ||
|
|
d2585097f3 | ||
|
|
e86063155f | ||
|
|
1f7b7b9132 | ||
|
|
7043982ae3 | ||
|
|
57914967cb | ||
|
|
8bfbc5ca99 | ||
|
|
cb40288fb8 | ||
|
|
33766c299f | ||
|
|
13468b3fa5 | ||
|
|
e69e46ba02 | ||
|
|
eaf0691752 | ||
|
|
5fd030a01b | ||
|
|
fc3c5a1e0c | ||
|
|
e375623abf | ||
|
|
288d3a1cce | ||
|
|
ac685ccaf5 | ||
|
|
bc50c9d013 | ||
|
|
dda20ae9a6 | ||
|
|
c34dcd3700 | ||
|
|
7ed4e7e5e3 | ||
|
|
d2ac6b2c97 | ||
|
|
8f9296140f | ||
|
|
8bc0a277f2 | ||
|
|
456c619379 | ||
|
|
18e698bed6 | ||
|
|
26b5e87457 | ||
|
|
c2c19a7cbe | ||
|
|
2a75d7d238 | ||
|
|
1c3f2810af | ||
|
|
7c48f9a85f | ||
|
|
980a46ffb0 | ||
|
|
824f763a77 | ||
|
|
9b3c96dc68 | ||
|
|
52f2bc69fb | ||
|
|
067cf8b46e | ||
|
|
df9d9e2e40 | ||
|
|
0f7f8a317e | ||
|
|
d2b080e88b | ||
|
|
d11bdb13f5 | ||
|
|
29e39b96e5 | ||
|
|
08e994ad86 | ||
|
|
bc2244c83f | ||
|
|
26abd19a28 | ||
|
|
c29b50394f | ||
|
|
3d23cc399f | ||
|
|
c1fd1cb44e | ||
|
|
7e13b6d039 | ||
|
|
c993490a82 | ||
|
|
4d91f39474 | ||
|
|
af6bd44d2a | ||
|
|
541ca3bc27 | ||
|
|
7f5d067af2 | ||
|
|
174ea22eeb | ||
|
|
9fc006546c | ||
|
|
b960ec1980 | ||
|
|
60732a2e8c | ||
|
|
9d1064badd | ||
|
|
2b5944fd6c | ||
|
|
43e95dfa45 | ||
|
|
ca82fd6e98 | ||
|
|
115116e8bf | ||
|
|
7947619cc3 |
+19
-5
@@ -18,7 +18,9 @@ ARG NPM_CONFIG_REGISTRY=
|
||||
# -----------------------------------------------------------------------------
|
||||
# Stage 1: Frontend Builder
|
||||
# -----------------------------------------------------------------------------
|
||||
FROM ${NODE_IMAGE} AS frontend-builder
|
||||
# --platform=$BUILDPLATFORM: the frontend output is JS (arch-neutral), so build
|
||||
# it on the native host arch instead of under QEMU emulation for the target.
|
||||
FROM --platform=${BUILDPLATFORM} ${NODE_IMAGE} AS frontend-builder
|
||||
ARG NPM_CONFIG_REGISTRY
|
||||
|
||||
WORKDIR /app/frontend
|
||||
@@ -44,7 +46,11 @@ RUN pnpm run build
|
||||
# -----------------------------------------------------------------------------
|
||||
# Stage 2: Backend Builder
|
||||
# -----------------------------------------------------------------------------
|
||||
FROM ${GOLANG_IMAGE} AS backend-builder
|
||||
# --platform=$BUILDPLATFORM: run the Go toolchain on the native host arch and
|
||||
# cross-compile to the target arch below. The binary is CGO_ENABLED=0, so this
|
||||
# is a clean pure-Go cross-compile — no QEMU emulation of go mod download / go
|
||||
# build (emulated networking here was dropping module fetches with EOF).
|
||||
FROM --platform=${BUILDPLATFORM} ${GOLANG_IMAGE} AS backend-builder
|
||||
|
||||
# Build arguments for version info (set by CI)
|
||||
ARG VERSION=
|
||||
@@ -52,6 +58,9 @@ ARG COMMIT=docker
|
||||
ARG DATE
|
||||
ARG GOPROXY
|
||||
ARG GOSUMDB
|
||||
# Populated by buildx from the --platform target (e.g. linux/amd64).
|
||||
ARG TARGETOS
|
||||
ARG TARGETARCH
|
||||
|
||||
ENV GOPROXY=${GOPROXY}
|
||||
ENV GOSUMDB=${GOSUMDB}
|
||||
@@ -63,7 +72,10 @@ WORKDIR /app/backend
|
||||
|
||||
# Copy go mod files first (better caching)
|
||||
COPY backend/go.mod backend/go.sum ./
|
||||
RUN go mod download
|
||||
# Cache mount keeps the module cache across builds so a transient CDN blip on
|
||||
# retry resumes instead of re-fetching every zip from scratch.
|
||||
RUN --mount=type=cache,id=sub2api-gomod,target=/go/pkg/mod \
|
||||
go mod download
|
||||
|
||||
# Copy backend source first
|
||||
COPY backend/ ./
|
||||
@@ -73,10 +85,12 @@ COPY --from=frontend-builder /app/backend/internal/web/dist ./internal/web/dist
|
||||
|
||||
# Build the binary (BuildType=release for CI builds, embed frontend)
|
||||
# Version precedence: build arg VERSION > exact git tag > cmd/server/VERSION
|
||||
RUN VERSION_VALUE="${VERSION}" && \
|
||||
RUN --mount=type=cache,id=sub2api-gomod,target=/go/pkg/mod \
|
||||
--mount=type=cache,id=sub2api-gobuild,target=/root/.cache/go-build \
|
||||
VERSION_VALUE="${VERSION}" && \
|
||||
if [ -z "${VERSION_VALUE}" ]; then VERSION_VALUE="$(./scripts/resolve-version.sh)"; fi && \
|
||||
DATE_VALUE="${DATE:-$(date -u +%Y-%m-%dT%H:%M:%SZ)}" && \
|
||||
CGO_ENABLED=0 GOOS=linux go build \
|
||||
CGO_ENABLED=0 GOOS=${TARGETOS:-linux} GOARCH=${TARGETARCH} go build \
|
||||
-tags embed \
|
||||
-ldflags="-s -w -X main.Version=${VERSION_VALUE} -X main.Commit=${COMMIT} -X main.Date=${DATE_VALUE} -X main.BuildType=release" \
|
||||
-trimpath \
|
||||
|
||||
@@ -169,6 +169,18 @@ Model authenticity: no content intervention or secondary filtering — experienc
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://console.claudeapi.com/agent/register/drTKjyn6wGLK061Z?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default"><img src="assets/partners/logos/claudeapi.jpg" alt="claudeapi" width="150"></a></td>
|
||||
<td>Thanks to Claude API for sponsoring this project! <a href="https://console.claudeapi.com/agent/register/drTKjyn6wGLK061Z?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default">Claude API</a> is an official-channel API provider focused on Claude models. Built on official Anthropic keys and the official AWS Bedrock channel, it delivers a stable integration experience for Claude Code and Agent applications, supports the full Claude model lineup, and retains official capabilities such as Tool Use and long context. The service involves no reverse engineering and no model degradation, making it a great fit for heavy Claude Code users, Agent engineers, and enterprise engineering teams. Register via the <a href="https://console.claudeapi.com/agent/register/drTKjyn6wGLK061Z?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default">exclusive link</a> and contact customer support to receive free trial credits; invoicing and team onboarding are also supported.
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://code0.ai/agent/register/LgpIgl9JHtVG53V1?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default"><img src="assets/partners/logos/code0.jpg" alt="code0" width="150"></a></td>
|
||||
<td>Thanks to code0.ai for sponsoring this project! <a href="https://code0.ai/agent/register/LgpIgl9JHtVG53V1?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default">code0.ai</a> is an AI coding workbench for developers and engineering teams, aggregating mainstream agent coding capabilities such as Claude Code and Codex, and covering common development scenarios including code generation, project understanding, debugging and fixing, code review, and documentation generation. It suits independent developers, Agent engineers, open-source maintainers, and enterprise R&D teams, with invoicing and team onboarding supported. Register via the <a href="https://code0.ai/agent/register/LgpIgl9JHtVG53V1?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default">exclusive link</a> and contact customer support to receive free trial credits and experience a more efficient AI coding workflow.
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
</table>
|
||||
|
||||
## Overview
|
||||
@@ -627,8 +639,20 @@ override this limit.
|
||||
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.
|
||||
the global cap.
|
||||
|
||||
Enable the v2 mode router before selecting an account-level WS mode such as
|
||||
`http_bridge`:
|
||||
|
||||
```yaml
|
||||
gateway:
|
||||
openai_ws:
|
||||
mode_router_v2_enabled: true
|
||||
```
|
||||
|
||||
Or set `GATEWAY_OPENAI_WS_MODE_ROUTER_V2_ENABLED=true` in the environment.
|
||||
Use `http_bridge` for client-WebSocket/upstream-HTTP operation when rolling out
|
||||
or mitigating upstream WebSocket issues.
|
||||
|
||||
#### ⚠️ Important: Creating the Admin Account
|
||||
|
||||
@@ -709,6 +733,7 @@ Sub2API supports both Grok subscription accounts through xAI OAuth and standard
|
||||
- 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`
|
||||
- JSON image-edit and video-generation requests accept image references in `image`, `images`, `reference_images`, and `mask` objects. Use `url` for xAI-compatible payloads; the legacy `image_url` field remains accepted and is normalized to `url` before forwarding.
|
||||
- Out of scope for this provider: TTS, transcription, browser automation, cookies, and Grok web scraping
|
||||
|
||||
### OAuth Configuration
|
||||
@@ -775,6 +800,10 @@ xAI quota is passive. Sub2API does not invent subscription quota values; it reco
|
||||
|
||||
`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.
|
||||
|
||||
New Grok image and video generation requests use a media-specific eligibility check. API-key accounts remain eligible. OAuth accounts require positive paid-entitlement evidence from the xAI billing probe; Free, forbidden, missing, malformed, and inconclusive billing observations are excluded from new media generation. Unobserved OAuth accounts are probed before the first media request is forwarded, and imports run the billing-first quota probe proactively. Chat requests and video status lookups are not affected by this media-only quarantine. If no eligible account remains, the media endpoint returns HTTP `503` with error type `grok_media_no_eligible_account`.
|
||||
|
||||
Administrators can override automatic media eligibility through the account create/update API by setting `extra.grok_media_eligible` to `false` (exclude) or `true` (force eligible). On update, set it to `null` to remove the override and return to automatic probe-based behavior; omitting the field preserves the current override. A weekly allowance period alone is not treated as a paid tier signal. Successful image responses must contain at least one actual image output; empty HTTP `200` responses trigger account failover instead of being counted and returned as successful generations.
|
||||
|
||||
---
|
||||
|
||||
## Antigravity Support
|
||||
|
||||
@@ -172,6 +172,18 @@
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://console.claudeapi.com/agent/register/drTKjyn6wGLK061Z?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default"><img src="assets/partners/logos/claudeapi.jpg" alt="claudeapi" width="150"></a></td>
|
||||
<td>感谢 Claude API 对本项目的赞助! <a href="https://console.claudeapi.com/agent/register/drTKjyn6wGLK061Z?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default">Claude API</a> 是专注 Claude 模型的官方渠道 API 服务商,基于 Anthropic 官方 Key 与 AWS Bedrock 官方渠道,提供稳定的 Claude Code 与 Agent 应用接入体验,支持 Claude 全系列模型,保留 Tool Use、长上下文等官方能力。服务非逆向、非降智,适合 Claude Code 深度用户、Agent 工程师与企业技术团队使用。通过<a href="https://console.claudeapi.com/agent/register/drTKjyn6wGLK061Z?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default">[专属链接]</a>注册后联系客服,可领取免费测试额度,并支持开票和团队对接。
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://code0.ai/agent/register/LgpIgl9JHtVG53V1?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default"><img src="assets/partners/logos/code0.jpg" alt="code0" width="150"></a></td>
|
||||
<td>感谢 code0.ai 对本项目的赞助! <a href="https://code0.ai/agent/register/LgpIgl9JHtVG53V1?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default">code0.ai</a> 是面向开发者与技术团队的 AI 编程工作台,聚合 Claude Code、Codex 等主流 Agent 编程能力,支持代码生成、项目理解、调试修复、代码审查与文档生成等常见研发场景。适合独立开发者、Agent 工程师、开源项目维护者和企业研发团队使用,支持开票和团队对接。通过<a href="https://code0.ai/agent/register/LgpIgl9JHtVG53V1?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default">[专属链接]</a>注册后联系客服,可领取免费测试额度,体验更高效的 AI 编程工作流。
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
</table>
|
||||
|
||||
## 项目概述
|
||||
|
||||
@@ -167,6 +167,18 @@
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://console.claudeapi.com/agent/register/drTKjyn6wGLK061Z?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default"><img src="assets/partners/logos/claudeapi.jpg" alt="claudeapi" width="150"></a></td>
|
||||
<td>Claude API のご支援に感謝します!<a href="https://console.claudeapi.com/agent/register/drTKjyn6wGLK061Z?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default">Claude API</a> は Claude モデルに特化した公式チャネルの API サービスプロバイダーで、Anthropic 公式キーと AWS Bedrock 公式チャネルをベースに、Claude Code や Agent アプリケーションへの安定した接続体験を提供します。Claude 全シリーズのモデルに対応し、Tool Use や長文コンテキストなどの公式機能もそのまま利用可能。リバースエンジニアリングやモデル劣化のないサービスで、Claude Code のヘビーユーザー、Agent エンジニア、企業の技術チームに最適です。<a href="https://console.claudeapi.com/agent/register/drTKjyn6wGLK061Z?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default">専用リンク</a>から登録後カスタマーサポートへご連絡いただくと、無料お試しクレジットを受け取れます。請求書発行やチーム導入にも対応しています。
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
<tr>
|
||||
<td width="180"><a href="https://code0.ai/agent/register/LgpIgl9JHtVG53V1?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default"><img src="assets/partners/logos/code0.jpg" alt="code0" width="150"></a></td>
|
||||
<td>code0.ai のご支援に感謝します!<a href="https://code0.ai/agent/register/LgpIgl9JHtVG53V1?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default">code0.ai</a> は開発者と技術チーム向けの AI プログラミングワークベンチで、Claude Code や Codex などの主要な Agent コーディング能力を集約し、コード生成、プロジェクト理解、デバッグと修正、コードレビュー、ドキュメント生成といった一般的な開発シーンをサポートします。個人開発者、Agent エンジニア、OSS メンテナー、企業の開発チームに最適で、請求書発行やチーム導入にも対応。<a href="https://code0.ai/agent/register/LgpIgl9JHtVG53V1?utm_source=zcf&utm_medium=partner&utm_campaign=zcf_2026&utm_content=default">専用リンク</a>から登録後カスタマーサポートへご連絡いただくと、無料お試しクレジットを受け取り、より効率的な AI プログラミングワークフローを体験できます。
|
||||
</td>
|
||||
</tr>
|
||||
|
||||
</table>
|
||||
|
||||
## 概要
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 3.6 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 3.2 KiB |
@@ -0,0 +1,20 @@
|
||||
# Ingress rejection log cleanup
|
||||
|
||||
This maintenance command removes historical admission rejections from
|
||||
`ops_error_logs` without matching unrelated authentication or upstream errors.
|
||||
It is a dry run unless `--execute` is supplied, and always requires an explicit
|
||||
RFC3339 cutoff.
|
||||
|
||||
```sh
|
||||
go run ./cmd/cleanup-ingress-reject-logs --before 2026-07-17T00:00:00Z
|
||||
go run ./cmd/cleanup-ingress-reject-logs --before 2026-07-17T00:00:00Z --execute
|
||||
```
|
||||
|
||||
Run the execute form only after every application instance has been upgraded so
|
||||
older instances cannot add new ingress rejection rows below the chosen cutoff.
|
||||
The classifier intentionally retains invariant failures such as
|
||||
`USER_NOT_FOUND`, database errors, quota/billing errors, and upstream failures.
|
||||
|
||||
After the rollout and cleanup are verified, run
|
||||
`backend/scripts/finalize-ingress-reject-cleanup.sql` in a maintenance window to
|
||||
remove the deprecated plaintext-key audit table and attribution columns.
|
||||
@@ -0,0 +1,218 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"flag"
|
||||
"fmt"
|
||||
"log"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
_ "github.com/Wei-Shaw/sub2api/ent/runtime"
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/repository"
|
||||
"github.com/lib/pq"
|
||||
)
|
||||
|
||||
const classifierVersion = "ingress-reject-v1"
|
||||
|
||||
type candidate struct {
|
||||
id int64
|
||||
statusCode int
|
||||
message string
|
||||
body string
|
||||
}
|
||||
|
||||
func main() {
|
||||
beforeRaw := flag.String("before", "", "required RFC3339 cutoff; only older rows are considered")
|
||||
execute := flag.Bool("execute", false, "delete matched rows (default is dry-run)")
|
||||
batchSize := flag.Int("batch-size", 5000, "scan/delete batch size (1-5000)")
|
||||
flag.Parse()
|
||||
|
||||
if *beforeRaw == "" {
|
||||
log.Fatal("--before is required")
|
||||
}
|
||||
before, err := time.Parse(time.RFC3339, *beforeRaw)
|
||||
if err != nil {
|
||||
log.Fatalf("invalid --before: %v", err)
|
||||
}
|
||||
if *batchSize < 1 || *batchSize > 5000 {
|
||||
log.Fatal("--batch-size must be between 1 and 5000")
|
||||
}
|
||||
|
||||
cfg, err := config.LoadForBootstrap()
|
||||
if err != nil {
|
||||
log.Fatalf("load config: %v", err)
|
||||
}
|
||||
client, db, err := repository.InitEnt(cfg)
|
||||
if err != nil {
|
||||
log.Fatalf("initialize database: %v", err)
|
||||
}
|
||||
defer func() { _ = client.Close() }()
|
||||
|
||||
ctx := context.Background()
|
||||
counts, scanned, matched, deleted, err := cleanup(ctx, db, before, *batchSize, *execute)
|
||||
if err != nil {
|
||||
log.Fatalf("cleanup failed: %v", err)
|
||||
}
|
||||
|
||||
digest := sha256.Sum256([]byte(classifierVersion))
|
||||
mode := "dry-run"
|
||||
if *execute {
|
||||
mode = "execute"
|
||||
}
|
||||
fmt.Printf("mode=%s before=%s classifier=%s scanned=%d matched=%d deleted=%d\n",
|
||||
mode, before.UTC().Format(time.RFC3339), hex.EncodeToString(digest[:]), scanned, matched, deleted)
|
||||
reasons := make([]string, 0, len(counts))
|
||||
for reason := range counts {
|
||||
reasons = append(reasons, reason)
|
||||
}
|
||||
sort.Strings(reasons)
|
||||
for _, reason := range reasons {
|
||||
fmt.Printf("reason=%s count=%d\n", reason, counts[reason])
|
||||
}
|
||||
if *execute && deleted > 0 {
|
||||
fmt.Println("cleanup complete; schedule VACUUM (ANALYZE) ops_error_logs during normal maintenance")
|
||||
}
|
||||
}
|
||||
|
||||
func cleanup(ctx context.Context, db *sql.DB, before time.Time, batchSize int, execute bool) (map[string]int64, int64, int64, int64, error) {
|
||||
counts := make(map[string]int64)
|
||||
var cursor, scanned, matched, deleted int64
|
||||
for {
|
||||
rows, err := db.QueryContext(ctx, `
|
||||
SELECT id, COALESCE(status_code, 0), COALESCE(error_message, ''), COALESCE(error_body, '')
|
||||
FROM ops_error_logs
|
||||
WHERE id > $1
|
||||
AND created_at < $2
|
||||
AND error_phase = 'auth'
|
||||
AND account_id IS NULL
|
||||
AND upstream_status_code IS NULL
|
||||
AND COALESCE(upstream_error_message, '') = ''
|
||||
AND COALESCE(upstream_error_detail, '') = ''
|
||||
ORDER BY id ASC
|
||||
LIMIT $3`, cursor, before, batchSize)
|
||||
if err != nil {
|
||||
return nil, scanned, matched, deleted, err
|
||||
}
|
||||
|
||||
batch := make([]candidate, 0, batchSize)
|
||||
for rows.Next() {
|
||||
var item candidate
|
||||
if err := rows.Scan(&item.id, &item.statusCode, &item.message, &item.body); err != nil {
|
||||
_ = rows.Close()
|
||||
return nil, scanned, matched, deleted, err
|
||||
}
|
||||
batch = append(batch, item)
|
||||
cursor = item.id
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
_ = rows.Close()
|
||||
return nil, scanned, matched, deleted, err
|
||||
}
|
||||
_ = rows.Close()
|
||||
if len(batch) == 0 {
|
||||
break
|
||||
}
|
||||
|
||||
ids := make([]int64, 0, len(batch))
|
||||
for _, item := range batch {
|
||||
scanned++
|
||||
if reason, ok := historicalIngressRejectReason(item); ok {
|
||||
matched++
|
||||
counts[reason]++
|
||||
ids = append(ids, item.id)
|
||||
}
|
||||
}
|
||||
if execute && len(ids) > 0 {
|
||||
result, err := db.ExecContext(ctx,
|
||||
`DELETE FROM ops_error_logs WHERE id = ANY($1) AND created_at < $2`, pq.Array(ids), before)
|
||||
if err != nil {
|
||||
return nil, scanned, matched, deleted, err
|
||||
}
|
||||
n, err := result.RowsAffected()
|
||||
if err != nil {
|
||||
return nil, scanned, matched, deleted, err
|
||||
}
|
||||
deleted += n
|
||||
}
|
||||
}
|
||||
return counts, scanned, matched, deleted, nil
|
||||
}
|
||||
|
||||
func historicalIngressRejectReason(item candidate) (string, bool) {
|
||||
code, message := parseErrorIdentity(item.body, item.message)
|
||||
switch code {
|
||||
case "API_KEY_REQUIRED":
|
||||
return "missing_key", true
|
||||
case "INVALID_API_KEY":
|
||||
return "invalid_key", true
|
||||
case "API_KEY_DISABLED":
|
||||
return "key_disabled", true
|
||||
case "USER_INACTIVE":
|
||||
return "user_inactive", true
|
||||
case "GROUP_DELETED":
|
||||
return "group_deleted", true
|
||||
case "GROUP_DISABLED":
|
||||
return "group_disabled", true
|
||||
case "GROUP_NOT_ALLOWED":
|
||||
return "group_forbidden", true
|
||||
case "ACCESS_DENIED":
|
||||
return "ip_acl_denied", true
|
||||
case "api_key_in_query_deprecated":
|
||||
return "query_key_deprecated", true
|
||||
}
|
||||
|
||||
normalized := strings.TrimSpace(message)
|
||||
switch {
|
||||
case normalized == "API key is required":
|
||||
return "missing_key", true
|
||||
case normalized == "Invalid API key":
|
||||
return "invalid_key", true
|
||||
case normalized == "API key is disabled":
|
||||
return "key_disabled", true
|
||||
case normalized == "User account is not active":
|
||||
return "user_inactive", true
|
||||
case normalized == "API Key 所属分组已删除":
|
||||
return "group_deleted", true
|
||||
case normalized == "API Key 所属分组已停用":
|
||||
return "group_disabled", true
|
||||
case normalized == "API Key 所属专属分组不再允许当前用户使用":
|
||||
return "group_forbidden", true
|
||||
case normalized == "API Key is not assigned to any group and cannot be used. Please contact the administrator to assign it to a group.":
|
||||
return "group_unassigned", true
|
||||
case strings.HasPrefix(normalized, "Access denied. Your IP is "):
|
||||
return "ip_acl_denied", true
|
||||
case normalized == "Query parameter api_key is deprecated. Use Authorization header or key instead.":
|
||||
return "query_key_deprecated", true
|
||||
default:
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
|
||||
func parseErrorIdentity(body, fallbackMessage string) (string, string) {
|
||||
var payload struct {
|
||||
Code string `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Error struct {
|
||||
Code json.RawMessage `json:"code"`
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(body), &payload); err != nil {
|
||||
return "", fallbackMessage
|
||||
}
|
||||
message := payload.Message
|
||||
if message == "" {
|
||||
message = payload.Error.Message
|
||||
}
|
||||
if message == "" {
|
||||
message = fallbackMessage
|
||||
}
|
||||
return strings.TrimSpace(payload.Code), message
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
package main
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestHistoricalIngressRejectReason(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
item candidate
|
||||
reason string
|
||||
match bool
|
||||
}{
|
||||
{name: "standard invalid key", item: candidate{body: `{"code":"INVALID_API_KEY","message":"Invalid API key"}`}, reason: "invalid_key", match: true},
|
||||
{name: "google missing key", item: candidate{body: `{"error":{"code":401,"message":"API key is required","status":"UNAUTHENTICATED"}}`}, reason: "missing_key", match: true},
|
||||
{name: "google group deleted", item: candidate{body: `{"error":{"code":403,"message":"API Key 所属分组已删除","status":"PERMISSION_DENIED"}}`}, reason: "group_deleted", match: true},
|
||||
{name: "ip acl", item: candidate{body: `{"code":"ACCESS_DENIED","message":"Access denied. Your IP is 192.0.2.1"}`}, reason: "ip_acl_denied", match: true},
|
||||
{name: "user not found remains", item: candidate{body: `{"code":"USER_NOT_FOUND","message":"User associated with API key not found"}`}, match: false},
|
||||
{name: "quota remains", item: candidate{body: `{"code":"API_KEY_QUOTA_EXHAUSTED","message":"quota"}`}, match: false},
|
||||
{name: "database failure remains", item: candidate{statusCode: 500, message: "Failed to validate API key", body: `{"code":"INTERNAL_ERROR","message":"Failed to validate API key"}`}, match: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
reason, ok := historicalIngressRejectReason(tt.item)
|
||||
if ok != tt.match || reason != tt.reason {
|
||||
t.Fatalf("got (%q, %v), want (%q, %v)", reason, ok, tt.reason, tt.match)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1 +1 @@
|
||||
0.1.156
|
||||
0.1.160
|
||||
|
||||
@@ -153,6 +153,14 @@ func runMainServer() {
|
||||
log.Fatalf("Failed to initialize application: %v", err)
|
||||
}
|
||||
defer app.Cleanup()
|
||||
if app.PromptAudit != nil {
|
||||
if err := app.PromptAudit.Start(context.Background()); err != nil {
|
||||
// Startup continues so unrelated APIs stay up, but Prompt Audit itself
|
||||
// fails closed (unavailable) until a later reload installs a trusted
|
||||
// snapshot—avoiding a silent ModeOff bypass of persisted blocking policy.
|
||||
log.Printf("Prompt Audit started in degraded fail-closed state: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 启动服务器
|
||||
go func() {
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler"
|
||||
"github.com/Wei-Shaw/sub2api/internal/payment"
|
||||
"github.com/Wei-Shaw/sub2api/internal/repository"
|
||||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
@@ -24,8 +25,9 @@ import (
|
||||
)
|
||||
|
||||
type Application struct {
|
||||
Server *http.Server
|
||||
Cleanup func()
|
||||
Server *http.Server
|
||||
PromptAudit *securityaudit.PromptService
|
||||
Cleanup func()
|
||||
}
|
||||
|
||||
func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
@@ -36,6 +38,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
// Business layer ProviderSets
|
||||
repository.ProviderSet,
|
||||
service.ProviderSet,
|
||||
securityaudit.ProviderSet,
|
||||
payment.ProviderSet,
|
||||
middleware.ProviderSet,
|
||||
handler.ProviderSet,
|
||||
@@ -53,7 +56,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
provideCleanup,
|
||||
|
||||
// Application struct
|
||||
wire.Struct(new(Application), "Server", "Cleanup"),
|
||||
wire.Struct(new(Application), "Server", "PromptAudit", "Cleanup"),
|
||||
)
|
||||
return nil, nil
|
||||
}
|
||||
@@ -78,6 +81,10 @@ func provideCleanup(
|
||||
opsCleanup *service.OpsCleanupService,
|
||||
opsScheduledReport *service.OpsScheduledReportService,
|
||||
opsSystemLogSink *service.OpsSystemLogSink,
|
||||
opsService *service.OpsService,
|
||||
opsIngressReject *service.OpsIngressRejectAggregator,
|
||||
apiKeyService *service.APIKeyService,
|
||||
authCacheInvalidationWorker *service.AuthCacheInvalidationWorker,
|
||||
schedulerSnapshot *service.SchedulerSnapshotService,
|
||||
tokenRefresh *service.TokenRefreshService,
|
||||
accountExpiry *service.AccountExpiryService,
|
||||
@@ -105,6 +112,7 @@ func provideCleanup(
|
||||
quotaFlusher *service.UserPlatformQuotaUsageFlusher,
|
||||
upstreamBillingProbe *service.UpstreamBillingProbeService,
|
||||
auditLog *service.AuditLogService,
|
||||
promptAudit *securityaudit.PromptService,
|
||||
) func() {
|
||||
return func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
@@ -117,6 +125,36 @@ func provideCleanup(
|
||||
|
||||
// 应用层清理步骤可并行执行,基础设施资源(Redis/Ent)最后按顺序关闭。
|
||||
parallelSteps := []cleanupStep{
|
||||
{"OpsIngressRejectAggregator", func() error {
|
||||
if opsIngressReject != nil {
|
||||
opsIngressReject.Stop()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"AuthCacheInvalidationWorker", func() error {
|
||||
if authCacheInvalidationWorker != nil {
|
||||
authCacheInvalidationWorker.Stop()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"AuthCacheInvalidationSubscriber", func() error {
|
||||
if apiKeyService != nil {
|
||||
apiKeyService.StopAuthCacheInvalidationSubscriber()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"OpsRuntimeSettingsRefresh", func() error {
|
||||
if opsService != nil {
|
||||
opsService.StopRuntimeSettingsRefresh()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"PromptAuditService", func() error {
|
||||
if promptAudit != nil {
|
||||
return promptAudit.Shutdown(ctx)
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"OpsScheduledReportService", func() error {
|
||||
if opsScheduledReport != nil {
|
||||
opsScheduledReport.Stop()
|
||||
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/admin"
|
||||
"github.com/Wei-Shaw/sub2api/internal/payment"
|
||||
"github.com/Wei-Shaw/sub2api/internal/repository"
|
||||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
@@ -156,7 +157,9 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
antigravityGatewayService := service.NewAntigravityGatewayService(accountRepository, gatewayCache, schedulerSnapshotService, antigravityTokenProvider, rateLimitService, httpUpstream, settingService, internal500CounterCache)
|
||||
geminiMessagesCompatService := service.NewGeminiMessagesCompatService(accountRepository, groupRepository, gatewayCache, schedulerSnapshotService, geminiTokenProvider, rateLimitService, httpUpstream, antigravityGatewayService, configConfig)
|
||||
opsSystemLogSink := service.ProvideOpsSystemLogSink(opsRepository)
|
||||
opsService := service.ProvideOpsService(opsRepository, settingRepository, configConfig, accountRepository, userRepository, concurrencyService, gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, opsSystemLogSink, settingService)
|
||||
authCacheInvalidationOutboxRepository := repository.NewAuthCacheInvalidationOutboxRepository(db)
|
||||
authCacheInvalidationWorker := service.ProvideAuthCacheInvalidationWorker(authCacheInvalidationOutboxRepository, apiKeyCache, apiKeyService)
|
||||
opsService := service.ProvideOpsService(opsRepository, settingRepository, configConfig, accountRepository, userRepository, concurrencyService, gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, opsSystemLogSink, settingService, authCacheInvalidationWorker, apiKeyService)
|
||||
usageHandler := handler.NewUsageHandler(usageService, apiKeyService, opsService, settingService)
|
||||
redeemHandler := handler.NewRedeemHandler(redeemService)
|
||||
subscriptionHandler := handler.NewSubscriptionHandler(subscriptionService)
|
||||
@@ -173,11 +176,12 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
leaderLockCache := repository.NewLeaderLockCache(redisClient)
|
||||
dashboardAggregationService := service.ProvideDashboardAggregationService(dashboardAggregationRepository, timingWheelService, leaderLockCache, db, configConfig)
|
||||
dashboardHandler := admin.NewDashboardHandler(dashboardService, dashboardAggregationService)
|
||||
adminGroupRepository := repository.NewAdminGroupRepository(client, db)
|
||||
adminAccountRepository := repository.NewAdminAccountRepository(client, db, schedulerCache)
|
||||
proxyExitInfoProber := repository.NewProxyExitInfoProber(configConfig)
|
||||
proxyLatencyCache := repository.NewProxyLatencyCache(redisClient)
|
||||
adminService := service.NewAdminService(userRepository, groupRepository, adminAccountRepository, proxyRepository, apiKeyRepository, redeemCodeRepository, userGroupRateRepository, userRPMCache, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator, client, settingService, subscriptionService, userSubscriptionRepository, privacyClientFactory, openAIGatewayService, affiliateService)
|
||||
adminUserHandler := admin.NewUserHandler(adminService, concurrencyService, serviceUserPlatformQuotaRepository, billingCache, totpService, userService)
|
||||
adminService := service.NewAdminService(userRepository, adminGroupRepository, adminAccountRepository, proxyRepository, apiKeyRepository, redeemCodeRepository, userGroupRateRepository, userRPMCache, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator, client, settingService, subscriptionService, userSubscriptionRepository, privacyClientFactory, openAIGatewayService, affiliateService)
|
||||
adminUserHandler := admin.NewUserHandler(adminService, concurrencyService, serviceUserPlatformQuotaRepository, billingCache, totpService, userService, settingService)
|
||||
groupCapacityService := service.NewGroupCapacityService(accountRepository, groupRepository, concurrencyService, sessionLimitCache, rpmCache)
|
||||
groupHandler := admin.NewGroupHandler(adminService, dashboardService, groupCapacityService)
|
||||
claudeUsageFetcher := repository.NewClaudeUsageFetcher(httpUpstream)
|
||||
@@ -214,7 +218,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
registry := payment.ProvideRegistry()
|
||||
defaultLoadBalancer := payment.ProvideDefaultLoadBalancer(client, encryptionKey)
|
||||
paymentService := service.ProvidePaymentService(client, registry, defaultLoadBalancer, redeemService, subscriptionService, paymentConfigService, userRepository, groupRepository, affiliateService, notificationEmailService)
|
||||
settingHandler := handler.ProvideAdminSettingHandler(settingService, emailService, turnstileService, opsService, paymentConfigService, paymentService, userAttributeService, notificationEmailService)
|
||||
settingHandler := handler.ProvideAdminSettingHandler(settingService, emailService, turnstileService, opsService, paymentConfigService, paymentService, userAttributeService, notificationEmailService, totpService, userService)
|
||||
opsHandler := admin.NewOpsHandler(opsService)
|
||||
updateCache := repository.NewUpdateCache(redisClient)
|
||||
gitHubReleaseClient := repository.ProvideGitHubReleaseClient(configConfig)
|
||||
@@ -247,6 +251,13 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
contentModerationHashCache := repository.NewContentModerationHashCache(redisClient)
|
||||
contentModerationService := service.NewContentModerationService(settingRepository, contentModerationRepository, contentModerationHashCache, groupRepository, userRepository, apiKeyAuthCacheInvalidator, emailService)
|
||||
contentModerationHandler := admin.NewContentModerationHandler(contentModerationService)
|
||||
configManager := securityaudit.NewConfigManager(db, settingRepository, redisClient, secretEncryptor)
|
||||
postgreSQLRepository := securityaudit.NewPostgreSQLRepository(db)
|
||||
redisPayloadStore := securityaudit.NewRedisPayloadStore(redisClient)
|
||||
openAICompatibleScanner := securityaudit.NewOpenAICompatibleScanner()
|
||||
atomicMetrics := securityaudit.NewAtomicMetrics()
|
||||
promptService := securityaudit.NewPromptService(configManager, postgreSQLRepository, redisPayloadStore, openAICompatibleScanner, atomicMetrics)
|
||||
promptAdminHandler := securityaudit.NewPromptAdminHandler(promptService)
|
||||
paymentHandler := admin.NewPaymentHandler(paymentService, paymentConfigService)
|
||||
affiliateHandler := admin.NewAffiliateHandler(affiliateService, adminService)
|
||||
complianceHandler := admin.NewComplianceHandler(settingService)
|
||||
@@ -254,12 +265,14 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
auditLogService := service.ProvideAuditLogService(auditLogRepository, settingService)
|
||||
auditLogHandler := admin.NewAuditLogHandler(auditLogService, totpService)
|
||||
upstreamBillingProbeService := service.ProvideUpstreamBillingProbeService(accountRepository, accountTestService, settingService, leaderLockCache, db)
|
||||
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService)
|
||||
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService)
|
||||
usageRecordWorkerPool := service.NewUsageRecordWorkerPool(configConfig)
|
||||
userMsgQueueCache := repository.NewUserMsgQueueCache(redisClient)
|
||||
userMessageQueueService := service.ProvideUserMessageQueueService(userMsgQueueCache, rpmCache, configConfig)
|
||||
gatewayHandler := handler.NewGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService)
|
||||
openAIGatewayHandler := handler.NewOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, configConfig)
|
||||
legacyEngine := securityaudit.NewLegacyModerationAdapter(contentModerationService)
|
||||
coordinator := securityaudit.NewCoordinator(legacyEngine, promptService)
|
||||
gatewayHandler := handler.ProvideGatewayHandler(gatewayService, openAIGatewayService, geminiMessagesCompatService, antigravityGatewayService, userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, userMessageQueueService, configConfig, settingService, coordinator)
|
||||
openAIGatewayHandler := handler.ProvideOpenAIGatewayHandler(openAIGatewayService, concurrencyService, billingCacheService, apiKeyService, usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, grokQuotaService, configConfig, coordinator)
|
||||
handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo, notificationEmailService)
|
||||
totpHandler := handler.NewTotpHandler(totpService)
|
||||
handlerPaymentHandler := handler.NewPaymentHandler(paymentService, paymentConfigService)
|
||||
@@ -279,7 +292,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
batchImageDownloadLimiter := repository.NewBatchImageDownloadLimiter(redisClient, configConfig)
|
||||
batchImageDownloadService := service.NewBatchImageDownloadService(batchImageRepository, accountRepository, batchImageDownloadLimiter, configConfig)
|
||||
batchImageCleanupService := service.ProvideBatchImageCleanupService(batchImageRepository, accountRepository, configConfig)
|
||||
batchImageHandler := handler.NewBatchImageHandler(batchImagePublicService, batchImageDownloadService, batchImageCleanupService)
|
||||
batchImageHandler := handler.ProvideBatchImageHandler(batchImagePublicService, batchImageDownloadService, batchImageCleanupService, openAIGatewayHandler)
|
||||
idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig)
|
||||
idempotencyCleanupService := service.ProvideIdempotencyCleanupService(idempotencyRepository, configConfig)
|
||||
handlers := handler.ProvideHandlers(authHandler, userHandler, apiKeyHandler, usageHandler, redeemHandler, subscriptionHandler, announcementHandler, channelMonitorUserHandler, adminHandlers, gatewayHandler, openAIGatewayHandler, handlerSettingHandler, totpHandler, handlerPaymentHandler, paymentWebhookHandler, availableChannelHandler, asyncImageHandler, batchImageHandler, idempotencyCoordinator, idempotencyCleanupService)
|
||||
@@ -287,7 +300,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
adminAuthMiddleware := middleware.NewAdminAuthMiddleware(authService, userService, settingService, auditLogService)
|
||||
apiKeyAuthMiddleware := middleware.NewAPIKeyAuthMiddleware(apiKeyService, subscriptionService, configConfig)
|
||||
auditLogMiddleware := middleware.NewAuditLogMiddleware(auditLogService)
|
||||
stepUpAuthMiddleware := middleware.NewStepUpAuthMiddleware(totpService, userService)
|
||||
stepUpAuthMiddleware := middleware.NewStepUpAuthMiddleware(totpService, userService, settingService)
|
||||
engine := server.ProvideRouter(configConfig, handlers, jwtAuthMiddleware, adminAuthMiddleware, apiKeyAuthMiddleware, auditLogMiddleware, stepUpAuthMiddleware, apiKeyService, subscriptionService, opsService, settingService, redisClient)
|
||||
httpServer := server.ProvideHTTPServer(configConfig, engine)
|
||||
opsMetricsCollector := service.ProvideOpsMetricsCollector(opsRepository, settingRepository, accountRepository, concurrencyService, db, redisClient, configConfig)
|
||||
@@ -295,6 +308,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
opsAlertEvaluatorService := service.ProvideOpsAlertEvaluatorService(opsService, opsRepository, emailService, redisClient, configConfig, proxyRepository)
|
||||
opsCleanupService := service.ProvideOpsCleanupService(opsRepository, db, redisClient, configConfig, channelMonitorService, settingRepository, opsService)
|
||||
opsScheduledReportService := service.ProvideOpsScheduledReportService(opsService, userService, emailService, redisClient, configConfig)
|
||||
opsIngressRejectAggregator := service.ProvideOpsIngressRejectAggregator(opsRepository, opsService)
|
||||
accountExpiryService := service.ProvideAccountExpiryService(accountRepository)
|
||||
proxyExpiryService := service.ProvideProxyExpiryService(proxyRepository)
|
||||
subscriptionExpiryService := service.ProvideSubscriptionExpiryService(userSubscriptionRepository, settingRepository, notificationEmailService, leaderLockCache, db)
|
||||
@@ -303,10 +317,11 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
paymentOrderExpiryService := service.ProvidePaymentOrderExpiryService(paymentService, leaderLockCache, db)
|
||||
channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService)
|
||||
userPlatformQuotaUsageFlusher := service.ProvideUserPlatformQuotaUsageFlusher(configConfig, billingCache, serviceUserPlatformQuotaRepository, timingWheelService)
|
||||
v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, schedulerSnapshotService, tokenRefreshService, accountExpiryService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, auditLogService)
|
||||
v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, auditLogService, promptService)
|
||||
application := &Application{
|
||||
Server: httpServer,
|
||||
Cleanup: v,
|
||||
Server: httpServer,
|
||||
PromptAudit: promptService,
|
||||
Cleanup: v,
|
||||
}
|
||||
return application, nil
|
||||
}
|
||||
@@ -314,8 +329,9 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
// wire.go:
|
||||
|
||||
type Application struct {
|
||||
Server *http.Server
|
||||
Cleanup func()
|
||||
Server *http.Server
|
||||
PromptAudit *securityaudit.PromptService
|
||||
Cleanup func()
|
||||
}
|
||||
|
||||
func providePrivacyClientFactory() service.PrivacyClientFactory {
|
||||
@@ -338,6 +354,10 @@ func provideCleanup(
|
||||
opsCleanup *service.OpsCleanupService,
|
||||
opsScheduledReport *service.OpsScheduledReportService,
|
||||
opsSystemLogSink *service.OpsSystemLogSink,
|
||||
opsService *service.OpsService,
|
||||
opsIngressReject *service.OpsIngressRejectAggregator,
|
||||
apiKeyService *service.APIKeyService,
|
||||
authCacheInvalidationWorker *service.AuthCacheInvalidationWorker,
|
||||
schedulerSnapshot *service.SchedulerSnapshotService,
|
||||
tokenRefresh *service.TokenRefreshService,
|
||||
accountExpiry *service.AccountExpiryService,
|
||||
@@ -365,6 +385,7 @@ func provideCleanup(
|
||||
quotaFlusher *service.UserPlatformQuotaUsageFlusher,
|
||||
upstreamBillingProbe *service.UpstreamBillingProbeService,
|
||||
auditLog *service.AuditLogService,
|
||||
promptAudit *securityaudit.PromptService,
|
||||
) func() {
|
||||
return func() {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
@@ -376,6 +397,36 @@ func provideCleanup(
|
||||
}
|
||||
|
||||
parallelSteps := []cleanupStep{
|
||||
{"OpsIngressRejectAggregator", func() error {
|
||||
if opsIngressReject != nil {
|
||||
opsIngressReject.Stop()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"AuthCacheInvalidationWorker", func() error {
|
||||
if authCacheInvalidationWorker != nil {
|
||||
authCacheInvalidationWorker.Stop()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"AuthCacheInvalidationSubscriber", func() error {
|
||||
if apiKeyService != nil {
|
||||
apiKeyService.StopAuthCacheInvalidationSubscriber()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"OpsRuntimeSettingsRefresh", func() error {
|
||||
if opsService != nil {
|
||||
opsService.StopRuntimeSettingsRefresh()
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"PromptAuditService", func() error {
|
||||
if promptAudit != nil {
|
||||
return promptAudit.Shutdown(ctx)
|
||||
}
|
||||
return nil
|
||||
}},
|
||||
{"OpsScheduledReportService", func() error {
|
||||
if opsScheduledReport != nil {
|
||||
opsScheduledReport.Stop()
|
||||
|
||||
@@ -58,6 +58,10 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) {
|
||||
&service.OpsCleanupService{},
|
||||
&service.OpsScheduledReportService{},
|
||||
opsSystemLogSinkSvc,
|
||||
nil, // opsService
|
||||
nil, // opsIngressRejectAggregator
|
||||
nil, // apiKeyService
|
||||
nil, // authCacheInvalidationWorker
|
||||
schedulerSnapshotSvc,
|
||||
tokenRefreshSvc,
|
||||
accountExpirySvc,
|
||||
@@ -85,6 +89,7 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) {
|
||||
nil, // quotaFlusher
|
||||
nil, // upstreamBillingProbe
|
||||
nil, // auditLog
|
||||
nil, // promptAudit
|
||||
)
|
||||
|
||||
require.NotPanics(t, func() {
|
||||
|
||||
+15
-1
@@ -43,6 +43,8 @@ type Group struct {
|
||||
IsExclusive bool `json:"is_exclusive,omitempty"`
|
||||
// Status holds the value of the "status" field.
|
||||
Status string `json:"status,omitempty"`
|
||||
// 内部幂等恢复标识,不对 API 暴露
|
||||
DuplicateOperationID *string `json:"duplicate_operation_id,omitempty"`
|
||||
// Platform holds the value of the "platform" field.
|
||||
Platform string `json:"platform,omitempty"`
|
||||
// SubscriptionType holds the value of the "subscription_type" field.
|
||||
@@ -229,7 +231,7 @@ func (*Group) scanValues(columns []string) ([]any, error) {
|
||||
values[i] = new(sql.NullFloat64)
|
||||
case group.FieldID, group.FieldDefaultValidityDays, group.FieldFallbackGroupID, group.FieldFallbackGroupIDOnInvalidRequest, group.FieldSortOrder, group.FieldRpmLimit:
|
||||
values[i] = new(sql.NullInt64)
|
||||
case group.FieldName, group.FieldDescription, group.FieldPeakStart, group.FieldPeakEnd, group.FieldStatus, group.FieldPlatform, group.FieldSubscriptionType, group.FieldDefaultMappedModel:
|
||||
case group.FieldName, group.FieldDescription, group.FieldPeakStart, group.FieldPeakEnd, group.FieldStatus, group.FieldDuplicateOperationID, group.FieldPlatform, group.FieldSubscriptionType, group.FieldDefaultMappedModel:
|
||||
values[i] = new(sql.NullString)
|
||||
case group.FieldCreatedAt, group.FieldUpdatedAt, group.FieldDeletedAt:
|
||||
values[i] = new(sql.NullTime)
|
||||
@@ -328,6 +330,13 @@ func (_m *Group) assignValues(columns []string, values []any) error {
|
||||
} else if value.Valid {
|
||||
_m.Status = value.String
|
||||
}
|
||||
case group.FieldDuplicateOperationID:
|
||||
if value, ok := values[i].(*sql.NullString); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field duplicate_operation_id", values[i])
|
||||
} else if value.Valid {
|
||||
_m.DuplicateOperationID = new(string)
|
||||
*_m.DuplicateOperationID = value.String
|
||||
}
|
||||
case group.FieldPlatform:
|
||||
if value, ok := values[i].(*sql.NullString); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field platform", values[i])
|
||||
@@ -680,6 +689,11 @@ func (_m *Group) String() string {
|
||||
builder.WriteString("status=")
|
||||
builder.WriteString(_m.Status)
|
||||
builder.WriteString(", ")
|
||||
if v := _m.DuplicateOperationID; v != nil {
|
||||
builder.WriteString("duplicate_operation_id=")
|
||||
builder.WriteString(*v)
|
||||
}
|
||||
builder.WriteString(", ")
|
||||
builder.WriteString("platform=")
|
||||
builder.WriteString(_m.Platform)
|
||||
builder.WriteString(", ")
|
||||
|
||||
@@ -40,6 +40,8 @@ const (
|
||||
FieldIsExclusive = "is_exclusive"
|
||||
// FieldStatus holds the string denoting the status field in the database.
|
||||
FieldStatus = "status"
|
||||
// FieldDuplicateOperationID holds the string denoting the duplicate_operation_id field in the database.
|
||||
FieldDuplicateOperationID = "duplicate_operation_id"
|
||||
// FieldPlatform holds the string denoting the platform field in the database.
|
||||
FieldPlatform = "platform"
|
||||
// FieldSubscriptionType holds the string denoting the subscription_type field in the database.
|
||||
@@ -199,6 +201,7 @@ var Columns = []string{
|
||||
FieldPeakRateMultiplier,
|
||||
FieldIsExclusive,
|
||||
FieldStatus,
|
||||
FieldDuplicateOperationID,
|
||||
FieldPlatform,
|
||||
FieldSubscriptionType,
|
||||
FieldDailyLimitUsd,
|
||||
@@ -292,6 +295,8 @@ var (
|
||||
DefaultStatus string
|
||||
// StatusValidator is a validator for the "status" field. It is called by the builders before save.
|
||||
StatusValidator func(string) error
|
||||
// DuplicateOperationIDValidator is a validator for the "duplicate_operation_id" field. It is called by the builders before save.
|
||||
DuplicateOperationIDValidator func(string) error
|
||||
// DefaultPlatform holds the default value on creation for the "platform" field.
|
||||
DefaultPlatform string
|
||||
// PlatformValidator is a validator for the "platform" field. It is called by the builders before save.
|
||||
@@ -414,6 +419,11 @@ func ByStatus(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldStatus, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByDuplicateOperationID orders the results by the duplicate_operation_id field.
|
||||
func ByDuplicateOperationID(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldDuplicateOperationID, opts...).ToFunc()
|
||||
}
|
||||
|
||||
// ByPlatform orders the results by the platform field.
|
||||
func ByPlatform(opts ...sql.OrderTermOption) OrderOption {
|
||||
return sql.OrderByField(FieldPlatform, opts...).ToFunc()
|
||||
|
||||
@@ -115,6 +115,11 @@ func Status(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldStatus, v))
|
||||
}
|
||||
|
||||
// DuplicateOperationID applies equality check predicate on the "duplicate_operation_id" field. It's identical to DuplicateOperationIDEQ.
|
||||
func DuplicateOperationID(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldDuplicateOperationID, v))
|
||||
}
|
||||
|
||||
// Platform applies equality check predicate on the "platform" field. It's identical to PlatformEQ.
|
||||
func Platform(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldPlatform, v))
|
||||
@@ -840,6 +845,81 @@ func StatusContainsFold(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldContainsFold(FieldStatus, v))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDEQ applies the EQ predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDEQ(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldDuplicateOperationID, v))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDNEQ applies the NEQ predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDNEQ(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldNEQ(FieldDuplicateOperationID, v))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDIn applies the In predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDIn(vs ...string) predicate.Group {
|
||||
return predicate.Group(sql.FieldIn(FieldDuplicateOperationID, vs...))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDNotIn applies the NotIn predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDNotIn(vs ...string) predicate.Group {
|
||||
return predicate.Group(sql.FieldNotIn(FieldDuplicateOperationID, vs...))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDGT applies the GT predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDGT(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldGT(FieldDuplicateOperationID, v))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDGTE applies the GTE predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDGTE(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldGTE(FieldDuplicateOperationID, v))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDLT applies the LT predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDLT(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldLT(FieldDuplicateOperationID, v))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDLTE applies the LTE predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDLTE(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldLTE(FieldDuplicateOperationID, v))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDContains applies the Contains predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDContains(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldContains(FieldDuplicateOperationID, v))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDHasPrefix applies the HasPrefix predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDHasPrefix(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldHasPrefix(FieldDuplicateOperationID, v))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDHasSuffix applies the HasSuffix predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDHasSuffix(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldHasSuffix(FieldDuplicateOperationID, v))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDIsNil applies the IsNil predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDIsNil() predicate.Group {
|
||||
return predicate.Group(sql.FieldIsNull(FieldDuplicateOperationID))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDNotNil applies the NotNil predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDNotNil() predicate.Group {
|
||||
return predicate.Group(sql.FieldNotNull(FieldDuplicateOperationID))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDEqualFold applies the EqualFold predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDEqualFold(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldEqualFold(FieldDuplicateOperationID, v))
|
||||
}
|
||||
|
||||
// DuplicateOperationIDContainsFold applies the ContainsFold predicate on the "duplicate_operation_id" field.
|
||||
func DuplicateOperationIDContainsFold(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldContainsFold(FieldDuplicateOperationID, v))
|
||||
}
|
||||
|
||||
// PlatformEQ applies the EQ predicate on the "platform" field.
|
||||
func PlatformEQ(v string) predicate.Group {
|
||||
return predicate.Group(sql.FieldEQ(FieldPlatform, v))
|
||||
|
||||
@@ -189,6 +189,20 @@ func (_c *GroupCreate) SetNillableStatus(v *string) *GroupCreate {
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetDuplicateOperationID sets the "duplicate_operation_id" field.
|
||||
func (_c *GroupCreate) SetDuplicateOperationID(v string) *GroupCreate {
|
||||
_c.mutation.SetDuplicateOperationID(v)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetNillableDuplicateOperationID sets the "duplicate_operation_id" field if the given value is not nil.
|
||||
func (_c *GroupCreate) SetNillableDuplicateOperationID(v *string) *GroupCreate {
|
||||
if v != nil {
|
||||
_c.SetDuplicateOperationID(*v)
|
||||
}
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetPlatform sets the "platform" field.
|
||||
func (_c *GroupCreate) SetPlatform(v string) *GroupCreate {
|
||||
_c.mutation.SetPlatform(v)
|
||||
@@ -993,6 +1007,11 @@ func (_c *GroupCreate) check() error {
|
||||
return &ValidationError{Name: "status", err: fmt.Errorf(`ent: validator failed for field "Group.status": %w`, err)}
|
||||
}
|
||||
}
|
||||
if v, ok := _c.mutation.DuplicateOperationID(); ok {
|
||||
if err := group.DuplicateOperationIDValidator(v); err != nil {
|
||||
return &ValidationError{Name: "duplicate_operation_id", err: fmt.Errorf(`ent: validator failed for field "Group.duplicate_operation_id": %w`, err)}
|
||||
}
|
||||
}
|
||||
if _, ok := _c.mutation.Platform(); !ok {
|
||||
return &ValidationError{Name: "platform", err: errors.New(`ent: missing required field "Group.platform"`)}
|
||||
}
|
||||
@@ -1152,6 +1171,10 @@ func (_c *GroupCreate) createSpec() (*Group, *sqlgraph.CreateSpec) {
|
||||
_spec.SetField(group.FieldStatus, field.TypeString, value)
|
||||
_node.Status = value
|
||||
}
|
||||
if value, ok := _c.mutation.DuplicateOperationID(); ok {
|
||||
_spec.SetField(group.FieldDuplicateOperationID, field.TypeString, value)
|
||||
_node.DuplicateOperationID = &value
|
||||
}
|
||||
if value, ok := _c.mutation.Platform(); ok {
|
||||
_spec.SetField(group.FieldPlatform, field.TypeString, value)
|
||||
_node.Platform = value
|
||||
@@ -2246,6 +2269,9 @@ func (u *GroupUpsertOne) UpdateNewValues() *GroupUpsertOne {
|
||||
if _, exists := u.create.mutation.CreatedAt(); exists {
|
||||
s.SetIgnore(group.FieldCreatedAt)
|
||||
}
|
||||
if _, exists := u.create.mutation.DuplicateOperationID(); exists {
|
||||
s.SetIgnore(group.FieldDuplicateOperationID)
|
||||
}
|
||||
}))
|
||||
return u
|
||||
}
|
||||
@@ -3366,6 +3392,9 @@ func (u *GroupUpsertBulk) UpdateNewValues() *GroupUpsertBulk {
|
||||
if _, exists := b.mutation.CreatedAt(); exists {
|
||||
s.SetIgnore(group.FieldCreatedAt)
|
||||
}
|
||||
if _, exists := b.mutation.DuplicateOperationID(); exists {
|
||||
s.SetIgnore(group.FieldDuplicateOperationID)
|
||||
}
|
||||
}
|
||||
}))
|
||||
return u
|
||||
|
||||
@@ -1273,6 +1273,9 @@ func (_u *GroupUpdate) sqlSave(ctx context.Context) (_node int, err error) {
|
||||
if value, ok := _u.mutation.Status(); ok {
|
||||
_spec.SetField(group.FieldStatus, field.TypeString, value)
|
||||
}
|
||||
if _u.mutation.DuplicateOperationIDCleared() {
|
||||
_spec.ClearField(group.FieldDuplicateOperationID, field.TypeString)
|
||||
}
|
||||
if value, ok := _u.mutation.Platform(); ok {
|
||||
_spec.SetField(group.FieldPlatform, field.TypeString, value)
|
||||
}
|
||||
@@ -3063,6 +3066,9 @@ func (_u *GroupUpdateOne) sqlSave(ctx context.Context) (_node *Group, err error)
|
||||
if value, ok := _u.mutation.Status(); ok {
|
||||
_spec.SetField(group.FieldStatus, field.TypeString, value)
|
||||
}
|
||||
if _u.mutation.DuplicateOperationIDCleared() {
|
||||
_spec.ClearField(group.FieldDuplicateOperationID, field.TypeString)
|
||||
}
|
||||
if value, ok := _u.mutation.Platform(); ok {
|
||||
_spec.SetField(group.FieldPlatform, field.TypeString, value)
|
||||
}
|
||||
|
||||
@@ -845,6 +845,7 @@ var (
|
||||
{Name: "peak_rate_multiplier", Type: field.TypeFloat64, Default: 1, SchemaType: map[string]string{"postgres": "decimal(10,4)"}},
|
||||
{Name: "is_exclusive", Type: field.TypeBool, Default: false},
|
||||
{Name: "status", Type: field.TypeString, Size: 20, Default: "active"},
|
||||
{Name: "duplicate_operation_id", Type: field.TypeString, Nullable: true, Size: 64},
|
||||
{Name: "platform", Type: field.TypeString, Size: 50, Default: "anthropic"},
|
||||
{Name: "subscription_type", Type: field.TypeString, Size: 20, Default: "standard"},
|
||||
{Name: "daily_limit_usd", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}},
|
||||
@@ -896,12 +897,12 @@ var (
|
||||
{
|
||||
Name: "group_platform",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{GroupsColumns[13]},
|
||||
Columns: []*schema.Column{GroupsColumns[14]},
|
||||
},
|
||||
{
|
||||
Name: "group_subscription_type",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{GroupsColumns[14]},
|
||||
Columns: []*schema.Column{GroupsColumns[15]},
|
||||
},
|
||||
{
|
||||
Name: "group_is_exclusive",
|
||||
@@ -916,7 +917,15 @@ var (
|
||||
{
|
||||
Name: "group_sort_order",
|
||||
Unique: false,
|
||||
Columns: []*schema.Column{GroupsColumns[41]},
|
||||
Columns: []*schema.Column{GroupsColumns[42]},
|
||||
},
|
||||
{
|
||||
Name: "idx_groups_duplicate_operation_id_active",
|
||||
Unique: true,
|
||||
Columns: []*schema.Column{GroupsColumns[13]},
|
||||
Annotation: &entsql.IndexAnnotation{
|
||||
Where: "duplicate_operation_id IS NOT NULL AND deleted_at IS NULL",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
+74
-1
@@ -20808,6 +20808,7 @@ type GroupMutation struct {
|
||||
addpeak_rate_multiplier *float64
|
||||
is_exclusive *bool
|
||||
status *string
|
||||
duplicate_operation_id *string
|
||||
platform *string
|
||||
subscription_type *string
|
||||
daily_limit_usd *float64
|
||||
@@ -21484,6 +21485,55 @@ func (m *GroupMutation) ResetStatus() {
|
||||
m.status = nil
|
||||
}
|
||||
|
||||
// SetDuplicateOperationID sets the "duplicate_operation_id" field.
|
||||
func (m *GroupMutation) SetDuplicateOperationID(s string) {
|
||||
m.duplicate_operation_id = &s
|
||||
}
|
||||
|
||||
// DuplicateOperationID returns the value of the "duplicate_operation_id" field in the mutation.
|
||||
func (m *GroupMutation) DuplicateOperationID() (r string, exists bool) {
|
||||
v := m.duplicate_operation_id
|
||||
if v == nil {
|
||||
return
|
||||
}
|
||||
return *v, true
|
||||
}
|
||||
|
||||
// OldDuplicateOperationID returns the old "duplicate_operation_id" 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) OldDuplicateOperationID(ctx context.Context) (v *string, err error) {
|
||||
if !m.op.Is(OpUpdateOne) {
|
||||
return v, errors.New("OldDuplicateOperationID is only allowed on UpdateOne operations")
|
||||
}
|
||||
if m.id == nil || m.oldValue == nil {
|
||||
return v, errors.New("OldDuplicateOperationID requires an ID field in the mutation")
|
||||
}
|
||||
oldValue, err := m.oldValue(ctx)
|
||||
if err != nil {
|
||||
return v, fmt.Errorf("querying old value for OldDuplicateOperationID: %w", err)
|
||||
}
|
||||
return oldValue.DuplicateOperationID, nil
|
||||
}
|
||||
|
||||
// ClearDuplicateOperationID clears the value of the "duplicate_operation_id" field.
|
||||
func (m *GroupMutation) ClearDuplicateOperationID() {
|
||||
m.duplicate_operation_id = nil
|
||||
m.clearedFields[group.FieldDuplicateOperationID] = struct{}{}
|
||||
}
|
||||
|
||||
// DuplicateOperationIDCleared returns if the "duplicate_operation_id" field was cleared in this mutation.
|
||||
func (m *GroupMutation) DuplicateOperationIDCleared() bool {
|
||||
_, ok := m.clearedFields[group.FieldDuplicateOperationID]
|
||||
return ok
|
||||
}
|
||||
|
||||
// ResetDuplicateOperationID resets all changes to the "duplicate_operation_id" field.
|
||||
func (m *GroupMutation) ResetDuplicateOperationID() {
|
||||
m.duplicate_operation_id = nil
|
||||
delete(m.clearedFields, group.FieldDuplicateOperationID)
|
||||
}
|
||||
|
||||
// SetPlatform sets the "platform" field.
|
||||
func (m *GroupMutation) SetPlatform(s string) {
|
||||
m.platform = &s
|
||||
@@ -23714,7 +23764,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, 48)
|
||||
fields := make([]string, 0, 49)
|
||||
if m.created_at != nil {
|
||||
fields = append(fields, group.FieldCreatedAt)
|
||||
}
|
||||
@@ -23751,6 +23801,9 @@ func (m *GroupMutation) Fields() []string {
|
||||
if m.status != nil {
|
||||
fields = append(fields, group.FieldStatus)
|
||||
}
|
||||
if m.duplicate_operation_id != nil {
|
||||
fields = append(fields, group.FieldDuplicateOperationID)
|
||||
}
|
||||
if m.platform != nil {
|
||||
fields = append(fields, group.FieldPlatform)
|
||||
}
|
||||
@@ -23891,6 +23944,8 @@ func (m *GroupMutation) Field(name string) (ent.Value, bool) {
|
||||
return m.IsExclusive()
|
||||
case group.FieldStatus:
|
||||
return m.Status()
|
||||
case group.FieldDuplicateOperationID:
|
||||
return m.DuplicateOperationID()
|
||||
case group.FieldPlatform:
|
||||
return m.Platform()
|
||||
case group.FieldSubscriptionType:
|
||||
@@ -23996,6 +24051,8 @@ func (m *GroupMutation) OldField(ctx context.Context, name string) (ent.Value, e
|
||||
return m.OldIsExclusive(ctx)
|
||||
case group.FieldStatus:
|
||||
return m.OldStatus(ctx)
|
||||
case group.FieldDuplicateOperationID:
|
||||
return m.OldDuplicateOperationID(ctx)
|
||||
case group.FieldPlatform:
|
||||
return m.OldPlatform(ctx)
|
||||
case group.FieldSubscriptionType:
|
||||
@@ -24161,6 +24218,13 @@ func (m *GroupMutation) SetField(name string, value ent.Value) error {
|
||||
}
|
||||
m.SetStatus(v)
|
||||
return nil
|
||||
case group.FieldDuplicateOperationID:
|
||||
v, ok := value.(string)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T for field %s", value, name)
|
||||
}
|
||||
m.SetDuplicateOperationID(v)
|
||||
return nil
|
||||
case group.FieldPlatform:
|
||||
v, ok := value.(string)
|
||||
if !ok {
|
||||
@@ -24704,6 +24768,9 @@ func (m *GroupMutation) ClearedFields() []string {
|
||||
if m.FieldCleared(group.FieldDescription) {
|
||||
fields = append(fields, group.FieldDescription)
|
||||
}
|
||||
if m.FieldCleared(group.FieldDuplicateOperationID) {
|
||||
fields = append(fields, group.FieldDuplicateOperationID)
|
||||
}
|
||||
if m.FieldCleared(group.FieldDailyLimitUsd) {
|
||||
fields = append(fields, group.FieldDailyLimitUsd)
|
||||
}
|
||||
@@ -24763,6 +24830,9 @@ func (m *GroupMutation) ClearField(name string) error {
|
||||
case group.FieldDescription:
|
||||
m.ClearDescription()
|
||||
return nil
|
||||
case group.FieldDuplicateOperationID:
|
||||
m.ClearDuplicateOperationID()
|
||||
return nil
|
||||
case group.FieldDailyLimitUsd:
|
||||
m.ClearDailyLimitUsd()
|
||||
return nil
|
||||
@@ -24846,6 +24916,9 @@ func (m *GroupMutation) ResetField(name string) error {
|
||||
case group.FieldStatus:
|
||||
m.ResetStatus()
|
||||
return nil
|
||||
case group.FieldDuplicateOperationID:
|
||||
m.ResetDuplicateOperationID()
|
||||
return nil
|
||||
case group.FieldPlatform:
|
||||
m.ResetPlatform()
|
||||
return nil
|
||||
|
||||
@@ -995,102 +995,106 @@ func init() {
|
||||
group.DefaultStatus = groupDescStatus.Default.(string)
|
||||
// group.StatusValidator is a validator for the "status" field. It is called by the builders before save.
|
||||
group.StatusValidator = groupDescStatus.Validators[0].(func(string) error)
|
||||
// groupDescDuplicateOperationID is the schema descriptor for duplicate_operation_id field.
|
||||
groupDescDuplicateOperationID := groupFields[9].Descriptor()
|
||||
// group.DuplicateOperationIDValidator is a validator for the "duplicate_operation_id" field. It is called by the builders before save.
|
||||
group.DuplicateOperationIDValidator = groupDescDuplicateOperationID.Validators[0].(func(string) error)
|
||||
// groupDescPlatform is the schema descriptor for platform field.
|
||||
groupDescPlatform := groupFields[9].Descriptor()
|
||||
groupDescPlatform := groupFields[10].Descriptor()
|
||||
// group.DefaultPlatform holds the default value on creation for the platform field.
|
||||
group.DefaultPlatform = groupDescPlatform.Default.(string)
|
||||
// group.PlatformValidator is a validator for the "platform" field. It is called by the builders before save.
|
||||
group.PlatformValidator = groupDescPlatform.Validators[0].(func(string) error)
|
||||
// groupDescSubscriptionType is the schema descriptor for subscription_type field.
|
||||
groupDescSubscriptionType := groupFields[10].Descriptor()
|
||||
groupDescSubscriptionType := groupFields[11].Descriptor()
|
||||
// group.DefaultSubscriptionType holds the default value on creation for the subscription_type field.
|
||||
group.DefaultSubscriptionType = groupDescSubscriptionType.Default.(string)
|
||||
// group.SubscriptionTypeValidator is a validator for the "subscription_type" field. It is called by the builders before save.
|
||||
group.SubscriptionTypeValidator = groupDescSubscriptionType.Validators[0].(func(string) error)
|
||||
// groupDescDefaultValidityDays is the schema descriptor for default_validity_days field.
|
||||
groupDescDefaultValidityDays := groupFields[14].Descriptor()
|
||||
groupDescDefaultValidityDays := groupFields[15].Descriptor()
|
||||
// group.DefaultDefaultValidityDays holds the default value on creation for the default_validity_days field.
|
||||
group.DefaultDefaultValidityDays = groupDescDefaultValidityDays.Default.(int)
|
||||
// groupDescAllowImageGeneration is the schema descriptor for allow_image_generation field.
|
||||
groupDescAllowImageGeneration := groupFields[15].Descriptor()
|
||||
groupDescAllowImageGeneration := groupFields[16].Descriptor()
|
||||
// group.DefaultAllowImageGeneration holds the default value on creation for the allow_image_generation field.
|
||||
group.DefaultAllowImageGeneration = groupDescAllowImageGeneration.Default.(bool)
|
||||
// groupDescAllowBatchImageGeneration is the schema descriptor for allow_batch_image_generation field.
|
||||
groupDescAllowBatchImageGeneration := groupFields[16].Descriptor()
|
||||
groupDescAllowBatchImageGeneration := groupFields[17].Descriptor()
|
||||
// group.DefaultAllowBatchImageGeneration holds the default value on creation for the allow_batch_image_generation field.
|
||||
group.DefaultAllowBatchImageGeneration = groupDescAllowBatchImageGeneration.Default.(bool)
|
||||
// groupDescImageRateIndependent is the schema descriptor for image_rate_independent field.
|
||||
groupDescImageRateIndependent := groupFields[17].Descriptor()
|
||||
groupDescImageRateIndependent := groupFields[18].Descriptor()
|
||||
// group.DefaultImageRateIndependent holds the default value on creation for the image_rate_independent field.
|
||||
group.DefaultImageRateIndependent = groupDescImageRateIndependent.Default.(bool)
|
||||
// groupDescImageRateMultiplier is the schema descriptor for image_rate_multiplier field.
|
||||
groupDescImageRateMultiplier := groupFields[18].Descriptor()
|
||||
groupDescImageRateMultiplier := groupFields[19].Descriptor()
|
||||
// group.DefaultImageRateMultiplier holds the default value on creation for the image_rate_multiplier field.
|
||||
group.DefaultImageRateMultiplier = groupDescImageRateMultiplier.Default.(float64)
|
||||
// groupDescBatchImageDiscountMultiplier is the schema descriptor for batch_image_discount_multiplier field.
|
||||
groupDescBatchImageDiscountMultiplier := groupFields[22].Descriptor()
|
||||
groupDescBatchImageDiscountMultiplier := groupFields[23].Descriptor()
|
||||
// group.DefaultBatchImageDiscountMultiplier holds the default value on creation for the batch_image_discount_multiplier field.
|
||||
group.DefaultBatchImageDiscountMultiplier = groupDescBatchImageDiscountMultiplier.Default.(float64)
|
||||
// groupDescBatchImageHoldMultiplier is the schema descriptor for batch_image_hold_multiplier field.
|
||||
groupDescBatchImageHoldMultiplier := groupFields[23].Descriptor()
|
||||
groupDescBatchImageHoldMultiplier := groupFields[24].Descriptor()
|
||||
// group.DefaultBatchImageHoldMultiplier holds the default value on creation for the batch_image_hold_multiplier field.
|
||||
group.DefaultBatchImageHoldMultiplier = groupDescBatchImageHoldMultiplier.Default.(float64)
|
||||
// groupDescVideoRateIndependent is the schema descriptor for video_rate_independent field.
|
||||
groupDescVideoRateIndependent := groupFields[24].Descriptor()
|
||||
groupDescVideoRateIndependent := groupFields[25].Descriptor()
|
||||
// group.DefaultVideoRateIndependent holds the default value on creation for the video_rate_independent field.
|
||||
group.DefaultVideoRateIndependent = groupDescVideoRateIndependent.Default.(bool)
|
||||
// groupDescVideoRateMultiplier is the schema descriptor for video_rate_multiplier field.
|
||||
groupDescVideoRateMultiplier := groupFields[25].Descriptor()
|
||||
groupDescVideoRateMultiplier := groupFields[26].Descriptor()
|
||||
// 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[30].Descriptor()
|
||||
groupDescClaudeCodeOnly := groupFields[31].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[34].Descriptor()
|
||||
groupDescModelRoutingEnabled := groupFields[35].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[35].Descriptor()
|
||||
groupDescMcpXMLInject := groupFields[36].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[36].Descriptor()
|
||||
groupDescSupportedModelScopes := groupFields[37].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[37].Descriptor()
|
||||
groupDescSortOrder := groupFields[38].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[38].Descriptor()
|
||||
groupDescAllowMessagesDispatch := groupFields[39].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[39].Descriptor()
|
||||
groupDescRequireOauthOnly := groupFields[40].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[40].Descriptor()
|
||||
groupDescRequirePrivacySet := groupFields[41].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[41].Descriptor()
|
||||
groupDescDefaultMappedModel := groupFields[42].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[42].Descriptor()
|
||||
groupDescMessagesDispatchModelConfig := groupFields[43].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[43].Descriptor()
|
||||
groupDescModelsListConfig := groupFields[44].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[44].Descriptor()
|
||||
groupDescRpmLimit := groupFields[45].Descriptor()
|
||||
// group.DefaultRpmLimit holds the default value on creation for the rpm_limit field.
|
||||
group.DefaultRpmLimit = groupDescRpmLimit.Default.(int)
|
||||
idempotencyrecordMixin := schema.IdempotencyRecord{}.Mixin()
|
||||
|
||||
@@ -66,6 +66,12 @@ func (Group) Fields() []ent.Field {
|
||||
field.String("status").
|
||||
MaxLen(20).
|
||||
Default(domain.StatusActive),
|
||||
field.String("duplicate_operation_id").
|
||||
MaxLen(64).
|
||||
Optional().
|
||||
Nillable().
|
||||
Immutable().
|
||||
Comment("内部幂等恢复标识,不对 API 暴露"),
|
||||
|
||||
// Subscription-related fields (added by migration 003)
|
||||
field.String("platform").
|
||||
@@ -244,5 +250,9 @@ func (Group) Indexes() []ent.Index {
|
||||
index.Fields("is_exclusive"),
|
||||
index.Fields("deleted_at"),
|
||||
index.Fields("sort_order"),
|
||||
index.Fields("duplicate_operation_id").
|
||||
Unique().
|
||||
StorageKey("idx_groups_duplicate_operation_id_active").
|
||||
Annotations(entsql.IndexWhere("duplicate_operation_id IS NOT NULL AND deleted_at IS NULL")),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -645,6 +645,7 @@ type ServerConfig struct {
|
||||
EnableServerTiming bool `mapstructure:"enable_server_timing"` // Admin UI Server-Timing response header
|
||||
FrontendURL string `mapstructure:"frontend_url"` // 前端基础 URL,用于生成邮件中的外部链接
|
||||
ReadHeaderTimeout int `mapstructure:"read_header_timeout"` // 读取请求头超时(秒)
|
||||
MaxHeaderBytes int `mapstructure:"max_header_bytes"` // 请求头最大字节数(HTTP/2 映射为 header-list 上限)
|
||||
IdleTimeout int `mapstructure:"idle_timeout"` // 空闲连接超时(秒)
|
||||
TrustedProxies []string `mapstructure:"trusted_proxies"` // 可信代理列表(CIDR/IP)
|
||||
MaxRequestBodySize int64 `mapstructure:"max_request_body_size"` // 全局最大请求体限制
|
||||
@@ -796,6 +797,8 @@ type GatewayConfig struct {
|
||||
OpenAIHighEffortFirstOutputTimeoutSeconds int `mapstructure:"openai_high_effort_first_output_timeout_seconds"`
|
||||
// 请求体最大字节数,用于网关请求体大小限制
|
||||
MaxBodySize int64 `mapstructure:"max_body_size"`
|
||||
// TextMaxBodySize limits endpoints that cannot carry inline image/video payloads.
|
||||
TextMaxBodySize int64 `mapstructure:"text_max_body_size"`
|
||||
// 非流式上游响应体读取上限(字节),用于防止无界读取导致内存放大
|
||||
UpstreamResponseReadMaxBytes int64 `mapstructure:"upstream_response_read_max_bytes"`
|
||||
// 代理探测响应体读取上限(字节)
|
||||
@@ -1419,12 +1422,22 @@ type RateLimitConfig struct {
|
||||
|
||||
// APIKeyAuthCacheConfig API Key 认证缓存配置
|
||||
type APIKeyAuthCacheConfig struct {
|
||||
L1Size int `mapstructure:"l1_size"`
|
||||
L1TTLSeconds int `mapstructure:"l1_ttl_seconds"`
|
||||
L2TTLSeconds int `mapstructure:"l2_ttl_seconds"`
|
||||
NegativeTTLSeconds int `mapstructure:"negative_ttl_seconds"`
|
||||
JitterPercent int `mapstructure:"jitter_percent"`
|
||||
Singleflight bool `mapstructure:"singleflight"`
|
||||
L1Size int `mapstructure:"l1_size"`
|
||||
L1TTLSeconds int `mapstructure:"l1_ttl_seconds"`
|
||||
L2TTLSeconds int `mapstructure:"l2_ttl_seconds"`
|
||||
NegativeTTLSeconds int `mapstructure:"negative_ttl_seconds"`
|
||||
JitterPercent int `mapstructure:"jitter_percent"`
|
||||
Singleflight bool `mapstructure:"singleflight"`
|
||||
LookupConcurrency int `mapstructure:"lookup_concurrency"`
|
||||
InvalidAbuse InvalidAuthAbuseConfig `mapstructure:"invalid_abuse"`
|
||||
}
|
||||
|
||||
type InvalidAuthAbuseConfig struct {
|
||||
Enabled bool `mapstructure:"enabled"`
|
||||
Threshold int `mapstructure:"threshold"`
|
||||
WindowSeconds int `mapstructure:"window_seconds"`
|
||||
BlockSeconds int `mapstructure:"block_seconds"`
|
||||
Capacity int `mapstructure:"capacity"`
|
||||
}
|
||||
|
||||
// SubscriptionCacheConfig 订阅认证 L1 缓存配置
|
||||
@@ -1698,8 +1711,9 @@ func setDefaults() {
|
||||
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秒空闲超时
|
||||
viper.SetDefault("server.read_header_timeout", 10) // 10秒读取请求头
|
||||
viper.SetDefault("server.max_header_bytes", 64*1024)
|
||||
viper.SetDefault("server.idle_timeout", 120) // 120秒空闲超时
|
||||
viper.SetDefault("server.trusted_proxies", []string{})
|
||||
viper.SetDefault("server.max_request_body_size", int64(256*1024*1024))
|
||||
// H2C 默认配置
|
||||
@@ -1983,6 +1997,12 @@ func setDefaults() {
|
||||
viper.SetDefault("api_key_auth_cache.negative_ttl_seconds", 30)
|
||||
viper.SetDefault("api_key_auth_cache.jitter_percent", 10)
|
||||
viper.SetDefault("api_key_auth_cache.singleflight", true)
|
||||
viper.SetDefault("api_key_auth_cache.lookup_concurrency", 64)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.enabled", true)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.threshold", 120)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.window_seconds", 60)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.block_seconds", 60)
|
||||
viper.SetDefault("api_key_auth_cache.invalid_abuse.capacity", 16384)
|
||||
|
||||
// Subscription auth L1 cache
|
||||
viper.SetDefault("subscription_cache.l1_size", 16384)
|
||||
@@ -2111,6 +2131,7 @@ func setDefaults() {
|
||||
viper.SetDefault("gateway.antigravity_fallback_cooldown_minutes", 1)
|
||||
viper.SetDefault("gateway.antigravity_extra_retries", 10)
|
||||
viper.SetDefault("gateway.max_body_size", int64(256*1024*1024))
|
||||
viper.SetDefault("gateway.text_max_body_size", int64(32*1024*1024))
|
||||
viper.SetDefault("gateway.upstream_response_read_max_bytes", DefaultUpstreamResponseReadMaxBytes)
|
||||
viper.SetDefault("gateway.proxy_probe_response_read_max_bytes", int64(1024*1024))
|
||||
viper.SetDefault("gateway.gemini_debug_response_headers", false)
|
||||
@@ -2208,6 +2229,49 @@ func setDefaults() {
|
||||
}
|
||||
|
||||
func (c *Config) Validate() error {
|
||||
if c.Server.ReadHeaderTimeout < 1 || c.Server.ReadHeaderTimeout > 60 {
|
||||
return fmt.Errorf("server.read_header_timeout must be between 1 and 60 seconds")
|
||||
}
|
||||
if c.Server.MaxHeaderBytes < 8*1024 || c.Server.MaxHeaderBytes > 1024*1024 {
|
||||
return fmt.Errorf("server.max_header_bytes must be between 8192 and 1048576 bytes")
|
||||
}
|
||||
if c.Server.IdleTimeout <= 0 {
|
||||
return fmt.Errorf("server.idle_timeout must be positive")
|
||||
}
|
||||
if c.Server.MaxRequestBodySize < 0 {
|
||||
return fmt.Errorf("server.max_request_body_size must be non-negative")
|
||||
}
|
||||
if c.Server.H2C.Enabled {
|
||||
if c.Server.H2C.MaxConcurrentStreams == 0 {
|
||||
return fmt.Errorf("server.h2c.max_concurrent_streams must be positive")
|
||||
}
|
||||
if c.Server.H2C.IdleTimeout <= 0 {
|
||||
return fmt.Errorf("server.h2c.idle_timeout must be positive")
|
||||
}
|
||||
if c.Server.H2C.MaxReadFrameSize < 16*1024 || c.Server.H2C.MaxReadFrameSize > 16*1024*1024-1 {
|
||||
return fmt.Errorf("server.h2c.max_read_frame_size must be between 16384 and 16777215 bytes")
|
||||
}
|
||||
if c.Server.H2C.MaxUploadBufferPerConnection < 65535 {
|
||||
return fmt.Errorf("server.h2c.max_upload_buffer_per_connection must be at least 65535 bytes")
|
||||
}
|
||||
if c.Server.H2C.MaxUploadBufferPerStream <= 0 {
|
||||
return fmt.Errorf("server.h2c.max_upload_buffer_per_stream must be positive")
|
||||
}
|
||||
}
|
||||
if c.APIKeyAuth.InvalidAbuse.Enabled {
|
||||
if c.APIKeyAuth.InvalidAbuse.Threshold < 10 {
|
||||
return fmt.Errorf("api_key_auth_cache.invalid_abuse.threshold must be at least 10")
|
||||
}
|
||||
if c.APIKeyAuth.InvalidAbuse.WindowSeconds < 1 || c.APIKeyAuth.InvalidAbuse.WindowSeconds > 3600 {
|
||||
return fmt.Errorf("api_key_auth_cache.invalid_abuse.window_seconds must be between 1 and 3600")
|
||||
}
|
||||
if c.APIKeyAuth.InvalidAbuse.BlockSeconds < 1 || c.APIKeyAuth.InvalidAbuse.BlockSeconds > 3600 {
|
||||
return fmt.Errorf("api_key_auth_cache.invalid_abuse.block_seconds must be between 1 and 3600")
|
||||
}
|
||||
if c.APIKeyAuth.InvalidAbuse.Capacity < 256 || c.APIKeyAuth.InvalidAbuse.Capacity > 1_000_000 {
|
||||
return fmt.Errorf("api_key_auth_cache.invalid_abuse.capacity must be between 256 and 1000000")
|
||||
}
|
||||
}
|
||||
jwtSecret := strings.TrimSpace(c.JWT.Secret)
|
||||
if jwtSecret == "" {
|
||||
return fmt.Errorf("jwt.secret is required")
|
||||
@@ -2733,6 +2797,9 @@ func (c *Config) Validate() error {
|
||||
if c.Gateway.MaxBodySize <= 0 {
|
||||
return fmt.Errorf("gateway.max_body_size must be positive")
|
||||
}
|
||||
if c.Gateway.TextMaxBodySize <= 0 || c.Gateway.TextMaxBodySize > c.Gateway.MaxBodySize {
|
||||
return fmt.Errorf("gateway.text_max_body_size must be positive and no greater than gateway.max_body_size")
|
||||
}
|
||||
if c.Gateway.UpstreamResponseReadMaxBytes <= 0 {
|
||||
return fmt.Errorf("gateway.upstream_response_read_max_bytes must be positive")
|
||||
}
|
||||
|
||||
@@ -35,6 +35,18 @@ func TestLoadServerTimingConfig(t *testing.T) {
|
||||
})
|
||||
}
|
||||
|
||||
func TestLoadHTTPIngressSafetyDefaults(t *testing.T) {
|
||||
resetViperWithJWTSecret(t)
|
||||
cfg, err := Load()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 10, cfg.Server.ReadHeaderTimeout)
|
||||
require.Equal(t, 64*1024, cfg.Server.MaxHeaderBytes)
|
||||
require.Equal(t, int64(32*1024*1024), cfg.Gateway.TextMaxBodySize)
|
||||
require.True(t, cfg.APIKeyAuth.InvalidAbuse.Enabled)
|
||||
require.Equal(t, 120, cfg.APIKeyAuth.InvalidAbuse.Threshold)
|
||||
require.Equal(t, 16384, cfg.APIKeyAuth.InvalidAbuse.Capacity)
|
||||
}
|
||||
|
||||
func TestLoadForBootstrapAllowsMissingJWTSecret(t *testing.T) {
|
||||
viper.Reset()
|
||||
t.Setenv("JWT_SECRET", "")
|
||||
@@ -1185,6 +1197,51 @@ func TestValidateConfigErrors(t *testing.T) {
|
||||
mutate func(*Config)
|
||||
wantErr string
|
||||
}{
|
||||
{
|
||||
name: "server read header timeout",
|
||||
mutate: func(c *Config) { c.Server.ReadHeaderTimeout = 0 },
|
||||
wantErr: "server.read_header_timeout",
|
||||
},
|
||||
{
|
||||
name: "server max header bytes too small",
|
||||
mutate: func(c *Config) { c.Server.MaxHeaderBytes = 4096 },
|
||||
wantErr: "server.max_header_bytes",
|
||||
},
|
||||
{
|
||||
name: "server max request body size",
|
||||
mutate: func(c *Config) { c.Server.MaxRequestBodySize = -1 },
|
||||
wantErr: "server.max_request_body_size",
|
||||
},
|
||||
{
|
||||
name: "h2c zero concurrent streams",
|
||||
mutate: func(c *Config) {
|
||||
c.Server.H2C.Enabled = true
|
||||
c.Server.H2C.MaxConcurrentStreams = 0
|
||||
},
|
||||
wantErr: "server.h2c.max_concurrent_streams",
|
||||
},
|
||||
{
|
||||
name: "h2c oversized read frame",
|
||||
mutate: func(c *Config) {
|
||||
c.Server.H2C.Enabled = true
|
||||
c.Server.H2C.MaxReadFrameSize = 16 * 1024 * 1024
|
||||
},
|
||||
wantErr: "server.h2c.max_read_frame_size",
|
||||
},
|
||||
{
|
||||
name: "invalid auth abuse threshold too small",
|
||||
mutate: func(c *Config) {
|
||||
c.APIKeyAuth.InvalidAbuse.Threshold = 9
|
||||
},
|
||||
wantErr: "api_key_auth_cache.invalid_abuse.threshold",
|
||||
},
|
||||
{
|
||||
name: "invalid auth abuse capacity too small",
|
||||
mutate: func(c *Config) {
|
||||
c.APIKeyAuth.InvalidAbuse.Capacity = 255
|
||||
},
|
||||
wantErr: "api_key_auth_cache.invalid_abuse.capacity",
|
||||
},
|
||||
{
|
||||
name: "jwt secret required",
|
||||
mutate: func(c *Config) { c.JWT.Secret = "" },
|
||||
@@ -1386,6 +1443,11 @@ func TestValidateConfigErrors(t *testing.T) {
|
||||
mutate: func(c *Config) { c.Gateway.MaxBodySize = 0 },
|
||||
wantErr: "gateway.max_body_size",
|
||||
},
|
||||
{
|
||||
name: "gateway text body exceeds media body",
|
||||
mutate: func(c *Config) { c.Gateway.TextMaxBodySize = c.Gateway.MaxBodySize + 1 },
|
||||
wantErr: "gateway.text_max_body_size",
|
||||
},
|
||||
{
|
||||
name: "gateway response header timeout",
|
||||
mutate: func(c *Config) { c.Gateway.ResponseHeaderTimeout = -1 },
|
||||
|
||||
@@ -61,7 +61,7 @@ type AccountHandler struct {
|
||||
sessionLimitCache service.SessionLimitCache
|
||||
rpmCache service.RPMCache
|
||||
tokenCacheInvalidator service.TokenCacheInvalidator
|
||||
grokImportProber grokUsageProber
|
||||
grokImportProber grokImportProber
|
||||
upstreamBillingProbe *service.UpstreamBillingProbeService
|
||||
}
|
||||
|
||||
@@ -121,6 +121,7 @@ type CreateAccountRequest struct {
|
||||
GroupIDs []int64 `json:"group_ids"`
|
||||
ExpiresAt *int64 `json:"expires_at"`
|
||||
AutoPauseOnExpired *bool `json:"auto_pause_on_expired"`
|
||||
ProbeEnabled *bool `json:"upstream_billing_probe_enabled"`
|
||||
ConfirmMixedChannelRisk *bool `json:"confirm_mixed_channel_risk"` // 用户确认混合渠道风险
|
||||
}
|
||||
|
||||
@@ -159,6 +160,7 @@ type BulkUpdateAccountsRequest struct {
|
||||
GroupIDs *[]int64 `json:"group_ids"`
|
||||
Credentials map[string]any `json:"credentials"`
|
||||
Extra map[string]any `json:"extra"`
|
||||
ProbeEnabled *bool `json:"upstream_billing_probe_enabled"`
|
||||
ConfirmMixedChannelRisk *bool `json:"confirm_mixed_channel_risk"` // 用户确认混合渠道风险
|
||||
}
|
||||
|
||||
@@ -828,6 +830,7 @@ func (h *AccountHandler) Create(c *gin.Context) {
|
||||
GroupIDs: req.GroupIDs,
|
||||
ExpiresAt: req.ExpiresAt,
|
||||
AutoPauseOnExpired: req.AutoPauseOnExpired,
|
||||
ProbeEnabled: req.ProbeEnabled,
|
||||
SkipMixedChannelCheck: skipCheck,
|
||||
})
|
||||
if execErr != nil {
|
||||
@@ -1907,7 +1910,8 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) {
|
||||
req.Schedulable != nil ||
|
||||
req.GroupIDs != nil ||
|
||||
len(req.Credentials) > 0 ||
|
||||
len(req.Extra) > 0
|
||||
len(req.Extra) > 0 ||
|
||||
req.ProbeEnabled != nil
|
||||
|
||||
if !hasUpdates {
|
||||
response.BadRequest(c, "No updates provided")
|
||||
@@ -1928,6 +1932,7 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) {
|
||||
GroupIDs: req.GroupIDs,
|
||||
Credentials: req.Credentials,
|
||||
Extra: req.Extra,
|
||||
ProbeEnabled: req.ProbeEnabled,
|
||||
SkipMixedChannelCheck: skipCheck,
|
||||
})
|
||||
if err != nil {
|
||||
|
||||
@@ -222,3 +222,22 @@ func TestBulkUpdateAcceptsFilterTargetRequest(t *testing.T) {
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
|
||||
require.Equal(t, float64(0), resp["code"])
|
||||
}
|
||||
|
||||
func TestBulkUpdateAcceptsDedicatedUpstreamBillingProbeSetting(t *testing.T) {
|
||||
adminSvc := newStubAdminService()
|
||||
router := setupAccountMixedChannelRouter(adminSvc)
|
||||
|
||||
body, _ := json.Marshal(map[string]any{
|
||||
"account_ids": []int64{1, 2},
|
||||
"upstream_billing_probe_enabled": false,
|
||||
})
|
||||
rec := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/bulk-update", bytes.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
router.ServeHTTP(rec, req)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.NotNil(t, adminSvc.lastBulkUpdateAccountInput)
|
||||
require.NotNil(t, adminSvc.lastBulkUpdateAccountInput.ProbeEnabled)
|
||||
require.False(t, *adminSvc.lastBulkUpdateAccountInput.ProbeEnabled)
|
||||
}
|
||||
|
||||
@@ -16,7 +16,7 @@ func setupAdminRouter() (*gin.Engine, *stubAdminService) {
|
||||
router := gin.New()
|
||||
adminSvc := newStubAdminService()
|
||||
|
||||
userHandler := NewUserHandler(adminSvc, nil, nil, nil, nil, nil)
|
||||
userHandler := NewUserHandler(adminSvc, nil, nil, nil, nil, nil, nil)
|
||||
groupHandler := NewGroupHandler(adminSvc, nil, nil)
|
||||
proxyHandler := NewProxyHandler(adminSvc)
|
||||
redeemHandler := NewRedeemHandler(adminSvc, nil)
|
||||
|
||||
@@ -33,6 +33,7 @@ type stubAdminService struct {
|
||||
createSparkShadowErr error
|
||||
updateAccountErr error
|
||||
bulkUpdateAccountErr error
|
||||
lastBulkUpdateAccountInput *service.BulkUpdateAccountsInput
|
||||
getAccountResult *service.Account
|
||||
updateAccountCalls int
|
||||
updateAccountExtraCalls int
|
||||
@@ -195,6 +196,10 @@ func (s *stubAdminService) BatchUpdateConcurrency(ctx context.Context, userIDs [
|
||||
return len(userIDs), nil
|
||||
}
|
||||
|
||||
func (s *stubAdminService) BatchUpdateLimits(ctx context.Context, userIDs []int64, concurrency, rpmLimit *int) (int, error) {
|
||||
return len(userIDs), nil
|
||||
}
|
||||
|
||||
func (s *stubAdminService) GetUserAPIKeys(ctx context.Context, userID int64, page, pageSize int, sortBy, sortOrder string) ([]service.APIKey, int64, error) {
|
||||
return s.apiKeys, int64(len(s.apiKeys)), nil
|
||||
}
|
||||
@@ -293,6 +298,15 @@ func (s *stubAdminService) CreateGroup(ctx context.Context, input *service.Creat
|
||||
return &group, nil
|
||||
}
|
||||
|
||||
func (s *stubAdminService) DuplicateGroup(ctx context.Context, id int64, actorScope, operationKey string) (*service.Group, error) {
|
||||
group := service.Group{ID: 201, Name: "group (Copy)", Status: "inactive"}
|
||||
return &group, nil
|
||||
}
|
||||
|
||||
func (s *stubAdminService) RecoverDuplicateGroup(ctx context.Context, id int64, actorScope, operationKey string) (*service.Group, error) {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (s *stubAdminService) UpdateGroup(ctx context.Context, id int64, input *service.UpdateGroupInput) (*service.Group, error) {
|
||||
group := service.Group{ID: id, Name: input.Name, Status: service.StatusActive}
|
||||
return &group, nil
|
||||
@@ -465,6 +479,7 @@ func (s *stubAdminService) SetAccountSchedulable(ctx context.Context, id int64,
|
||||
}
|
||||
|
||||
func (s *stubAdminService) BulkUpdateAccounts(ctx context.Context, input *service.BulkUpdateAccountsInput) (*service.BulkUpdateAccountsResult, error) {
|
||||
s.lastBulkUpdateAccountInput = input
|
||||
if s.bulkUpdateAccountErr != nil {
|
||||
return nil, s.bulkUpdateAccountErr
|
||||
}
|
||||
|
||||
@@ -6,7 +6,6 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
@@ -150,7 +149,7 @@ func (h *AuditLogHandler) Clear(c *gin.Context) {
|
||||
CredentialMasked: middleware.MaskedRequestCredential(c),
|
||||
Method: http.MethodPost,
|
||||
Path: c.FullPath(),
|
||||
ClientIP: ip.GetTrustedClientIP(c),
|
||||
ClientIP: middleware.SecurityClientIP(c),
|
||||
UserAgent: c.Request.UserAgent(),
|
||||
StatusCode: http.StatusOK,
|
||||
}
|
||||
|
||||
@@ -15,12 +15,12 @@ const (
|
||||
grokImportProbeTimeout = 25 * time.Second
|
||||
)
|
||||
|
||||
type grokUsageProber interface {
|
||||
ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error)
|
||||
type grokImportProber interface {
|
||||
QueryQuota(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error)
|
||||
}
|
||||
|
||||
type grokImportProbeTask struct {
|
||||
prober grokUsageProber
|
||||
prober grokImportProber
|
||||
accountID int64
|
||||
}
|
||||
|
||||
@@ -51,7 +51,7 @@ func newGrokImportProbeScheduler(concurrency int, timeout time.Duration) *grokIm
|
||||
}
|
||||
}
|
||||
|
||||
func (s *grokImportProbeScheduler) schedule(prober grokUsageProber, account *service.Account) {
|
||||
func (s *grokImportProbeScheduler) schedule(prober grokImportProber, account *service.Account) {
|
||||
if s == nil || prober == nil || account == nil || account.ID <= 0 {
|
||||
return
|
||||
}
|
||||
@@ -97,7 +97,7 @@ func (s *grokImportProbeScheduler) nextTask() (grokImportProbeTask, bool) {
|
||||
return task, true
|
||||
}
|
||||
|
||||
func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64) {
|
||||
func (s *grokImportProbeScheduler) run(prober grokImportProber, accountID int64) {
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
slog.Error(
|
||||
@@ -112,7 +112,7 @@ func (s *grokImportProbeScheduler) run(prober grokUsageProber, accountID int64)
|
||||
// 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)
|
||||
result, err := prober.QueryQuota(ctx, accountID)
|
||||
if err != nil {
|
||||
slog.Warn(
|
||||
"grok_import_active_probe_failed",
|
||||
|
||||
@@ -36,7 +36,7 @@ func newGrokImportProbeStub(buffer int) *grokImportProbeStub {
|
||||
}
|
||||
}
|
||||
|
||||
func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) {
|
||||
func (s *grokImportProbeStub) QueryQuota(ctx context.Context, accountID int64) (*service.GrokQuotaProbeResult, error) {
|
||||
_, deadlineSeen := ctx.Deadline()
|
||||
s.mu.Lock()
|
||||
s.calls[accountID]++
|
||||
@@ -69,7 +69,7 @@ func (s *grokImportProbeStub) ProbeUsage(ctx context.Context, accountID int64) (
|
||||
return nil, failure
|
||||
}
|
||||
return &service.GrokQuotaProbeResult{
|
||||
Source: "active_probe",
|
||||
Source: "hybrid_probe",
|
||||
Model: "grok-4.5",
|
||||
StatusCode: 200,
|
||||
ResetSupported: false,
|
||||
|
||||
@@ -23,7 +23,7 @@ type GrokOAuthHandler struct {
|
||||
grokOAuthService *service.GrokOAuthService
|
||||
adminService service.AdminService
|
||||
quotaService *service.GrokQuotaService
|
||||
importProber grokUsageProber
|
||||
importProber grokImportProber
|
||||
reconciler service.GrokOAuthReconciler
|
||||
}
|
||||
|
||||
@@ -373,8 +373,7 @@ func (h *GrokOAuthHandler) createAccountFromSSOToken(ctx context.Context, req Gr
|
||||
return grokSSOImportWorkerResult{item: GrokSSOToOAuthItemResult{Index: index, Error: grokSSOImportErrorMessage(err)}}
|
||||
}
|
||||
|
||||
credentials := h.grokOAuthService.BuildAccountCredentials(tokenInfo)
|
||||
credentials = service.MergeCredentials(cloneGrokSSOMap(req.Credentials), credentials)
|
||||
credentials := grokSSOImportCredentials(h.grokOAuthService.BuildAccountCredentials(tokenInfo), req.Credentials)
|
||||
name := grokSSOImportAccountName(req.Name, tokenInfo, index, total)
|
||||
expiresAt, autoPauseOnExpired := grokSSOImportExpiry(req.ExpiresAt, req.AutoPauseOnExpired, tokenInfo)
|
||||
account, err := h.adminService.CreateAccount(ctx, &service.CreateAccountInput{
|
||||
@@ -408,6 +407,18 @@ func (h *GrokOAuthHandler) createAccountFromSSOToken(ctx context.Context, req Gr
|
||||
}
|
||||
}
|
||||
|
||||
// grokSSOImportCredentials 合并 SSO 兑换出的凭据与导入请求携带的运营侧配置。
|
||||
// token 字段以 BuildAccountCredentials 为准(请求不可覆盖);但 base_url 是运营侧
|
||||
// 配置且 Build 恒写官方地址,会吞掉导入时指定的自定义转发地址——与
|
||||
// RefreshAccountToken 的保留逻辑对齐,请求显式提供时以请求为准。
|
||||
func grokSSOImportCredentials(built map[string]any, reqCredentials map[string]any) map[string]any {
|
||||
credentials := service.MergeCredentials(cloneGrokSSOMap(reqCredentials), built)
|
||||
if reqBaseURL, ok := reqCredentials["base_url"].(string); ok && strings.TrimSpace(reqBaseURL) != "" {
|
||||
credentials["base_url"] = strings.TrimSpace(reqBaseURL)
|
||||
}
|
||||
return credentials
|
||||
}
|
||||
|
||||
func grokSSOImportExpiry(requestExpiresAt *int64, requestAutoPause *bool, tokenInfo *service.GrokTokenInfo) (*int64, *bool) {
|
||||
if tokenInfo == nil || strings.TrimSpace(tokenInfo.RefreshToken) != "" || tokenInfo.ExpiresAt <= 0 {
|
||||
return requestExpiresAt, requestAutoPause
|
||||
|
||||
@@ -136,8 +136,12 @@ func TestGrokOAuthHandlerQueryQuotaProbesUpstream(t *testing.T) {
|
||||
for i, upstreamReq := range requests {
|
||||
require.Equal(t, "Bearer access-token", upstreamReq.Header.Get("Authorization"))
|
||||
if upstreamReq.URL.String() == xai.DefaultCLIBaseURL+"/responses" {
|
||||
require.Equal(t, "application/json, text/event-stream", upstreamReq.Header.Get("Accept"))
|
||||
require.Contains(t, string(bodies[i]), `"model":"grok-4.5"`)
|
||||
require.Contains(t, string(bodies[i]), `"store":false`)
|
||||
require.Contains(t, string(bodies[i]), `"input":"hi"`)
|
||||
require.Contains(t, string(bodies[i]), `"stream":true`)
|
||||
require.NotContains(t, string(bodies[i]), `"max_output_tokens"`)
|
||||
require.NotContains(t, string(bodies[i]), `"store"`)
|
||||
}
|
||||
}
|
||||
require.NotNil(t, repo.updates[42])
|
||||
@@ -223,6 +227,45 @@ func TestGrokSSOImportExpiryPreservesRequestSettingsWithRefreshToken(t *testing.
|
||||
require.Same(t, &requestedAutoPause, autoPause)
|
||||
}
|
||||
|
||||
func TestGrokSSOImportCredentialsPreservesRequestedBaseURL(t *testing.T) {
|
||||
built := map[string]any{
|
||||
"access_token": "at-1",
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
}
|
||||
reqCredentials := map[string]any{
|
||||
"base_url": "https://relay.example.com/v1",
|
||||
"header_override_enabled": true,
|
||||
"header_overrides": map[string]any{"x-relay-key": "k"},
|
||||
}
|
||||
|
||||
credentials := grokSSOImportCredentials(built, reqCredentials)
|
||||
|
||||
// token 字段以兑换结果为准;base_url 是运营侧配置,必须保留请求里的自定义地址
|
||||
require.Equal(t, "at-1", credentials["access_token"])
|
||||
require.Equal(t, "https://relay.example.com/v1", credentials["base_url"])
|
||||
require.Equal(t, true, credentials["header_override_enabled"])
|
||||
require.Equal(t, map[string]any{"x-relay-key": "k"}, credentials["header_overrides"])
|
||||
// 入参不被污染(req.Credentials 会被多个 worker 并发读取)
|
||||
require.Equal(t, "https://relay.example.com/v1", reqCredentials["base_url"])
|
||||
}
|
||||
|
||||
func TestGrokSSOImportCredentialsDefaultsToOfficialBaseURL(t *testing.T) {
|
||||
built := map[string]any{
|
||||
"access_token": "at-1",
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
}
|
||||
|
||||
credentials := grokSSOImportCredentials(built, nil)
|
||||
require.Equal(t, xai.DefaultCLIBaseURL, credentials["base_url"])
|
||||
|
||||
credentials = grokSSOImportCredentials(map[string]any{
|
||||
"access_token": "at-2",
|
||||
"base_url": xai.DefaultCLIBaseURL,
|
||||
}, map[string]any{"base_url": " "})
|
||||
require.Equal(t, xai.DefaultCLIBaseURL, credentials["base_url"])
|
||||
require.Equal(t, "at-2", credentials["access_token"])
|
||||
}
|
||||
|
||||
func TestGrokSSOImportWorkerRecoversPanic(t *testing.T) {
|
||||
h := &GrokOAuthHandler{}
|
||||
result := h.safeCreateAccountFromSSOToken(context.Background(), GrokSSOToOAuthRequest{}, "token", 2, 3)
|
||||
|
||||
@@ -2,12 +2,15 @@ package admin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"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/timezone"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
@@ -361,6 +364,53 @@ func (h *GroupHandler) Create(c *gin.Context) {
|
||||
response.Success(c, dto.GroupFromServiceAdmin(group))
|
||||
}
|
||||
|
||||
// Duplicate handles creating an inactive group copy with the source account bindings.
|
||||
// POST /api/v1/admin/groups/:id/duplicate
|
||||
func (h *GroupHandler) Duplicate(c *gin.Context) {
|
||||
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil || groupID <= 0 {
|
||||
response.BadRequest(c, "Invalid group ID")
|
||||
return
|
||||
}
|
||||
actorScope := adminActorScope(c)
|
||||
|
||||
result, err := executeAdminIdempotent(
|
||||
c,
|
||||
"admin.groups.duplicate",
|
||||
struct {
|
||||
GroupID int64 `json:"group_id"`
|
||||
}{GroupID: groupID},
|
||||
service.DefaultWriteIdempotencyTTL(),
|
||||
func(ctx context.Context) (any, error) {
|
||||
group, execErr := h.adminService.DuplicateGroup(ctx, groupID, actorScope, c.GetHeader("Idempotency-Key"))
|
||||
if execErr != nil {
|
||||
return nil, execErr
|
||||
}
|
||||
return dto.GroupFromServiceAdmin(group), nil
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
reason := infraerrors.Reason(err)
|
||||
if reason == infraerrors.Reason(service.ErrIdempotencyInProgress) || reason == infraerrors.Reason(service.ErrIdempotencyStoreUnavail) {
|
||||
recovered, recoverErr := h.adminService.RecoverDuplicateGroup(c.Request.Context(), groupID, actorScope, c.GetHeader("Idempotency-Key"))
|
||||
if recoverErr != nil {
|
||||
slog.Warn("group_duplicate_recovery_failed", "group_id", groupID, "actor_scope", actorScope, "reason", reason, "error", recoverErr)
|
||||
} else if recovered != nil {
|
||||
c.Header("X-Idempotency-Recovered", "true")
|
||||
response.Success(c, dto.GroupFromServiceAdmin(recovered))
|
||||
return
|
||||
}
|
||||
}
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
if result != nil && result.Replayed {
|
||||
c.Header("X-Idempotency-Replayed", "true")
|
||||
}
|
||||
response.Success(c, result.Data)
|
||||
}
|
||||
|
||||
// Update handles updating a group
|
||||
// PUT /api/v1/admin/groups/:id
|
||||
func (h *GroupHandler) Update(c *gin.Context) {
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
//go:build unit
|
||||
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
middleware2 "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"
|
||||
)
|
||||
|
||||
type duplicateGroupAdminServiceStub struct {
|
||||
service.AdminService
|
||||
group *service.Group
|
||||
calls int
|
||||
recoverCalls int
|
||||
groupID int64
|
||||
actorScope string
|
||||
operationKey string
|
||||
recoverScope string
|
||||
recoverKey string
|
||||
created bool
|
||||
}
|
||||
|
||||
func (s *duplicateGroupAdminServiceStub) DuplicateGroup(_ context.Context, groupID int64, actorScope, operationKey string) (*service.Group, error) {
|
||||
s.calls++
|
||||
s.groupID = groupID
|
||||
s.actorScope = actorScope
|
||||
s.operationKey = operationKey
|
||||
s.created = true
|
||||
return s.group, nil
|
||||
}
|
||||
|
||||
func (s *duplicateGroupAdminServiceStub) RecoverDuplicateGroup(_ context.Context, _ int64, actorScope, operationKey string) (*service.Group, error) {
|
||||
s.recoverCalls++
|
||||
s.recoverScope = actorScope
|
||||
s.recoverKey = operationKey
|
||||
if !s.created {
|
||||
return nil, nil
|
||||
}
|
||||
return s.group, nil
|
||||
}
|
||||
|
||||
func setupDuplicateGroupRouter(t *testing.T, svc service.AdminService) *gin.Engine {
|
||||
t.Helper()
|
||||
previousCoordinator := service.DefaultIdempotencyCoordinator()
|
||||
service.SetDefaultIdempotencyCoordinator(nil)
|
||||
t.Cleanup(func() { service.SetDefaultIdempotencyCoordinator(previousCoordinator) })
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 77})
|
||||
c.Next()
|
||||
})
|
||||
handler := NewGroupHandler(svc, nil, nil)
|
||||
router.POST("/api/v1/admin/groups/:id/duplicate", handler.Duplicate)
|
||||
return router
|
||||
}
|
||||
|
||||
func duplicateGroupHandlerFixture() *service.Group {
|
||||
return &service.Group{
|
||||
ID: 43,
|
||||
Name: "primary (Copy)",
|
||||
Platform: service.PlatformAnthropic,
|
||||
Status: "inactive",
|
||||
RateMultiplier: 1,
|
||||
AccountCount: 3,
|
||||
ActiveAccountCount: 2,
|
||||
DuplicateOperationID: "internal-operation-must-not-leak",
|
||||
ModelRouting: map[string][]int64{"claude-*": {7}},
|
||||
}
|
||||
}
|
||||
|
||||
func TestDuplicateGroupHandlerReturnsAdminDTOWithoutOperationMetadata(t *testing.T) {
|
||||
svc := &duplicateGroupAdminServiceStub{group: duplicateGroupHandlerFixture()}
|
||||
router := setupDuplicateGroupRouter(t, svc)
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPost, "/api/v1/admin/groups/42/duplicate", nil)
|
||||
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
require.Equal(t, 1, svc.calls)
|
||||
require.Contains(t, recorder.Body.String(), `"name":"primary (Copy)"`)
|
||||
require.Contains(t, recorder.Body.String(), `"status":"inactive"`)
|
||||
require.Contains(t, recorder.Body.String(), `"account_count":3`)
|
||||
require.NotContains(t, recorder.Body.String(), "duplicate_operation_id")
|
||||
require.NotContains(t, recorder.Body.String(), "internal-operation-must-not-leak")
|
||||
}
|
||||
|
||||
func TestDuplicateGroupHandlerRejectsInvalidID(t *testing.T) {
|
||||
for _, id := range []string{"not-a-number", "0", "-1"} {
|
||||
t.Run(id, func(t *testing.T) {
|
||||
svc := &duplicateGroupAdminServiceStub{}
|
||||
router := setupDuplicateGroupRouter(t, svc)
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPost, "/api/v1/admin/groups/"+id+"/duplicate", nil)
|
||||
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
require.Equal(t, http.StatusBadRequest, recorder.Code)
|
||||
require.Zero(t, svc.calls)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestDuplicateGroupHandlerReplaysSameIdempotencyKey(t *testing.T) {
|
||||
svc := &duplicateGroupAdminServiceStub{group: duplicateGroupHandlerFixture()}
|
||||
router := setupDuplicateGroupRouter(t, svc)
|
||||
service.SetDefaultIdempotencyCoordinator(service.NewIdempotencyCoordinator(newMemoryIdempotencyRepoStub(), service.DefaultIdempotencyConfig()))
|
||||
|
||||
call := func() *httptest.ResponseRecorder {
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPost, "/api/v1/admin/groups/42/duplicate", nil)
|
||||
request.Header.Set("Idempotency-Key", "duplicate-group-42")
|
||||
router.ServeHTTP(recorder, request)
|
||||
return recorder
|
||||
}
|
||||
|
||||
first := call()
|
||||
second := call()
|
||||
|
||||
require.Equal(t, http.StatusOK, first.Code)
|
||||
require.Equal(t, http.StatusOK, second.Code)
|
||||
require.Equal(t, 1, svc.calls)
|
||||
require.Equal(t, int64(42), svc.groupID)
|
||||
require.Equal(t, "admin:77", svc.actorScope)
|
||||
require.Equal(t, "duplicate-group-42", svc.operationKey)
|
||||
require.Equal(t, "true", second.Header().Get("X-Idempotency-Replayed"))
|
||||
}
|
||||
|
||||
func TestDuplicateGroupHandlerRecoversAfterMarkSucceededFailure(t *testing.T) {
|
||||
svc := &duplicateGroupAdminServiceStub{group: duplicateGroupHandlerFixture()}
|
||||
router := setupDuplicateGroupRouter(t, svc)
|
||||
repo := &failOnceMarkSucceededRepo{memoryIdempotencyRepoStub: newMemoryIdempotencyRepoStub(), failNext: true}
|
||||
service.SetDefaultIdempotencyCoordinator(service.NewIdempotencyCoordinator(repo, service.DefaultIdempotencyConfig()))
|
||||
|
||||
call := func() *httptest.ResponseRecorder {
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPost, "/api/v1/admin/groups/42/duplicate", nil)
|
||||
request.Header.Set("Idempotency-Key", "duplicate-group-42-recovery")
|
||||
router.ServeHTTP(recorder, request)
|
||||
return recorder
|
||||
}
|
||||
|
||||
first := call()
|
||||
second := call()
|
||||
|
||||
require.Equal(t, http.StatusOK, first.Code)
|
||||
require.Equal(t, http.StatusOK, second.Code)
|
||||
require.Equal(t, "true", first.Header().Get("X-Idempotency-Recovered"))
|
||||
require.Equal(t, "true", second.Header().Get("X-Idempotency-Recovered"))
|
||||
require.Equal(t, 1, svc.calls, "ambiguous retries must not repeat the create side effect")
|
||||
require.Equal(t, 2, svc.recoverCalls)
|
||||
require.Equal(t, "admin:77", svc.recoverScope)
|
||||
require.Equal(t, "duplicate-group-42-recovery", svc.recoverKey)
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// GetAuthCacheInvalidationHealth exposes durable outbox lag and subscriber health.
|
||||
func (h *OpsHandler) GetAuthCacheInvalidationHealth(c *gin.Context) {
|
||||
if h.opsService == nil {
|
||||
response.Error(c, http.StatusServiceUnavailable, "Ops service not available")
|
||||
return
|
||||
}
|
||||
if err := h.opsService.RequireMonitoringEnabled(c.Request.Context()); err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, h.opsService.GetAuthCacheInvalidationHealth(c.Request.Context()))
|
||||
}
|
||||
@@ -0,0 +1,131 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/netip"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
var ingressRejectReasons = map[string]struct{}{
|
||||
"query_api_key_deprecated": {}, "api_key_required": {}, "invalid_api_key": {},
|
||||
"invalid_auth_rate_limited": {},
|
||||
"api_key_auth_overloaded": {},
|
||||
"api_key_disabled": {}, "ip_restricted": {}, "user_inactive": {}, "group_deleted": {},
|
||||
"group_disabled": {}, "group_not_allowed": {}, "group_unassigned": {}, "other": {},
|
||||
}
|
||||
|
||||
var ingressRejectRouteFamilies = map[string]struct{}{
|
||||
"antigravity": {}, "gemini": {}, "codex": {}, "messages": {}, "responses": {},
|
||||
"chat_completions": {}, "images": {}, "videos": {}, "embeddings": {}, "models": {}, "other": {},
|
||||
}
|
||||
|
||||
var ingressRejectProtocols = map[string]struct{}{
|
||||
"google": {}, "anthropic": {}, "openai": {}, "gateway": {}, "other": {},
|
||||
}
|
||||
|
||||
// ListIngressRejects returns bounded security aggregates, never raw credentials or request bodies.
|
||||
func (h *OpsHandler) ListIngressRejects(c *gin.Context) {
|
||||
if h.opsService == nil {
|
||||
response.Error(c, http.StatusServiceUnavailable, "Ops service not available")
|
||||
return
|
||||
}
|
||||
page, pageSize := response.ParsePagination(c)
|
||||
if pageSize > 200 {
|
||||
pageSize = 200
|
||||
}
|
||||
startTime, endTime, err := parseOpsTimeRange(c, "1h")
|
||||
if err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
filter := &service.OpsIngressRejectFilter{Page: page, PageSize: pageSize}
|
||||
if !startTime.IsZero() {
|
||||
filter.StartTime = &startTime
|
||||
}
|
||||
if !endTime.IsZero() {
|
||||
filter.EndTime = &endTime
|
||||
}
|
||||
if filter.RejectReason, err = parseIngressRejectEnum(c, "reason", ingressRejectReasons); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
if filter.RouteFamily, err = parseIngressRejectEnum(c, "route_family", ingressRejectRouteFamilies); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
if filter.Protocol, err = parseIngressRejectEnum(c, "protocol", ingressRejectProtocols); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
if raw := strings.TrimSpace(c.Query("client_ip")); raw != "" {
|
||||
addr, parseErr := netip.ParseAddr(raw)
|
||||
if parseErr != nil {
|
||||
response.BadRequest(c, "Invalid client_ip")
|
||||
return
|
||||
}
|
||||
addr = addr.Unmap()
|
||||
if addr.Is6() {
|
||||
addr = netip.PrefixFrom(addr, 64).Masked().Addr()
|
||||
}
|
||||
filter.ClientIP = addr.String()
|
||||
}
|
||||
if filter.UserID, err = parseOptionalPositiveID(c, "user_id"); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
if filter.APIKeyID, err = parseOptionalPositiveID(c, "api_key_id"); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
result, err := h.opsService.ListIngressRejects(c.Request.Context(), filter)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, result)
|
||||
}
|
||||
|
||||
func (h *OpsHandler) GetIngressRejectHealth(c *gin.Context) {
|
||||
if h.opsService == nil {
|
||||
response.Error(c, http.StatusServiceUnavailable, "Ops service not available")
|
||||
return
|
||||
}
|
||||
if err := h.opsService.RequireMonitoringEnabled(c.Request.Context()); err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, h.opsService.GetIngressRejectHealth())
|
||||
}
|
||||
|
||||
func parseIngressRejectEnum(c *gin.Context, name string, allowed map[string]struct{}) (string, error) {
|
||||
value := strings.TrimSpace(c.Query(name))
|
||||
if value == "" {
|
||||
return "", nil
|
||||
}
|
||||
if _, ok := allowed[value]; !ok {
|
||||
return "", &ingressRejectQueryError{message: "Invalid " + name}
|
||||
}
|
||||
return value, nil
|
||||
}
|
||||
|
||||
func parseOptionalPositiveID(c *gin.Context, name string) (*int64, error) {
|
||||
raw := strings.TrimSpace(c.Query(name))
|
||||
if raw == "" {
|
||||
return nil, nil
|
||||
}
|
||||
value, err := strconv.ParseInt(raw, 10, 64)
|
||||
if err != nil || value <= 0 {
|
||||
return nil, &ingressRejectQueryError{message: "Invalid " + name}
|
||||
}
|
||||
return &value, nil
|
||||
}
|
||||
|
||||
type ingressRejectQueryError struct{ message string }
|
||||
|
||||
func (e *ingressRejectQueryError) Error() string { return e.message }
|
||||
@@ -0,0 +1,41 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newIngressRejectHandlerForTest() *OpsHandler {
|
||||
return NewOpsHandler(service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil))
|
||||
}
|
||||
|
||||
func TestListIngressRejectsValidatesFilters(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, path := range []string{
|
||||
"/api/v1/admin/ops/ingress-rejections?reason=not-valid",
|
||||
"/api/v1/admin/ops/ingress-rejections?client_ip=not-an-ip",
|
||||
"/api/v1/admin/ops/ingress-rejections?user_id=0",
|
||||
"/api/v1/admin/ops/ingress-rejections?api_key_id=-1",
|
||||
} {
|
||||
t.Run(path, func(t *testing.T) {
|
||||
context, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
context.Request = httptest.NewRequest(http.MethodGet, path, nil)
|
||||
newIngressRejectHandlerForTest().ListIngressRejects(context)
|
||||
require.Equal(t, http.StatusBadRequest, context.Writer.Status())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetIngressRejectHealthAvailable(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
context, _ := gin.CreateTestContext(recorder)
|
||||
context.Request = httptest.NewRequest(http.MethodGet, "/api/v1/admin/ops/ingress-rejections/health", nil)
|
||||
newIngressRejectHandlerForTest().GetIngressRejectHealth(context)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), `"capacity":8192`)
|
||||
}
|
||||
@@ -59,6 +59,8 @@ type SettingHandler struct {
|
||||
paymentService *service.PaymentService
|
||||
userAttributeService *service.UserAttributeService
|
||||
notificationEmailService *service.NotificationEmailService
|
||||
totpService *service.TotpService
|
||||
userService *service.UserService
|
||||
}
|
||||
|
||||
// NewSettingHandler 创建系统设置处理器
|
||||
@@ -80,6 +82,15 @@ func (h *SettingHandler) SetNotificationEmailService(notificationEmailService *s
|
||||
h.notificationEmailService = notificationEmailService
|
||||
}
|
||||
|
||||
// SetStepUpDeps attaches the services backing the step-up switch preconditions
|
||||
// (enable requires the acting admin to have TOTP enabled; disable is itself a
|
||||
// step-up gated operation), without changing the constructor signature used by
|
||||
// existing unit tests.
|
||||
func (h *SettingHandler) SetStepUpDeps(totpService *service.TotpService, userService *service.UserService) {
|
||||
h.totpService = totpService
|
||||
h.userService = userService
|
||||
}
|
||||
|
||||
// GetSettings 获取所有系统设置
|
||||
// GET /api/v1/admin/settings
|
||||
func (h *SettingHandler) GetSettings(c *gin.Context) {
|
||||
@@ -124,6 +135,7 @@ func (h *SettingHandler) GetSettings(c *gin.Context) {
|
||||
TotpEnabled: settings.TotpEnabled,
|
||||
TotpEncryptionKeyConfigured: h.settingService.IsTotpEncryptionKeyConfigured(),
|
||||
SessionBindingEnabled: settings.SessionBindingEnabled,
|
||||
StepUpEnabled: settings.StepUpEnabled,
|
||||
AuditLogRetentionDays: settings.AuditLogRetentionDays,
|
||||
LoginAgreementEnabled: settings.LoginAgreementEnabled,
|
||||
LoginAgreementMode: settings.LoginAgreementMode,
|
||||
|
||||
@@ -56,6 +56,12 @@ func diffSettings(before *service.SystemSettings, after *service.SystemSettings,
|
||||
if before.TotpEnabled != after.TotpEnabled {
|
||||
changed = append(changed, "totp_enabled")
|
||||
}
|
||||
if before.SessionBindingEnabled != after.SessionBindingEnabled {
|
||||
changed = append(changed, "session_binding_enabled")
|
||||
}
|
||||
if before.StepUpEnabled != after.StepUpEnabled {
|
||||
changed = append(changed, "step_up_enabled")
|
||||
}
|
||||
if before.LoginAgreementEnabled != after.LoginAgreementEnabled {
|
||||
changed = append(changed, "login_agreement_enabled")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"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"
|
||||
)
|
||||
|
||||
// step-up 开关转换的门控测试。
|
||||
// 测试环境不注入认证上下文/userService,因此一旦触发校验会以 401/403/500 中止;
|
||||
// 借此区分「触发了转换校验」与「直接放行到常规保存(200)」。
|
||||
|
||||
func newStepUpSwitchTestHandler(t *testing.T, stored map[string]string) (*SettingHandler, *settingHandlerRepoStub) {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
repo := &settingHandlerRepoStub{values: stored}
|
||||
svc := service.NewSettingService(repo, &config.Config{Default: config.DefaultConfig{UserConcurrency: 5}})
|
||||
return NewSettingHandler(svc, nil, nil, nil, nil, nil, nil), repo
|
||||
}
|
||||
|
||||
func doUpdateSettings(t *testing.T, h *SettingHandler, body map[string]any, prepare func(c *gin.Context)) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
rawBody, err := json.Marshal(body)
|
||||
require.NoError(t, err)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPut, "/api/v1/admin/settings", bytes.NewReader(rawBody))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
if prepare != nil {
|
||||
prepare(c)
|
||||
}
|
||||
|
||||
h.UpdateSettings(c)
|
||||
return rec
|
||||
}
|
||||
|
||||
// 开启开关(false→true):无认证上下文时拒绝,且带专用错误标记。
|
||||
func TestUpdateSettingsEnableStepUpRejectsWithoutSession(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"step_up_enabled": true}, nil)
|
||||
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "STEP_UP_ENABLE_REQUIRES_TOTP")
|
||||
require.NotEqual(t, "true", repo.values[service.SettingKeyStepUpEnabled])
|
||||
}
|
||||
|
||||
// 开启开关:admin API key(机器凭证)一律拒绝,reason 与门控保持一致便于前端分流。
|
||||
func TestUpdateSettingsEnableStepUpRejectsAdminAPIKey(t *testing.T) {
|
||||
h, _ := newStepUpSwitchTestHandler(t, map[string]string{})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"step_up_enabled": true}, func(c *gin.Context) {
|
||||
c.Set("auth_method", service.AuditAuthMethodAdminAPIKey)
|
||||
})
|
||||
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "STEP_UP_ADMIN_API_KEY_FORBIDDEN")
|
||||
}
|
||||
|
||||
// 开启开关:有认证会话但 userService 未注入时 fail-closed(500),不得放行。
|
||||
func TestUpdateSettingsEnableStepUpFailsClosedWithoutUserService(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"step_up_enabled": true}, func(c *gin.Context) {
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: 1})
|
||||
})
|
||||
|
||||
require.Equal(t, http.StatusInternalServerError, rec.Code)
|
||||
require.NotEqual(t, "true", repo.values[service.SettingKeyStepUpEnabled])
|
||||
}
|
||||
|
||||
// 关闭开关(true→false)本身是敏感操作:无认证上下文时被 step-up 门控以 401 拦截。
|
||||
func TestUpdateSettingsDisableStepUpRequiresStepUp(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeyStepUpEnabled: "true",
|
||||
})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"step_up_enabled": false}, nil)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, rec.Code)
|
||||
require.Equal(t, "true", repo.values[service.SettingKeyStepUpEnabled])
|
||||
}
|
||||
|
||||
// 关闭开关:admin API key 被 step-up 门控以 403 拦截。
|
||||
func TestUpdateSettingsDisableStepUpRejectsAdminAPIKey(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeyStepUpEnabled: "true",
|
||||
})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"step_up_enabled": false}, func(c *gin.Context) {
|
||||
c.Set("auth_method", service.AuditAuthMethodAdminAPIKey)
|
||||
})
|
||||
|
||||
require.Equal(t, http.StatusForbidden, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "STEP_UP_ADMIN_API_KEY_FORBIDDEN")
|
||||
require.Equal(t, "true", repo.values[service.SettingKeyStepUpEnabled])
|
||||
}
|
||||
|
||||
// 无状态转换(false→false):不触发任何转换校验,常规保存成功且默认持久化为 false。
|
||||
func TestUpdateSettingsStepUpNoTransitionSkipsGate(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"step_up_enabled": false}, nil)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, "false", repo.values[service.SettingKeyStepUpEnabled])
|
||||
// 会话 IP/UA 绑定默认关闭:未显式提交时持久化 false。
|
||||
require.Equal(t, "false", repo.values[service.SettingKeySessionBindingEnabled])
|
||||
}
|
||||
|
||||
// 保持开启(true→true):不触发转换校验,常规保存不被打断。
|
||||
func TestUpdateSettingsStepUpKeepEnabledSkipsGate(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeyStepUpEnabled: "true",
|
||||
})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"step_up_enabled": true}, nil)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, "true", repo.values[service.SettingKeyStepUpEnabled])
|
||||
}
|
||||
|
||||
// 省略字段=保持现值:不含 step_up_enabled/session_binding_enabled 的旧客户端全量保存
|
||||
// 不得把已开启的安全开关静默重置,也不触发任何转换门控。
|
||||
func TestUpdateSettingsOmittedSecuritySwitchesKeepStoredValues(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{
|
||||
service.SettingKeyStepUpEnabled: "true",
|
||||
service.SettingKeySessionBindingEnabled: "true",
|
||||
})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"registration_enabled": true}, nil)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, "true", repo.values[service.SettingKeyStepUpEnabled])
|
||||
require.Equal(t, "true", repo.values[service.SettingKeySessionBindingEnabled])
|
||||
}
|
||||
|
||||
// 省略字段在开关本就关闭时同样保持关闭(默认值路径)。
|
||||
func TestUpdateSettingsOmittedSecuritySwitchesKeepDisabled(t *testing.T) {
|
||||
h, repo := newStepUpSwitchTestHandler(t, map[string]string{})
|
||||
|
||||
rec := doUpdateSettings(t, h, map[string]any{"registration_enabled": true}, nil)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Equal(t, "false", repo.values[service.SettingKeyStepUpEnabled])
|
||||
require.Equal(t, "false", repo.values[service.SettingKeySessionBindingEnabled])
|
||||
}
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -27,7 +28,8 @@ type UpdateSettingsRequest struct {
|
||||
FrontendURL string `json:"frontend_url"`
|
||||
InvitationCodeEnabled bool `json:"invitation_code_enabled"`
|
||||
TotpEnabled bool `json:"totp_enabled"` // TOTP 双因素认证
|
||||
SessionBindingEnabled bool `json:"session_binding_enabled"` // 会话 IP/UA 绑定
|
||||
SessionBindingEnabled *bool `json:"session_binding_enabled"` // 会话 IP/UA 绑定(省略=保持现值)
|
||||
StepUpEnabled *bool `json:"step_up_enabled"` // 敏感操作 step-up 2FA(省略=保持现值)
|
||||
AuditLogRetentionDays int `json:"audit_log_retention_days"` // 审计日志保留天数
|
||||
LoginAgreementEnabled bool `json:"login_agreement_enabled"`
|
||||
LoginAgreementMode string `json:"login_agreement_mode"`
|
||||
@@ -335,6 +337,41 @@ type UpdateSettingsRequest struct {
|
||||
|
||||
// UpdateSettings 更新系统设置
|
||||
// PUT /api/v1/admin/settings
|
||||
// ensureActorTotpForStepUp 校验当前操作者具备开启 step-up 门控的条件:
|
||||
// 必须是真人管理员会话(admin API key 无法完成 TOTP step-up,拒绝)且本人已启用 TOTP。
|
||||
// 校验失败时写入错误响应并返回 false。
|
||||
func (h *SettingHandler) ensureActorTotpForStepUp(c *gin.Context) bool {
|
||||
if c.GetString("auth_method") == service.AuditAuthMethodAdminAPIKey {
|
||||
response.ErrorWithDetails(c, http.StatusForbidden,
|
||||
"Admin API key cannot enable step-up verification; use an admin session with TOTP enabled",
|
||||
"STEP_UP_ADMIN_API_KEY_FORBIDDEN", nil)
|
||||
return false
|
||||
}
|
||||
subject, ok := middleware.GetAuthSubjectFromContext(c)
|
||||
if !ok || subject.UserID <= 0 {
|
||||
response.ErrorWithDetails(c, http.StatusForbidden,
|
||||
"Enabling step-up verification requires an authenticated admin session",
|
||||
"STEP_UP_ENABLE_REQUIRES_TOTP", nil)
|
||||
return false
|
||||
}
|
||||
if h.userService == nil {
|
||||
response.InternalError(c, "Step-up precondition check unavailable")
|
||||
return false
|
||||
}
|
||||
user, err := h.userService.GetByID(c.Request.Context(), subject.UserID)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return false
|
||||
}
|
||||
if !user.TotpEnabled {
|
||||
response.ErrorWithDetails(c, http.StatusBadRequest,
|
||||
"Enable two-factor authentication (TOTP) for your account before turning on step-up verification",
|
||||
"STEP_UP_ENABLE_REQUIRES_TOTP", nil)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
var req UpdateSettingsRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
@@ -353,6 +390,34 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 两个安全开关的请求字段为指针:省略字段=保持现值,避免旧客户端/脚本
|
||||
// 用不含新字段的全量 payload 保存设置时把安全开关静默重置。
|
||||
sessionBindingEnabled := previousSettings.SessionBindingEnabled
|
||||
if req.SessionBindingEnabled != nil {
|
||||
sessionBindingEnabled = *req.SessionBindingEnabled
|
||||
}
|
||||
stepUpEnabled := previousSettings.StepUpEnabled
|
||||
if req.StepUpEnabled != nil {
|
||||
stepUpEnabled = *req.StepUpEnabled
|
||||
}
|
||||
|
||||
// 开启敏感操作 step-up 门控属自锁风险操作:仅允许本人已启用 TOTP 的管理员会话开启,
|
||||
// 否则开启后操作者立即被挡在所有敏感操作之外。仅在 false→true 的开启瞬间校验,
|
||||
// 保持开启状态的常规设置保存不受影响。
|
||||
if stepUpEnabled && !previousSettings.StepUpEnabled {
|
||||
if !h.ensureActorTotpForStepUp(c) {
|
||||
return
|
||||
}
|
||||
}
|
||||
// 关闭 step-up 门控本身就是敏感操作:防止拿到管理员会话的攻击者先关闸再执行导出/备份。
|
||||
// previousSettings 已证实开关处于开启状态,使用无条件门控变体,
|
||||
// 避免门控内部二次读取开关时因存储故障 fail-open(前端捕获 STEP_UP_REQUIRED 弹码重试)。
|
||||
if !stepUpEnabled && previousSettings.StepUpEnabled {
|
||||
if !middleware.EnforceStepUpAlways(c, h.totpService, h.userService) {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// 验证参数
|
||||
if req.DefaultConcurrency < 1 {
|
||||
req.DefaultConcurrency = 1
|
||||
@@ -1181,7 +1246,8 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
FrontendURL: req.FrontendURL,
|
||||
InvitationCodeEnabled: req.InvitationCodeEnabled,
|
||||
TotpEnabled: req.TotpEnabled,
|
||||
SessionBindingEnabled: req.SessionBindingEnabled,
|
||||
SessionBindingEnabled: sessionBindingEnabled,
|
||||
StepUpEnabled: stepUpEnabled,
|
||||
AuditLogRetentionDays: req.AuditLogRetentionDays,
|
||||
LoginAgreementEnabled: req.LoginAgreementEnabled,
|
||||
LoginAgreementMode: loginAgreementMode,
|
||||
@@ -1622,6 +1688,9 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
if h.opsService != nil {
|
||||
h.opsService.SetMonitoringEnabled(settings.OpsMonitoringEnabled)
|
||||
}
|
||||
|
||||
// Update OpenAI fast policy (stored under dedicated key, only when provided).
|
||||
if req.OpenAIFastPolicySettings != nil {
|
||||
@@ -1710,6 +1779,7 @@ func (h *SettingHandler) UpdateSettings(c *gin.Context) {
|
||||
TotpEnabled: updatedSettings.TotpEnabled,
|
||||
TotpEncryptionKeyConfigured: h.settingService.IsTotpEncryptionKeyConfigured(),
|
||||
SessionBindingEnabled: updatedSettings.SessionBindingEnabled,
|
||||
StepUpEnabled: updatedSettings.StepUpEnabled,
|
||||
AuditLogRetentionDays: updatedSettings.AuditLogRetentionDays,
|
||||
LoginAgreementEnabled: updatedSettings.LoginAgreementEnabled,
|
||||
LoginAgreementMode: updatedSettings.LoginAgreementMode,
|
||||
|
||||
@@ -33,6 +33,7 @@ type UserHandler struct {
|
||||
billingCache service.BillingCache // T17/T18 缓存失效(PUT/POST 路径)
|
||||
totpService *service.TotpService // 角色提升为管理员的 step-up 门控
|
||||
userService *service.UserService
|
||||
settingService *service.SettingService // step-up 功能开关
|
||||
}
|
||||
|
||||
// NewUserHandler creates a new admin user handler
|
||||
@@ -43,6 +44,7 @@ func NewUserHandler(
|
||||
billingCache service.BillingCache,
|
||||
totpService *service.TotpService,
|
||||
userService *service.UserService,
|
||||
settingService *service.SettingService,
|
||||
) *UserHandler {
|
||||
return &UserHandler{
|
||||
adminService: adminService,
|
||||
@@ -51,6 +53,7 @@ func NewUserHandler(
|
||||
billingCache: billingCache,
|
||||
totpService: totpService,
|
||||
userService: userService,
|
||||
settingService: settingService,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -275,7 +278,7 @@ func (h *UserHandler) Create(c *gin.Context) {
|
||||
|
||||
// 创建管理员账号属权限敏感操作:需最近完成 step-up 2FA 验证。
|
||||
if req.Role == service.RoleAdmin {
|
||||
if !middleware.EnforceStepUp(c, h.totpService, h.userService) {
|
||||
if !middleware.EnforceStepUp(c, h.totpService, h.userService, h.settingService) {
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -331,7 +334,7 @@ func (h *UserHandler) Update(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
if target.Role != service.RoleAdmin {
|
||||
if !middleware.EnforceStepUp(c, h.totpService, h.userService) {
|
||||
if !middleware.EnforceStepUp(c, h.totpService, h.userService, h.settingService) {
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -606,6 +609,73 @@ func (h *UserHandler) BatchUpdateConcurrency(c *gin.Context) {
|
||||
response.Success(c, gin.H{"affected": affected})
|
||||
}
|
||||
|
||||
// BatchUpdateLimits overwrites concurrency and/or RPM limits for multiple users.
|
||||
// POST /api/v1/admin/users/batch-limits
|
||||
type BatchUpdateLimitsRequest struct {
|
||||
UserIDs []int64 `json:"user_ids"`
|
||||
All bool `json:"all"`
|
||||
Concurrency *int `json:"concurrency" binding:"omitempty,min=0"`
|
||||
RPMLimit *int `json:"rpm_limit" binding:"omitempty,min=0"`
|
||||
}
|
||||
|
||||
func (h *UserHandler) BatchUpdateLimits(c *gin.Context) {
|
||||
var req BatchUpdateLimitsRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
if req.Concurrency == nil && req.RPMLimit == nil {
|
||||
response.BadRequest(c, "at least one of concurrency or rpm_limit is required")
|
||||
return
|
||||
}
|
||||
if !req.All && len(req.UserIDs) == 0 {
|
||||
response.BadRequest(c, "user_ids is required unless all=true")
|
||||
return
|
||||
}
|
||||
if !req.All && len(req.UserIDs) > 500 {
|
||||
response.BadRequest(c, "user_ids cannot exceed 500")
|
||||
return
|
||||
}
|
||||
|
||||
userIDs := req.UserIDs
|
||||
if req.All {
|
||||
userIDs = nil
|
||||
page := 1
|
||||
const pageSize = 500
|
||||
for {
|
||||
users, _, err := h.adminService.ListUsers(c.Request.Context(), page, pageSize, service.UserListFilters{}, "id", "asc")
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
for _, user := range users {
|
||||
userIDs = append(userIDs, user.ID)
|
||||
}
|
||||
if len(users) < pageSize {
|
||||
break
|
||||
}
|
||||
page++
|
||||
}
|
||||
}
|
||||
|
||||
if len(userIDs) == 0 {
|
||||
response.Success(c, gin.H{"affected": 0})
|
||||
return
|
||||
}
|
||||
|
||||
affected, err := h.adminService.BatchUpdateLimits(
|
||||
c.Request.Context(),
|
||||
userIDs,
|
||||
req.Concurrency,
|
||||
req.RPMLimit,
|
||||
)
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
response.Success(c, gin.H{"affected": affected})
|
||||
}
|
||||
|
||||
// GetUserPlatformQuotas GET /admin/users/:id/platform-quotas
|
||||
// admin 视角:D14 lazy 归零 + 暴露 *_window_start 调试字段
|
||||
func (h *UserHandler) GetUserPlatformQuotas(c *gin.Context) {
|
||||
|
||||
@@ -35,7 +35,7 @@ func TestUserHandlerListIncludesActivityFieldsAndSortParams(t *testing.T) {
|
||||
UpdatedAt: lastLoginAt,
|
||||
},
|
||||
}
|
||||
handler := NewUserHandler(adminSvc, nil, nil, nil, nil, nil)
|
||||
handler := NewUserHandler(adminSvc, nil, nil, nil, nil, nil, nil)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
@@ -89,7 +89,7 @@ func TestUserHandlerGetByIDIncludesActivityFields(t *testing.T) {
|
||||
UpdatedAt: lastLoginAt,
|
||||
},
|
||||
}
|
||||
handler := NewUserHandler(adminSvc, nil, nil, nil, nil, nil)
|
||||
handler := NewUserHandler(adminSvc, nil, nil, nil, nil, nil, nil)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type batchLimitsAdminServiceStub struct {
|
||||
*stubAdminService
|
||||
calls []batchLimitsAdminServiceCall
|
||||
}
|
||||
|
||||
type batchLimitsAdminServiceCall struct {
|
||||
userIDs []int64
|
||||
concurrency *int
|
||||
rpmLimit *int
|
||||
}
|
||||
|
||||
func cloneIntPointer(value *int) *int {
|
||||
if value == nil {
|
||||
return nil
|
||||
}
|
||||
cloned := *value
|
||||
return &cloned
|
||||
}
|
||||
|
||||
func (s *batchLimitsAdminServiceStub) BatchUpdateLimits(_ context.Context, userIDs []int64, concurrency, rpmLimit *int) (int, error) {
|
||||
s.calls = append(s.calls, batchLimitsAdminServiceCall{
|
||||
userIDs: append([]int64(nil), userIDs...),
|
||||
concurrency: cloneIntPointer(concurrency),
|
||||
rpmLimit: cloneIntPointer(rpmLimit),
|
||||
})
|
||||
return len(userIDs), nil
|
||||
}
|
||||
|
||||
func setupBatchLimitsRouter(serviceStub service.AdminService) *gin.Engine {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
handler := NewUserHandler(serviceStub, nil, nil, nil, nil, nil, nil)
|
||||
router.POST("/api/v1/admin/users/batch-limits", handler.BatchUpdateLimits)
|
||||
return router
|
||||
}
|
||||
|
||||
func postBatchLimits(t *testing.T, router *gin.Engine, body []byte) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(
|
||||
http.MethodPost,
|
||||
"/api/v1/admin/users/batch-limits",
|
||||
bytes.NewReader(body),
|
||||
)
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
router.ServeHTTP(recorder, request)
|
||||
return recorder
|
||||
}
|
||||
|
||||
func TestUserHandlerBatchUpdateLimitsAcceptsPartialAndZeroValues(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
expectedConcurrency *int
|
||||
expectedRPMLimit *int
|
||||
}{
|
||||
{name: "concurrency only", body: `{"user_ids":[1,2],"concurrency":10}`, expectedConcurrency: pointerTo(10)},
|
||||
{name: "both limits", body: `{"user_ids":[1,2],"concurrency":8,"rpm_limit":60}`, expectedConcurrency: pointerTo(8), expectedRPMLimit: pointerTo(60)},
|
||||
{name: "explicit zero", body: `{"user_ids":[1,2],"concurrency":0,"rpm_limit":0}`, expectedConcurrency: pointerTo(0), expectedRPMLimit: pointerTo(0)},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
serviceStub := &batchLimitsAdminServiceStub{stubAdminService: newStubAdminService()}
|
||||
recorder := postBatchLimits(t, setupBatchLimitsRouter(serviceStub), []byte(test.body))
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
require.Len(t, serviceStub.calls, 1)
|
||||
require.Equal(t, []int64{1, 2}, serviceStub.calls[0].userIDs)
|
||||
require.Equal(t, test.expectedConcurrency, serviceStub.calls[0].concurrency)
|
||||
require.Equal(t, test.expectedRPMLimit, serviceStub.calls[0].rpmLimit)
|
||||
|
||||
var response struct {
|
||||
Data struct {
|
||||
Affected int `json:"affected"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
require.Equal(t, 2, response.Data.Affected)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUserHandlerBatchUpdateLimitsRejectsInvalidRequests(t *testing.T) {
|
||||
tooManyIDs := make([]int64, 501)
|
||||
for index := range tooManyIDs {
|
||||
tooManyIDs[index] = int64(index + 1)
|
||||
}
|
||||
tooManyBody, err := json.Marshal(map[string]any{"user_ids": tooManyIDs, "rpm_limit": 10})
|
||||
require.NoError(t, err)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
body []byte
|
||||
}{
|
||||
{name: "no limits", body: []byte(`{"user_ids":[1]}`)},
|
||||
{name: "invalid json", body: []byte(`{"user_ids":`)},
|
||||
{name: "missing user ids", body: []byte(`{"rpm_limit":10}`)},
|
||||
{name: "more than 500 ids", body: tooManyBody},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
serviceStub := &batchLimitsAdminServiceStub{stubAdminService: newStubAdminService()}
|
||||
recorder := postBatchLimits(t, setupBatchLimitsRouter(serviceStub), test.body)
|
||||
|
||||
require.Equal(t, http.StatusBadRequest, recorder.Code)
|
||||
require.Empty(t, serviceStub.calls)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUserHandlerBatchUpdateLimitsAllUsesEveryListedUser(t *testing.T) {
|
||||
base := newStubAdminService()
|
||||
base.users = []service.User{{ID: 11}, {ID: 12}, {ID: 13}}
|
||||
serviceStub := &batchLimitsAdminServiceStub{stubAdminService: base}
|
||||
recorder := postBatchLimits(
|
||||
t,
|
||||
setupBatchLimitsRouter(serviceStub),
|
||||
[]byte(`{"all":true,"user_ids":[999],"rpm_limit":0}`),
|
||||
)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
require.Len(t, serviceStub.calls, 1)
|
||||
require.Equal(t, []int64{11, 12, 13}, serviceStub.calls[0].userIDs)
|
||||
require.Equal(t, 1, base.lastListUsers.calls)
|
||||
}
|
||||
|
||||
func pointerTo(value int) *int {
|
||||
return &value
|
||||
}
|
||||
@@ -26,7 +26,7 @@ func (s *getByIDAdminStub) GetUserIncludeDeleted(_ context.Context, id int64) (*
|
||||
func setupGetByIDRouter(svc service.AdminService) *gin.Engine {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
h := NewUserHandler(svc, nil, nil, nil, nil, nil)
|
||||
h := NewUserHandler(svc, nil, nil, nil, nil, nil, nil)
|
||||
r.GET("/admin/users/:id", h.GetByID)
|
||||
return r
|
||||
}
|
||||
|
||||
@@ -39,7 +39,7 @@ func TestAdminUserList_ParsesAPIKeyGroupID(t *testing.T) {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
stub := &listUsersFilterStub{AdminService: newStubAdminService()}
|
||||
r := gin.New()
|
||||
h := NewUserHandler(stub, nil, nil, nil, nil, nil)
|
||||
h := NewUserHandler(stub, nil, nil, nil, nil, nil, nil)
|
||||
r.GET("/admin/users", h.List)
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
@@ -29,7 +29,7 @@ func setupRoleStepUpRouter(t *testing.T) (*gin.Engine, *stubAdminService) {
|
||||
Status: service.StatusActive,
|
||||
})
|
||||
|
||||
h := NewUserHandler(adminSvc, nil, nil, nil, nil, nil)
|
||||
h := NewUserHandler(adminSvc, nil, nil, nil, nil, nil, nil)
|
||||
router.POST("/api/v1/admin/users", h.Create)
|
||||
router.PUT("/api/v1/admin/users/:id", h.Update)
|
||||
return router, adminSvc
|
||||
|
||||
@@ -3168,6 +3168,9 @@ func (r *oauthPendingFlowUserRepo) BatchSetConcurrency(context.Context, []int64,
|
||||
func (r *oauthPendingFlowUserRepo) BatchAddConcurrency(context.Context, []int64, int) (int, error) {
|
||||
panic("unexpected BatchAddConcurrency call")
|
||||
}
|
||||
func (r *oauthPendingFlowUserRepo) BatchUpdateLimits(context.Context, []int64, *int, *int) (int, error) {
|
||||
panic("unexpected BatchUpdateLimits call")
|
||||
}
|
||||
|
||||
func (r *oauthPendingFlowUserRepo) GetLatestUsedAtByUserIDs(context.Context, []int64) (map[int64]*time.Time, error) {
|
||||
return map[int64]*time.Time{}, nil
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
@@ -20,6 +21,7 @@ type BatchImageHandler struct {
|
||||
service *service.BatchImagePublicService
|
||||
download *service.BatchImageDownloadService
|
||||
cleanup *service.BatchImageCleanupService
|
||||
openAI *OpenAIGatewayHandler
|
||||
}
|
||||
|
||||
func NewBatchImageHandler(service *service.BatchImagePublicService, download *service.BatchImageDownloadService, cleanup *service.BatchImageCleanupService) *BatchImageHandler {
|
||||
@@ -37,6 +39,9 @@ func (h *BatchImageHandler) Submit(c *gin.Context) {
|
||||
batchImageError(c, infraerrors.New(http.StatusUnauthorized, "API_KEY_REQUIRED", "API key is required"))
|
||||
return
|
||||
}
|
||||
if !h.checkSecurityAuditBeforeSubmit(c, &req) {
|
||||
return
|
||||
}
|
||||
got, err := h.service.Submit(c.Request.Context(), owner, req, c.GetHeader("Idempotency-Key"))
|
||||
if err != nil {
|
||||
batchImageError(c, err)
|
||||
@@ -45,6 +50,44 @@ func (h *BatchImageHandler) Submit(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, got)
|
||||
}
|
||||
|
||||
func (h *BatchImageHandler) checkSecurityAuditBeforeSubmit(c *gin.Context, req *service.BatchImageSubmitRequest) bool {
|
||||
if h == nil || h.openAI == nil || req == nil {
|
||||
return true
|
||||
}
|
||||
apiKey, ok := middleware.GetAPIKeyFromContext(c)
|
||||
if !ok || apiKey == nil {
|
||||
batchImageError(c, infraerrors.New(http.StatusUnauthorized, "API_KEY_REQUIRED", "API key is required"))
|
||||
return false
|
||||
}
|
||||
subject, ok := middleware.GetAuthSubjectFromContext(c)
|
||||
if !ok {
|
||||
batchImageError(c, infraerrors.New(http.StatusInternalServerError, "USER_CONTEXT_REQUIRED", "User context not found"))
|
||||
return false
|
||||
}
|
||||
items := make([]map[string]string, 0, len(req.Items))
|
||||
for _, item := range req.Items {
|
||||
if prompt := strings.TrimSpace(item.Prompt); prompt != "" {
|
||||
items = append(items, map[string]string{"prompt": prompt})
|
||||
}
|
||||
}
|
||||
if len(items) == 0 {
|
||||
return true
|
||||
}
|
||||
body, err := json.Marshal(map[string]any{"request": map[string]any{"items": items}})
|
||||
if err != nil {
|
||||
batchImageError(c, infraerrors.New(http.StatusBadRequest, "INVALID_BATCH_PROMPT", "batch prompts are invalid"))
|
||||
return false
|
||||
}
|
||||
reqLog := requestLogger(c, "handler.batch_image.security_audit",
|
||||
zap.Int64("user_id", subject.UserID), zap.Int64("api_key_id", apiKey.ID), zap.String("model", req.Model))
|
||||
decision := h.openAI.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, req.Model, body)
|
||||
if decision != nil && !decision.AllowNextStage {
|
||||
h.openAI.openAISecurityAuditError(c, decision)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (h *BatchImageHandler) Get(c *gin.Context) {
|
||||
owner, ok := batchImageOwnerFromContext(c)
|
||||
if !ok {
|
||||
|
||||
@@ -12,13 +12,6 @@ import (
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
func (h *GatewayHandler) checkContentModeration(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol string, model string, body []byte) *service.ContentModerationDecision {
|
||||
if h == nil || h.contentModerationService == nil {
|
||||
return nil
|
||||
}
|
||||
return runContentModeration(c, reqLog, h.contentModerationService, apiKey, subject, protocol, model, body)
|
||||
}
|
||||
|
||||
func contentModerationStatus(decision *service.ContentModerationDecision) int {
|
||||
if decision == nil || decision.StatusCode < 400 || decision.StatusCode > 599 {
|
||||
return http.StatusForbidden
|
||||
@@ -30,13 +23,6 @@ func contentModerationErrorCode(decision *service.ContentModerationDecision) str
|
||||
return "content_policy_violation"
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) checkContentModeration(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol string, model string, body []byte) *service.ContentModerationDecision {
|
||||
if h == nil || h.contentModerationService == nil {
|
||||
return nil
|
||||
}
|
||||
return runContentModeration(c, reqLog, h.contentModerationService, apiKey, subject, protocol, model, body)
|
||||
}
|
||||
|
||||
func runContentModeration(c *gin.Context, reqLog *zap.Logger, svc *service.ContentModerationService, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol string, model string, body []byte) *service.ContentModerationDecision {
|
||||
if svc == nil || c == nil || c.Request == nil {
|
||||
return nil
|
||||
|
||||
@@ -37,6 +37,7 @@ type SystemSettings struct {
|
||||
TotpEnabled bool `json:"totp_enabled"` // TOTP 双因素认证
|
||||
TotpEncryptionKeyConfigured bool `json:"totp_encryption_key_configured"` // TOTP 加密密钥是否已配置
|
||||
SessionBindingEnabled bool `json:"session_binding_enabled"` // 会话 IP/UA 绑定
|
||||
StepUpEnabled bool `json:"step_up_enabled"` // 敏感操作 step-up 2FA
|
||||
AuditLogRetentionDays int `json:"audit_log_retention_days"` // 审计日志保留天数
|
||||
LoginAgreementEnabled bool `json:"login_agreement_enabled"`
|
||||
LoginAgreementMode string `json:"login_agreement_mode"`
|
||||
|
||||
@@ -25,6 +25,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
||||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
@@ -49,6 +50,7 @@ type GatewayHandler struct {
|
||||
usageRecordWorkerPool *service.UsageRecordWorkerPool
|
||||
errorPassthroughService *service.ErrorPassthroughService
|
||||
contentModerationService *service.ContentModerationService
|
||||
securityAuditCoordinator *securityaudit.Coordinator
|
||||
concurrencyHelper *ConcurrencyHelper
|
||||
userMsgQueueHelper *UserMsgQueueHelper
|
||||
maxAccountSwitches int
|
||||
@@ -164,6 +166,7 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||||
return
|
||||
}
|
||||
body = parsedReq.Body.Bytes()
|
||||
reqModel := parsedReq.Model
|
||||
reqStream := parsedReq.Stream
|
||||
reqLog = reqLog.With(zap.String("model", reqModel), zap.Bool("stream", reqStream))
|
||||
@@ -199,8 +202,8 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolAnthropicMessages, reqModel, body); decision != nil && decision.Blocked {
|
||||
h.errorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message)
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolAnthropicMessages, reqModel, body); decision != nil && !decision.AllowNextStage {
|
||||
h.anthropicSecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1020,7 +1023,7 @@ func (h *GatewayHandler) Models(c *gin.Context) {
|
||||
}
|
||||
|
||||
if len(availableModels) > 0 {
|
||||
writeModelsList(c, availableModels)
|
||||
writeModelsList(c, platform, availableModels)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1040,6 +1043,10 @@ func (h *GatewayHandler) Models(c *gin.Context) {
|
||||
})
|
||||
return
|
||||
}
|
||||
if platform == service.PlatformGrok {
|
||||
writeGrokModelsList(c, xai.DefaultModelIDs())
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"object": "list",
|
||||
@@ -1047,7 +1054,11 @@ func (h *GatewayHandler) Models(c *gin.Context) {
|
||||
})
|
||||
}
|
||||
|
||||
func writeModelsList(c *gin.Context, modelIDs []string) {
|
||||
func writeModelsList(c *gin.Context, platform string, modelIDs []string) {
|
||||
if platform == service.PlatformGrok {
|
||||
writeGrokModelsList(c, modelIDs)
|
||||
return
|
||||
}
|
||||
models := make([]claude.Model, 0, len(modelIDs))
|
||||
for _, modelID := range modelIDs {
|
||||
models = append(models, claude.Model{
|
||||
@@ -1068,7 +1079,66 @@ func writeCustomModelsList(c *gin.Context, platform string, modelIDs []string) {
|
||||
writeOpenAIModelsList(c, modelIDs)
|
||||
return
|
||||
}
|
||||
writeModelsList(c, modelIDs)
|
||||
writeModelsList(c, platform, modelIDs)
|
||||
}
|
||||
|
||||
type grokReasoningEffortOption struct {
|
||||
Value string `json:"value"`
|
||||
Label string `json:"label"`
|
||||
Default bool `json:"default,omitempty"`
|
||||
}
|
||||
|
||||
type grokModelListItem struct {
|
||||
xai.Model
|
||||
SupportsReasoningEffort bool `json:"supportsReasoningEffort,omitempty"`
|
||||
ReasoningEffort string `json:"reasoningEffort,omitempty"`
|
||||
ReasoningEfforts []grokReasoningEffortOption `json:"reasoningEfforts,omitempty"`
|
||||
}
|
||||
|
||||
func writeGrokModelsList(c *gin.Context, modelIDs []string) {
|
||||
defaults := xai.DefaultModels()
|
||||
defaultsByID := make(map[string]xai.Model, len(defaults))
|
||||
for _, model := range defaults {
|
||||
defaultsByID[model.ID] = model
|
||||
}
|
||||
|
||||
models := make([]grokModelListItem, 0, len(modelIDs))
|
||||
for _, modelID := range modelIDs {
|
||||
model, ok := defaultsByID[modelID]
|
||||
if !ok {
|
||||
model = xai.Model{
|
||||
ID: modelID,
|
||||
Object: "model",
|
||||
OwnedBy: "xai",
|
||||
DisplayName: modelID,
|
||||
}
|
||||
}
|
||||
item := grokModelListItem{Model: model}
|
||||
if grokModelSupportsConfigurableReasoning(modelID) {
|
||||
item.SupportsReasoningEffort = true
|
||||
item.ReasoningEffort = "high"
|
||||
item.ReasoningEfforts = []grokReasoningEffortOption{
|
||||
{Value: "low", Label: "Low"},
|
||||
{Value: "medium", Label: "Medium"},
|
||||
{Value: "high", Label: "High", Default: true},
|
||||
}
|
||||
}
|
||||
models = append(models, item)
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"object": "list",
|
||||
"data": models,
|
||||
})
|
||||
}
|
||||
|
||||
func grokModelSupportsConfigurableReasoning(modelID string) bool {
|
||||
switch strings.ToLower(strings.TrimSpace(modelID)) {
|
||||
case "grok-4.5", "grok-4.5-latest", "grok", "grok-latest", "grok-build", "grok-build-latest", "grok-build-0.1":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func writeOpenAIModelsList(c *gin.Context, modelIDs []string) {
|
||||
@@ -1813,6 +1883,7 @@ func (h *GatewayHandler) CountTokens(c *gin.Context) {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||||
return
|
||||
}
|
||||
body = parsedReq.Body.Bytes()
|
||||
// count_tokens 走 messages 严格校验时,复用已解析请求,避免二次反序列化。
|
||||
SetClaudeCodeClientContext(c, body, parsedReq)
|
||||
reqLog = reqLog.With(zap.String("model", parsedReq.Model), zap.Bool("stream", parsedReq.Stream))
|
||||
|
||||
@@ -99,8 +99,8 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIChat, reqModel, body); decision != nil && decision.Blocked {
|
||||
h.chatCompletionsErrorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message)
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIChat, reqModel, body); decision != nil && !decision.AllowNextStage {
|
||||
h.openAISecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
|
||||
|
||||
@@ -104,8 +104,8 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIResponses, reqModel, body); decision != nil && decision.Blocked {
|
||||
h.responsesErrorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message)
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIResponses, reqModel, body); decision != nil && !decision.AllowNextStage {
|
||||
h.responsesSecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
|
||||
|
||||
@@ -25,11 +25,20 @@ type gatewayModelsResponseForTest struct {
|
||||
}
|
||||
|
||||
type gatewayModelItemForTest struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
OwnedBy string `json:"owned_by"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
OwnedBy string `json:"owned_by"`
|
||||
CreatedAt string `json:"created_at"`
|
||||
SupportsReasoningEffort bool `json:"supportsReasoningEffort"`
|
||||
ReasoningEffort string `json:"reasoningEffort"`
|
||||
ReasoningEfforts []gatewayReasoningEffortOptionForTest `json:"reasoningEfforts"`
|
||||
}
|
||||
|
||||
type gatewayReasoningEffortOptionForTest struct {
|
||||
Value string `json:"value"`
|
||||
Label string `json:"label"`
|
||||
Default bool `json:"default"`
|
||||
}
|
||||
|
||||
func (s *gatewayModelsAccountRepoStub) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]service.Account, error) {
|
||||
@@ -84,6 +93,50 @@ func TestGatewayModels_GeminiGroupFallsBackToGeminiModels(t *testing.T) {
|
||||
require.NotContains(t, modelIDsForTest(got.Data), "claude-sonnet-4-6")
|
||||
}
|
||||
|
||||
func TestGatewayModels_Grok45AdvertisesReasoningEffortForGrokBuild(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
groupID := int64(4409)
|
||||
h := newGatewayModelsHandlerForTest(
|
||||
&gatewayModelsAccountRepoStub{
|
||||
byGroup: map[int64][]service.Account{
|
||||
groupID: {
|
||||
{
|
||||
ID: 1,
|
||||
Platform: service.PlatformGrok,
|
||||
Credentials: map[string]any{
|
||||
"model_mapping": map[string]any{"grok-4.5": "grok-4.5"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformGrok},
|
||||
})
|
||||
|
||||
h.Models(c)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var got gatewayModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
require.Len(t, got.Data, 1)
|
||||
model := got.Data[0]
|
||||
require.Equal(t, "grok-4.5", model.ID)
|
||||
require.True(t, model.SupportsReasoningEffort)
|
||||
require.Equal(t, "high", model.ReasoningEffort)
|
||||
require.Equal(t, []gatewayReasoningEffortOptionForTest{
|
||||
{Value: "low", Label: "Low"},
|
||||
{Value: "medium", Label: "Medium"},
|
||||
{Value: "high", Label: "High", Default: true},
|
||||
}, model.ReasoningEfforts)
|
||||
}
|
||||
|
||||
func TestGatewayModels_GeminiGroupFiltersMappedModelsByPlatform(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
|
||||
@@ -187,8 +187,8 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
|
||||
setOpsRequestContext(c, modelName, stream)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(stream, false)))
|
||||
|
||||
if decision := h.checkContentModeration(c, reqLog, apiKey, authSubject, service.ContentModerationProtocolGemini, modelName, body); decision != nil && decision.Blocked {
|
||||
googleError(c, contentModerationStatus(decision), decision.Message)
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, authSubject, service.ContentModerationProtocolGemini, modelName, body); decision != nil && !decision.AllowNextStage {
|
||||
googleSecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
|
||||
|
||||
@@ -46,6 +46,11 @@ func (h *OpenAIGatewayHandler) GrokVideoStatus(c *gin.Context) {
|
||||
h.handleGrokMedia(c, service.GrokMediaEndpointVideoStatus, c.Param("request_id"))
|
||||
}
|
||||
|
||||
// GrokVideoContent proxies downloadable video content through the task's upstream account.
|
||||
func (h *OpenAIGatewayHandler) GrokVideoContent(c *gin.Context) {
|
||||
h.handleGrokMedia(c, service.GrokMediaEndpointVideoContent, c.Param("request_id"))
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.GrokMediaEndpoint, requestID string) {
|
||||
streamStarted := false
|
||||
defer h.recoverResponsesPanic(c, &streamStarted)
|
||||
@@ -95,11 +100,12 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
contentType := c.GetHeader("Content-Type")
|
||||
requestInfo := service.ParseGrokMediaRequest(contentType, body)
|
||||
requestModel := requestInfo.Model
|
||||
routingModel := service.NormalizeGrokMediaModelForEndpoint(endpoint, requestModel, requestInfo.HasInputImage())
|
||||
if endpoint.IsGenerationRequest() && strings.TrimSpace(requestModel) == "" {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||||
return
|
||||
}
|
||||
if endpoint == service.GrokMediaEndpointVideoStatus && strings.TrimSpace(requestID) == "" {
|
||||
if endpoint.IsVideoLookupRequest() && strings.TrimSpace(requestID) == "" {
|
||||
h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "request_id is required")
|
||||
return
|
||||
}
|
||||
@@ -114,9 +120,9 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
return
|
||||
}
|
||||
if moderationBody := requestInfo.ModerationBody(); len(moderationBody) > 0 {
|
||||
decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, requestModel, moderationBody)
|
||||
if decision != nil && decision.Blocked {
|
||||
h.errorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message)
|
||||
decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, requestModel, moderationBody)
|
||||
if decision != nil && !decision.AllowNextStage {
|
||||
h.openAISecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -159,20 +165,31 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
sessionSeed = []byte(requestID)
|
||||
}
|
||||
sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, sessionSeed)
|
||||
if endpoint == service.GrokMediaEndpointVideoStatus {
|
||||
sessionHash = service.GrokMediaVideoRequestSessionHash(requestID)
|
||||
boundLookupAccountID := int64(0)
|
||||
if endpoint.IsVideoLookupRequest() {
|
||||
sessionHash = service.GrokMediaVideoRequestSessionHash(requestID, subject.UserID, apiKey.ID)
|
||||
boundLookupAccountID, err = h.gatewayService.ResolveGrokMediaVideoRequestAccount(
|
||||
c.Request.Context(), apiKey.GroupID, requestID, subject.UserID, apiKey.ID,
|
||||
)
|
||||
if err != nil || boundLookupAccountID <= 0 {
|
||||
reqLog.Info("grok_media.video_lookup_owner_binding_missing", zap.Error(err))
|
||||
h.errorResponse(c, http.StatusNotFound, "not_found_error", "Video request not found")
|
||||
return
|
||||
}
|
||||
}
|
||||
requestCtx := c.Request.Context()
|
||||
failedAccountIDs := make(map[int64]struct{})
|
||||
sameAccountRetryCount := make(map[int64]int)
|
||||
var lastFailoverErr *service.UpstreamFailoverError
|
||||
var oauth429FailoverState service.OpenAIOAuth429FailoverState
|
||||
mediaEligibilityRejected := false
|
||||
switchCount := 0
|
||||
maxAccountSwitches := h.maxAccountSwitches
|
||||
if maxAccountSwitches <= 0 {
|
||||
maxAccountSwitches = 3
|
||||
}
|
||||
routingStart := time.Now()
|
||||
requiredCapability := grokMediaRequiredCapability(endpoint)
|
||||
|
||||
for {
|
||||
if failoverClientGone(c) {
|
||||
@@ -183,10 +200,10 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
apiKey.GroupID,
|
||||
"",
|
||||
sessionHash,
|
||||
requestModel,
|
||||
routingModel,
|
||||
failedAccountIDs,
|
||||
service.OpenAIUpstreamTransportHTTPSSE,
|
||||
"",
|
||||
requiredCapability,
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
@@ -201,8 +218,14 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
zap.Error(err),
|
||||
zap.Int("excluded_account_count", len(failedAccountIDs)),
|
||||
)
|
||||
if endpoint.IsGenerationRequest() && errors.Is(err, service.ErrNoAvailableAccounts) &&
|
||||
(len(failedAccountIDs) == 0 || (mediaEligibilityRejected && lastFailoverErr == nil)) {
|
||||
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts")
|
||||
return
|
||||
}
|
||||
if len(failedAccountIDs) == 0 {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, requestModel, service.PlatformGrok)
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, routingModel, service.PlatformGrok)
|
||||
if !cls.ModelNotFound {
|
||||
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
|
||||
}
|
||||
@@ -217,13 +240,26 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
return
|
||||
}
|
||||
if selection == nil || selection.Account == nil {
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, requestModel, service.PlatformGrok)
|
||||
if endpoint.IsGenerationRequest() {
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts")
|
||||
return
|
||||
}
|
||||
cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, routingModel, service.PlatformGrok)
|
||||
if !cls.ModelNotFound {
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
}
|
||||
h.errorResponse(c, cls.Status, cls.ErrType, cls.Message)
|
||||
return
|
||||
}
|
||||
if boundLookupAccountID > 0 && selection.Account.ID != boundLookupAccountID {
|
||||
reqLog.Warn("grok_media.video_lookup_bound_account_unavailable",
|
||||
zap.Int64("bound_account_id", boundLookupAccountID),
|
||||
zap.Int64("selected_account_id", selection.Account.ID),
|
||||
)
|
||||
h.errorResponse(c, http.StatusNotFound, "not_found_error", "Video request not found")
|
||||
return
|
||||
}
|
||||
|
||||
reqLog.Debug("grok_media.account_schedule_decision",
|
||||
zap.String("layer", scheduleDecision.Layer),
|
||||
@@ -235,6 +271,25 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
)
|
||||
|
||||
account := selection.Account
|
||||
if endpoint.IsGenerationRequest() {
|
||||
eligible, eligibilityReason, eligibilityErr := h.ensureGrokMediaAccountEligibility(requestCtx, account)
|
||||
if !eligible {
|
||||
mediaEligibilityRejected = true
|
||||
failedAccountIDs[account.ID] = struct{}{}
|
||||
reqLog.Warn("grok_media.account_eligibility_rejected",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.String("reason", eligibilityReason),
|
||||
zap.Bool("probe_failed", eligibilityErr != nil),
|
||||
)
|
||||
if switchCount >= maxAccountSwitches {
|
||||
markOpsRoutingCapacityLimited(c)
|
||||
h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts")
|
||||
return
|
||||
}
|
||||
switchCount++
|
||||
continue
|
||||
}
|
||||
}
|
||||
sessionHash = ensureOpenAIPoolModeSessionHash(sessionHash, account)
|
||||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||||
|
||||
@@ -274,7 +329,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
return
|
||||
}
|
||||
if failoverErr.ShouldReportAccountScheduleFailure() {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil)
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, grokMediaScheduleModel(account, routingModel, nil), false, nil)
|
||||
}
|
||||
if c.Writer.Size() != writerSizeBeforeForward {
|
||||
h.handleFailoverExhausted(c, failoverErr, true)
|
||||
@@ -284,6 +339,10 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
h.handleFailoverExhausted(c, failoverErr, false)
|
||||
return
|
||||
}
|
||||
if endpoint.IsVideoLookupRequest() {
|
||||
h.handleFailoverExhausted(c, failoverErr, false)
|
||||
return
|
||||
}
|
||||
if failoverErr.RetryableOnSameAccount {
|
||||
retryLimit := account.GetPoolModeRetryCount()
|
||||
if sameAccountRetryCount[account.ID] < retryLimit {
|
||||
@@ -322,8 +381,8 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
)
|
||||
continue
|
||||
}
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), false, nil)
|
||||
if c.Writer.Size() == writerSizeBeforeForward {
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, grokMediaScheduleModel(account, routingModel, nil), false, nil)
|
||||
if !service.IsResponseCommitted(c) && c.Writer.Size() == writerSizeBeforeForward {
|
||||
h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed")
|
||||
}
|
||||
reqLog.Warn("grok_media.forward_failed",
|
||||
@@ -333,9 +392,11 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
return
|
||||
}
|
||||
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(requestModel), true, nil)
|
||||
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, grokMediaScheduleModel(account, routingModel, result), true, nil)
|
||||
if endpoint.IsGenerationRequest() && strings.TrimSpace(result.ResponseID) != "" {
|
||||
if err := h.gatewayService.BindGrokMediaVideoRequestAccount(requestCtx, apiKey.GroupID, result.ResponseID, account.ID); err != nil {
|
||||
if err := h.gatewayService.BindGrokMediaVideoRequestAccount(
|
||||
requestCtx, apiKey.GroupID, result.ResponseID, subject.UserID, apiKey.ID, account.ID,
|
||||
); err != nil {
|
||||
reqLog.Warn("grok_media.bind_video_request_account_failed",
|
||||
zap.Int64("account_id", account.ID),
|
||||
zap.String("request_id", result.ResponseID),
|
||||
@@ -354,6 +415,37 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
|
||||
}
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) ensureGrokMediaAccountEligibility(ctx context.Context, account *service.Account) (bool, string, error) {
|
||||
if account == nil {
|
||||
return false, "missing_account", errors.New("grok media account is required")
|
||||
}
|
||||
eligible, reason := account.GrokMediaGenerationEligibility()
|
||||
if eligible || reason != "billing_unobserved" {
|
||||
return eligible, reason, nil
|
||||
}
|
||||
if h == nil || h.grokMediaEligibilityProber == nil {
|
||||
return false, "billing_probe_unavailable", errors.New("grok media eligibility probe is not configured")
|
||||
}
|
||||
return h.grokMediaEligibilityProber.ProbeMediaEligibility(ctx, account.ID)
|
||||
}
|
||||
|
||||
func grokMediaRequiredCapability(endpoint service.GrokMediaEndpoint) service.OpenAIEndpointCapability {
|
||||
if endpoint.IsGenerationRequest() {
|
||||
return service.OpenAIEndpointCapabilityGrokMediaGeneration
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func grokMediaScheduleModel(account *service.Account, routingModel string, result *service.OpenAIForwardResult) string {
|
||||
if result != nil && strings.TrimSpace(result.UpstreamModel) != "" {
|
||||
return result.UpstreamModel
|
||||
}
|
||||
if account == nil {
|
||||
return strings.TrimSpace(routingModel)
|
||||
}
|
||||
return account.GetMappedModel(routingModel)
|
||||
}
|
||||
|
||||
func shouldRecordGrokMediaUsage(endpoint service.GrokMediaEndpoint, requestModel string) bool {
|
||||
return endpoint.IsGenerationRequest() && strings.TrimSpace(requestModel) != ""
|
||||
}
|
||||
|
||||
@@ -1,12 +1,26 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type grokMediaEligibilityProberStub struct {
|
||||
eligible bool
|
||||
reason string
|
||||
err error
|
||||
calls int
|
||||
}
|
||||
|
||||
func (s *grokMediaEligibilityProberStub) ProbeMediaEligibility(context.Context, int64) (bool, string, error) {
|
||||
s.calls++
|
||||
return s.eligible, s.reason, s.err
|
||||
}
|
||||
|
||||
func TestShouldRecordGrokMediaUsage(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
@@ -38,6 +52,12 @@ func TestShouldRecordGrokMediaUsage(t *testing.T) {
|
||||
model: "",
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "video content skips usage",
|
||||
endpoint: service.GrokMediaEndpointVideoContent,
|
||||
model: "",
|
||||
want: false,
|
||||
},
|
||||
{
|
||||
name: "generation skips usage without model",
|
||||
endpoint: service.GrokMediaEndpointImagesGenerations,
|
||||
@@ -52,3 +72,96 @@ func TestShouldRecordGrokMediaUsage(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokMediaRequiredCapability(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
endpoint service.GrokMediaEndpoint
|
||||
want service.OpenAIEndpointCapability
|
||||
}{
|
||||
{name: "image generation", endpoint: service.GrokMediaEndpointImagesGenerations, want: service.OpenAIEndpointCapabilityGrokMediaGeneration},
|
||||
{name: "image edit", endpoint: service.GrokMediaEndpointImagesEdits, want: service.OpenAIEndpointCapabilityGrokMediaGeneration},
|
||||
{name: "video generation", endpoint: service.GrokMediaEndpointVideosGenerations, want: service.OpenAIEndpointCapabilityGrokMediaGeneration},
|
||||
{name: "video edit", endpoint: service.GrokMediaEndpointVideosEdits, want: service.OpenAIEndpointCapabilityGrokMediaGeneration},
|
||||
{name: "video extension", endpoint: service.GrokMediaEndpointVideosExtensions, want: service.OpenAIEndpointCapabilityGrokMediaGeneration},
|
||||
{name: "video status preserves lookup", endpoint: service.GrokMediaEndpointVideoStatus, want: ""},
|
||||
{name: "video content preserves lookup", endpoint: service.GrokMediaEndpointVideoContent, want: ""},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
require.Equal(t, tt.want, grokMediaRequiredCapability(tt.endpoint))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokMediaScheduleModelUsesNormalizedMappedUpstream(t *testing.T) {
|
||||
account := &service.Account{
|
||||
Platform: service.PlatformGrok,
|
||||
Credentials: map[string]any{
|
||||
"model_mapping": map[string]any{
|
||||
"grok-imagine-video-1.5": "wrong-raw-model",
|
||||
"grok-imagine-video": "mapped-video-model",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
require.Equal(t, "mapped-video-model", grokMediaScheduleModel(account, "grok-imagine-video", nil))
|
||||
require.Equal(t, "actual-upstream-model", grokMediaScheduleModel(account, "grok-imagine-video", &service.OpenAIForwardResult{
|
||||
UpstreamModel: "actual-upstream-model",
|
||||
}))
|
||||
require.Equal(t, "mapped-video-model", grokMediaScheduleModel(account, "grok-imagine-video", &service.OpenAIForwardResult{}))
|
||||
require.Equal(t, "grok-imagine-video", grokMediaScheduleModel(nil, " grok-imagine-video ", nil))
|
||||
}
|
||||
|
||||
func TestEnsureGrokMediaAccountEligibility(t *testing.T) {
|
||||
t.Run("non oauth account does not probe", func(t *testing.T) {
|
||||
prober := &grokMediaEligibilityProberStub{}
|
||||
h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober}
|
||||
account := &service.Account{Platform: service.PlatformGrok, Type: service.AccountTypeAPIKey}
|
||||
|
||||
eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, eligible)
|
||||
require.Equal(t, "non_oauth", reason)
|
||||
require.Zero(t, prober.calls)
|
||||
})
|
||||
|
||||
t.Run("unobserved oauth is probed before forwarding", func(t *testing.T) {
|
||||
prober := &grokMediaEligibilityProberStub{eligible: true, reason: "eligible"}
|
||||
h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober}
|
||||
account := &service.Account{ID: 7, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth}
|
||||
|
||||
eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account)
|
||||
|
||||
require.NoError(t, err)
|
||||
require.True(t, eligible)
|
||||
require.Equal(t, "eligible", reason)
|
||||
require.Equal(t, 1, prober.calls)
|
||||
})
|
||||
|
||||
t.Run("missing prober fails closed", func(t *testing.T) {
|
||||
h := &OpenAIGatewayHandler{}
|
||||
account := &service.Account{ID: 8, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth}
|
||||
|
||||
eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account)
|
||||
|
||||
require.Error(t, err)
|
||||
require.False(t, eligible)
|
||||
require.Equal(t, "billing_probe_unavailable", reason)
|
||||
})
|
||||
|
||||
t.Run("probe failure fails closed", func(t *testing.T) {
|
||||
probeErr := errors.New("probe failed")
|
||||
prober := &grokMediaEligibilityProberStub{reason: "billing_unobserved", err: probeErr}
|
||||
h := &OpenAIGatewayHandler{grokMediaEligibilityProber: prober}
|
||||
account := &service.Account{ID: 9, Platform: service.PlatformGrok, Type: service.AccountTypeOAuth}
|
||||
|
||||
eligible, reason, err := h.ensureGrokMediaAccountEligibility(context.Background(), account)
|
||||
|
||||
require.ErrorIs(t, err, probeErr)
|
||||
require.False(t, eligible)
|
||||
require.Equal(t, "billing_unobserved", reason)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ package handler
|
||||
|
||||
import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/admin"
|
||||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||||
)
|
||||
|
||||
// AdminHandlers contains all admin-related HTTP handlers
|
||||
@@ -35,6 +36,7 @@ type AdminHandlers struct {
|
||||
ChannelMonitor *admin.ChannelMonitorHandler
|
||||
ChannelMonitorTemplate *admin.ChannelMonitorRequestTemplateHandler
|
||||
ContentModeration *admin.ContentModerationHandler
|
||||
PromptAudit *securityaudit.PromptAdminHandler
|
||||
Payment *admin.PaymentHandler
|
||||
Affiliate *admin.AffiliateHandler
|
||||
Compliance *admin.ComplianceHandler
|
||||
|
||||
@@ -89,6 +89,9 @@ func (h *AsyncImageHandler) Submit(c *gin.Context) {
|
||||
imageTaskJSONError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
return
|
||||
}
|
||||
if !h.checkSecurityAuditBeforeSubmit(c, apiKey, platform, body) {
|
||||
return
|
||||
}
|
||||
|
||||
taskCtx, recorder, cancel := newAsyncImageContext(c, body, h.tasks.ExecutionTimeout())
|
||||
task, err := h.tasks.Create(c.Request.Context(), service.ImageTaskOwner{UserID: apiKey.UserID, APIKeyID: apiKey.ID})
|
||||
@@ -115,6 +118,42 @@ func (h *AsyncImageHandler) Submit(c *gin.Context) {
|
||||
go h.run(task.ID, platform, taskCtx, recorder, cancel)
|
||||
}
|
||||
|
||||
func (h *AsyncImageHandler) checkSecurityAuditBeforeSubmit(c *gin.Context, apiKey *service.APIKey, platform string, body []byte) bool {
|
||||
if h == nil || h.openAI == nil {
|
||||
return true
|
||||
}
|
||||
subject, ok := middleware2.GetAuthSubjectFromContext(c)
|
||||
if !ok {
|
||||
imageTaskJSONError(c, http.StatusInternalServerError, "api_error", "User context not found")
|
||||
return false
|
||||
}
|
||||
model := ""
|
||||
moderationBody := body
|
||||
if platform == service.PlatformGrok {
|
||||
parsed := service.ParseGrokMediaRequest(c.GetHeader("Content-Type"), body)
|
||||
model, moderationBody = parsed.Model, parsed.ModerationBody()
|
||||
} else if h.openAI.gatewayService != nil {
|
||||
parsed, err := h.openAI.gatewayService.ParseOpenAIImagesRequest(c, body)
|
||||
if err != nil {
|
||||
imageTaskJSONError(c, http.StatusBadRequest, "invalid_request_error", err.Error())
|
||||
return false
|
||||
}
|
||||
model, moderationBody = parsed.Model, parsed.ModerationBody()
|
||||
}
|
||||
if len(moderationBody) == 0 {
|
||||
c.Set(securityAuditCompletedContextKey, true)
|
||||
return true
|
||||
}
|
||||
reqLog := requestLogger(c, "handler.async_image.security_audit",
|
||||
zap.Int64("user_id", subject.UserID), zap.Int64("api_key_id", apiKey.ID), zap.String("model", model))
|
||||
decision := h.openAI.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, model, moderationBody)
|
||||
if decision != nil && !decision.AllowNextStage {
|
||||
h.openAI.openAISecurityAuditError(c, decision)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (h *AsyncImageHandler) Get(c *gin.Context) {
|
||||
if !h.enabled() {
|
||||
imageTaskJSONError(c, http.StatusNotFound, "not_found_error", "async image tasks are not enabled")
|
||||
|
||||
@@ -38,11 +38,12 @@ type noAccountErrorClassification struct {
|
||||
// The classifier intentionally does not consume the original error: the
|
||||
// selection layer never tells us *why* the pool came up empty (rate-limited
|
||||
// vs. unsupported model are both wrapped as ErrNoAvailableAccounts). Instead
|
||||
// we re-check pool composition through DiagnoseModelAvailabilityForPlatform,
|
||||
// which only inspects model_mapping configuration and ignores transient
|
||||
// state. That guarantees a 404 is only returned when no operator action
|
||||
// short of editing the account's model_mapping could make this request
|
||||
// succeed.
|
||||
// we re-check pool composition through DiagnoseModelAvailabilityForPlatform.
|
||||
// Its dedicated database query considers only persistent eligibility
|
||||
// (active status + schedulable setting) and model_mapping, bypassing scheduler
|
||||
// snapshots and transient filters. That guarantees a 404 is only returned
|
||||
// when persistent account/group/model configuration must change before the
|
||||
// request can succeed.
|
||||
//
|
||||
// routingModel is the model name that account selection actually compared
|
||||
// against (i.e. after group-level dispatch mapping). displayModel is the
|
||||
|
||||
@@ -154,6 +154,20 @@ func TestClassifyNoAccountError_HasModelSupport_KeepsRoutingMessageGenerationToC
|
||||
require.False(t, cls.ModelNotFound)
|
||||
}
|
||||
|
||||
func TestClassifyNoAccountError_ModelSupportedOnlyByRateLimitedAccount_Returns503(t *testing.T) {
|
||||
c := newTestGinContextWithRequest()
|
||||
// The diagnoser's configured-state lookup still sees the model-supporting
|
||||
// account even though normal scheduling has excluded it during cooldown.
|
||||
fd := &fakeDiagnoser{resp: service.ModelAvailabilityDiagnosis{HasAccountsInPool: true, HasModelSupport: true}}
|
||||
apiKey := &service.APIKey{GroupID: ptrInt64(7)}
|
||||
|
||||
cls := classifyNoAccountErrorFromGin(c, fd, apiKey, "claude-opus-4-8", "claude-opus-4-8", service.PlatformAnthropic)
|
||||
|
||||
require.Equal(t, http.StatusServiceUnavailable, cls.Status)
|
||||
require.Equal(t, "api_error", cls.ErrType)
|
||||
require.False(t, cls.ModelNotFound, "temporary account cooldown must remain retryable")
|
||||
}
|
||||
|
||||
func TestClassifyNoAccountError_NoAccountsInPool_Stays503(t *testing.T) {
|
||||
c := newTestGinContextWithRequest()
|
||||
fd := &fakeDiagnoser{resp: service.ModelAvailabilityDiagnosis{HasAccountsInPool: false, HasModelSupport: false}}
|
||||
|
||||
@@ -78,6 +78,10 @@ func (h *OpenAIGatewayHandler) AlphaSearch(c *gin.Context) {
|
||||
reqLog = reqLog.With(zap.String("model", requestedModel))
|
||||
setOpsRequestContext(c, requestedModel, false)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeSync))
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, "openai_alpha_search", requestedModel, body); decision != nil && !decision.AllowNextStage {
|
||||
h.openAISecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, requestedModel)
|
||||
forwardBody := openAIModelMappedBody(body, channelMapping.Mapped, channelMapping.MappedModel, h.gatewayService.ReplaceModelInBody)
|
||||
|
||||
@@ -90,8 +90,8 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
setOpsRequestContext(c, reqModel, reqStream)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(reqStream, false)))
|
||||
|
||||
if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIChat, reqModel, body); decision != nil && decision.Blocked {
|
||||
h.errorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message)
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIChat, reqModel, body); decision != nil && !decision.AllowNextStage {
|
||||
h.openAISecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
if h.rejectIfCyberSessionBlocked(c, apiKey, body, reqModel, cyberBlockFormatChat) {
|
||||
@@ -305,7 +305,10 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
|
||||
upstreamErrorAlreadyCommunicated := openAIForwardErrorAlreadyCommunicated(c, writerSizeBeforeForward, err)
|
||||
wroteFallback := false
|
||||
if !upstreamErrorAlreadyCommunicated {
|
||||
wroteFallback = h.ensureForwardErrorResponse(c, streamStarted)
|
||||
wroteFallback = h.ensureOpenAIStreamReadErrorResponse(c, err, streamStarted)
|
||||
if !wroteFallback {
|
||||
wroteFallback = h.ensureForwardErrorResponse(c, streamStarted)
|
||||
}
|
||||
}
|
||||
reqLog.Warn("openai_chat_completions.forward_failed",
|
||||
zap.Int64("account_id", account.ID),
|
||||
|
||||
@@ -49,6 +49,7 @@ type codexModelsFailoverHTTPUpstream struct {
|
||||
accountIDs []int64
|
||||
firstErr error
|
||||
firstStatus int
|
||||
firstBody string
|
||||
statuses map[int64]int
|
||||
}
|
||||
|
||||
@@ -62,6 +63,14 @@ func (u *codexModelsFailoverHTTPUpstream) Do(_ *http.Request, _ string, accountI
|
||||
if u.firstErr != nil {
|
||||
return nil, u.firstErr
|
||||
}
|
||||
if u.firstBody != "" && !hasStatus {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Status: "200 OK",
|
||||
Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(u.firstBody)),
|
||||
}, nil
|
||||
}
|
||||
if !hasStatus {
|
||||
status = u.firstStatus
|
||||
}
|
||||
@@ -150,6 +159,22 @@ func TestCodexModelsFailsOverFromUpstreamTransportError(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestCodexModelsFailsOverFromInvalidManifestEnvelope(t *testing.T) {
|
||||
handler, upstream, groupID := newCodexModelsFailoverTestHandler(http.StatusOK)
|
||||
upstream.firstBody = `{"object":"list","data":[]}`
|
||||
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 TestCodexModelsDoesNotFailOverFromPermanentUpstreamStatus(t *testing.T) {
|
||||
statuses := []int{
|
||||
http.StatusBadRequest,
|
||||
|
||||
@@ -74,6 +74,10 @@ func (h *OpenAIGatewayHandler) Embeddings(c *gin.Context) {
|
||||
reqLog = reqLog.With(zap.String("model", reqModel))
|
||||
setOpsRequestContext(c, reqModel, false)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeSync))
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, "openai_embeddings", reqModel, body); decision != nil && !decision.AllowNextStage {
|
||||
h.openAISecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
|
||||
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, reqModel)
|
||||
|
||||
|
||||
@@ -67,6 +67,7 @@ func (h *OpenAIGatewayHandler) CountTokens(c *gin.Context) {
|
||||
h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to parse request body")
|
||||
return
|
||||
}
|
||||
body = parsedReq.Body.Bytes()
|
||||
if parsedReq.Model == "" {
|
||||
h.anthropicErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||||
return
|
||||
|
||||
@@ -645,7 +645,8 @@ func TestGrokMedia429FailoverIsBounded(t *testing.T) {
|
||||
_, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "first_429")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/openai/v1/videos/request-1", nil)
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/videos/generations", bytes.NewBufferString(`{"model":"grok-imagine-video","prompt":"waves"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
@@ -657,7 +658,8 @@ func TestGrokMedia429FailoverIsBounded(t *testing.T) {
|
||||
_, _, upstream, router, cleanup := newGrokCredentialFailoverHandler(t, "all_429")
|
||||
defer cleanup()
|
||||
recorder := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/openai/v1/videos/request-1", nil)
|
||||
req := httptest.NewRequest(http.MethodPost, "/openai/v1/videos/generations", bytes.NewBufferString(`{"model":"grok-imagine-video","prompt":"waves"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
router.ServeHTTP(recorder, req)
|
||||
|
||||
@@ -678,7 +680,7 @@ func TestGrokOAuthCredentialFailoverAcrossHTTPHandlers(t *testing.T) {
|
||||
{name: "messages", method: http.MethodPost, path: "/openai/v1/messages", body: `{"model":"grok","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`},
|
||||
{name: "chat completions", method: http.MethodPost, path: "/openai/v1/chat/completions", body: `{"model":"grok","messages":[{"role":"user","content":"hello"}],"stream":false}`},
|
||||
{name: "chat completions raw fallback", method: http.MethodPost, path: "/openai/v1/chat/completions", body: `{"model":"grok","messages":[{"role":"user","content":"hello"}],"stop":["END"],"stream":false}`},
|
||||
{name: "grok media", method: http.MethodGet, path: "/openai/v1/videos/request-1"},
|
||||
{name: "grok media", method: http.MethodPost, path: "/openai/v1/videos/generations", body: `{"model":"grok-imagine-video","prompt":"waves"}`},
|
||||
}
|
||||
|
||||
for _, endpoint := range endpoints {
|
||||
@@ -822,6 +824,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
"access_token": "expired", "refresh_token": "revoked-refresh",
|
||||
"expires_at": time.Now().Add(-time.Minute).UTC().Format(time.RFC3339),
|
||||
},
|
||||
Extra: map[string]any{service.GrokMediaEligibleExtraKey: true},
|
||||
},
|
||||
{
|
||||
ID: 802, Name: "healthy", Platform: service.PlatformGrok, Type: service.AccountTypeOAuth,
|
||||
@@ -830,6 +833,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
"access_token": "healthy-access", "refresh_token": "healthy-refresh",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
Extra: map[string]any{service.GrokMediaEligibleExtraKey: true},
|
||||
},
|
||||
}
|
||||
if mode == "postmap_cancel" || mode == "first_429" || mode == "all_429" || mode == "mixed_429_500" || mode == "mixed_500_429" || mode == "oauth_429_apikey_500" {
|
||||
@@ -843,6 +847,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
"access_token": "untried-healthy-access", "refresh_token": "untried-healthy-refresh",
|
||||
"expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
|
||||
},
|
||||
Extra: map[string]any{service.GrokMediaEligibleExtraKey: true},
|
||||
})
|
||||
}
|
||||
if mode == "oauth_429_apikey_500" {
|
||||
@@ -904,7 +909,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
apiKey := &service.APIKey{
|
||||
ID: 902, GroupID: &groupID,
|
||||
User: &service.User{ID: 903, Status: service.StatusActive},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformGrok, Status: service.StatusActive},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformGrok, Status: service.StatusActive, AllowImageGeneration: true},
|
||||
}
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
@@ -916,6 +921,7 @@ func newGrokCredentialFailoverHandler(t *testing.T, mode string) (*OpenAIGateway
|
||||
router.GET("/openai/v1/responses", h.ResponsesWebSocket)
|
||||
router.POST("/openai/v1/messages", h.Messages)
|
||||
router.POST("/openai/v1/chat/completions", h.ChatCompletions)
|
||||
router.POST("/openai/v1/videos/generations", h.GrokVideoGeneration)
|
||||
router.GET("/openai/v1/videos/:request_id", h.GrokVideoStatus)
|
||||
handlerRefresherStarted.Store(router, refresher.started)
|
||||
cleanup := func() {
|
||||
|
||||
@@ -206,6 +206,10 @@ func TestOpsRecoveredCredentialFailoverUsesAccountAuthAttribution(t *testing.T)
|
||||
require.NotContains(t, job.entry.ErrorMessage, "earlier inference failure")
|
||||
require.NotNil(t, job.entry.UpstreamStatusCode)
|
||||
require.Zero(t, *job.entry.UpstreamStatusCode)
|
||||
require.Len(t, job.entry.UpstreamErrors, 2)
|
||||
require.Equal(t, http.StatusForbidden, job.entry.UpstreamErrors[0].UpstreamStatusCode)
|
||||
require.Nil(t, job.entry.UpstreamErrors)
|
||||
require.NotNil(t, job.entry.UpstreamErrorsJSON)
|
||||
events, err := service.ParseOpsUpstreamErrors(*job.entry.UpstreamErrorsJSON)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, events, 2)
|
||||
require.Equal(t, http.StatusForbidden, events[0].UpstreamStatusCode)
|
||||
}
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
@@ -27,17 +28,23 @@ import (
|
||||
|
||||
// OpenAIGatewayHandler handles OpenAI API gateway requests
|
||||
type OpenAIGatewayHandler struct {
|
||||
gatewayService *service.OpenAIGatewayService
|
||||
billingCacheService *service.BillingCacheService
|
||||
apiKeyService *service.APIKeyService
|
||||
usageRecordWorkerPool *service.UsageRecordWorkerPool
|
||||
errorPassthroughService *service.ErrorPassthroughService
|
||||
contentModerationService *service.ContentModerationService
|
||||
opsService *service.OpsService
|
||||
concurrencyHelper *ConcurrencyHelper
|
||||
imageLimiter *imageConcurrencyLimiter
|
||||
maxAccountSwitches int
|
||||
cfg *config.Config
|
||||
gatewayService *service.OpenAIGatewayService
|
||||
billingCacheService *service.BillingCacheService
|
||||
apiKeyService *service.APIKeyService
|
||||
usageRecordWorkerPool *service.UsageRecordWorkerPool
|
||||
errorPassthroughService *service.ErrorPassthroughService
|
||||
contentModerationService *service.ContentModerationService
|
||||
securityAuditCoordinator *securityaudit.Coordinator
|
||||
grokMediaEligibilityProber grokMediaEligibilityProber
|
||||
opsService *service.OpsService
|
||||
concurrencyHelper *ConcurrencyHelper
|
||||
imageLimiter *imageConcurrencyLimiter
|
||||
maxAccountSwitches int
|
||||
cfg *config.Config
|
||||
}
|
||||
|
||||
type grokMediaEligibilityProber interface {
|
||||
ProbeMediaEligibility(ctx context.Context, accountID int64) (bool, string, error)
|
||||
}
|
||||
|
||||
const maxOpenAIFirstOutputTimeoutSwitches = 1
|
||||
@@ -271,12 +278,15 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
setOpsRequestContext(c, reqModel, reqStream)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(reqStream, false)))
|
||||
|
||||
if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIResponses, reqModel, body); decision != nil && decision.Blocked {
|
||||
h.errorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message)
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIResponses, reqModel, body); decision != nil && !decision.AllowNextStage {
|
||||
h.openAISecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
|
||||
imageIntent := service.IsImageGenerationIntentForPlatform("/v1/responses", reqModel, body, openAICompatibleRequestPlatform(apiKey))
|
||||
// 使用 IsExplicitImageGenerationIntent 排除被动 image_gen namespace 声明。
|
||||
// Codex 在所有请求中被动声明 image_gen namespace,宽泛检测会导致禁了生图的
|
||||
// 分组中所有 Codex 请求被 403(#4447),并误占生图并发槽位。
|
||||
imageIntent := service.IsExplicitImageGenerationIntent("/v1/responses", reqModel, body)
|
||||
if imageIntent && !service.GroupAllowsImageGeneration(apiKey.Group) {
|
||||
h.errorResponse(c, http.StatusForbidden, "permission_error", service.ImageGenerationPermissionMessage())
|
||||
return
|
||||
@@ -353,8 +363,10 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
// 生图意图的 /v1/responses 请求必须调度到确实支持 Responses API 的账号,否则
|
||||
// 会在 forward 阶段被静默降级为无法生图的 Chat Completions 直转(#4417)。
|
||||
// 仅对 OpenAI 平台生效:Grok 生图走独立的 forwardGrokResponses 路径,不应被过滤。
|
||||
// 使用 IsExplicitImageGenerationIntent 排除被动 image_gen namespace 声明,
|
||||
// 避免 Codex 的被动工具目录使 CC-only 账号被误过滤(#4476)。
|
||||
requiredCapability := service.OpenAIEndpointCapabilityChatCompletions
|
||||
if imageIntent && requestPlatform == service.PlatformOpenAI {
|
||||
if service.IsExplicitImageGenerationIntent("/v1/responses", reqModel, body) && requestPlatform == service.PlatformOpenAI {
|
||||
requiredCapability = service.OpenAIEndpointCapabilityResponses
|
||||
}
|
||||
|
||||
@@ -849,8 +861,8 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
|
||||
setOpsRequestContext(c, reqModel, reqStream)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeFromLegacy(reqStream, false)))
|
||||
|
||||
if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolAnthropicMessages, reqModel, body); decision != nil && decision.Blocked {
|
||||
h.anthropicErrorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message)
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolAnthropicMessages, reqModel, body); decision != nil && !decision.AllowNextStage {
|
||||
h.anthropicSecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1473,13 +1485,13 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
setOpsRequestContext(c, reqModel, true)
|
||||
setOpsEndpointContext(c, "", int16(service.RequestTypeWSV2))
|
||||
|
||||
if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIResponses, reqModel, firstMessage); decision != nil && decision.Blocked {
|
||||
writeContentModerationWSError(ctx, wsConn, decision)
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, decision.Message)
|
||||
if decision := h.checkSecurityAuditStage(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIResponses, reqModel, firstMessage, "first_turn"); decision != nil && !decision.AllowNextStage {
|
||||
writeSecurityAuditWSError(ctx, wsConn, decision)
|
||||
closeOpenAIClientWS(wsConn, securityAuditWSCloseStatus(decision), securityAuditWSCloseReason(decision))
|
||||
return
|
||||
}
|
||||
|
||||
imageIntent := service.IsImageGenerationIntentForPlatform("/v1/responses", reqModel, firstMessage, openAICompatibleRequestPlatform(apiKey))
|
||||
imageIntent := service.IsExplicitImageGenerationIntent("/v1/responses", reqModel, firstMessage)
|
||||
if imageIntent && !service.GroupAllowsImageGeneration(apiKey.Group) {
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, service.ImageGenerationPermissionMessage())
|
||||
return
|
||||
@@ -1609,8 +1621,9 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
|
||||
// 与 HTTP Responses 路径保持一致:生图意图请求要求账号支持 Responses API(#4417)。
|
||||
// WSv2 传输本身已隐含 Responses 支持,此处为防御性对齐。
|
||||
// 使用 IsExplicitImageGenerationIntent 排除被动 namespace 声明(#4476)。
|
||||
requiredCapability := service.OpenAIEndpointCapabilityChatCompletions
|
||||
if imageIntent && requestPlatform == service.PlatformOpenAI {
|
||||
if service.IsExplicitImageGenerationIntent("/v1/responses", reqModel, firstMessage) && requestPlatform == service.PlatformOpenAI {
|
||||
requiredCapability = service.OpenAIEndpointCapabilityResponses
|
||||
}
|
||||
|
||||
@@ -1727,9 +1740,9 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
if model == "" {
|
||||
model = reqModel
|
||||
}
|
||||
if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIResponses, model, payload); decision != nil && decision.Blocked {
|
||||
writeContentModerationWSError(ctx, wsConn, decision)
|
||||
return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, decision.Message, nil)
|
||||
if decision := h.checkSecurityAuditStage(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIResponses, model, payload, "subsequent_turn"); decision != nil && !decision.AllowNextStage {
|
||||
writeSecurityAuditWSError(ctx, wsConn, decision)
|
||||
return service.NewOpenAIWSClientCloseError(securityAuditWSCloseStatus(decision), securityAuditWSCloseReason(decision), nil)
|
||||
}
|
||||
return nil
|
||||
},
|
||||
@@ -2223,12 +2236,29 @@ func (h *OpenAIGatewayHandler) mapUpstreamError(statusCode int) (int, string, st
|
||||
|
||||
// handleStreamingAwareError handles errors that may occur after streaming has started
|
||||
func (h *OpenAIGatewayHandler) handleStreamingAwareError(c *gin.Context, status int, errType, message string, streamStarted bool) {
|
||||
h.handleStreamingAwareErrorWithCode(c, status, errType, "", message, streamStarted, false)
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) handleStreamingAwareErrorWithCode(
|
||||
c *gin.Context,
|
||||
status int,
|
||||
errType string,
|
||||
code string,
|
||||
message string,
|
||||
streamStarted bool,
|
||||
countTowardsSLA bool,
|
||||
) {
|
||||
// body-signal compact 心跳可能已把响应头提交为 200:先停心跳(建立
|
||||
// happens-before,接管 ResponseWriter),并升级为流内错误处理。
|
||||
if service.StopOpenAICompactSSEKeepaliveCommitted(c) {
|
||||
streamStarted = true
|
||||
}
|
||||
if streamStarted {
|
||||
if countTowardsSLA {
|
||||
service.MarkOpsStreamFailure(c, errType, code, message, status)
|
||||
} else {
|
||||
service.MarkOpsStreamError(c, errType, message, status)
|
||||
}
|
||||
// /v1/responses 的严格 SDK(Codex CLI)要求终止事件必须属于
|
||||
// response.completed/failed/incomplete/cancelled 集合。
|
||||
// 通用 `event: error` 帧不被识别为终止事件,会导致
|
||||
@@ -2241,8 +2271,15 @@ func (h *OpenAIGatewayHandler) handleStreamingAwareError(c *gin.Context, status
|
||||
// Stream already started, send error as SSE event then close
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if ok {
|
||||
// SSE 错误事件固定 schema,使用 Quote 直拼可避免额外 Marshal 分配。
|
||||
errorEvent := "event: error\ndata: " + `{"error":{"type":` + strconv.Quote(errType) + `,"message":` + strconv.Quote(message) + `}}` + "\n\n"
|
||||
errorObject := gin.H{"type": errType, "message": message}
|
||||
if code != "" {
|
||||
errorObject["code"] = code
|
||||
}
|
||||
payload, err := json.Marshal(gin.H{"error": errorObject})
|
||||
if err != nil {
|
||||
payload = []byte(`{"error":{"type":"upstream_error","message":"Upstream request failed"}}`)
|
||||
}
|
||||
errorEvent := "event: error\ndata: " + string(payload) + "\n\n"
|
||||
if _, err := fmt.Fprint(c.Writer, errorEvent); err != nil {
|
||||
_ = c.Error(err)
|
||||
}
|
||||
@@ -2252,7 +2289,27 @@ func (h *OpenAIGatewayHandler) handleStreamingAwareError(c *gin.Context, status
|
||||
}
|
||||
|
||||
// Normal case: return JSON response with proper status code
|
||||
h.errorResponse(c, status, errType, message)
|
||||
if code == "" {
|
||||
h.errorResponse(c, status, errType, message)
|
||||
return
|
||||
}
|
||||
c.JSON(status, gin.H{"error": gin.H{
|
||||
"type": errType, "code": code, "message": message,
|
||||
}})
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) ensureOpenAIStreamReadErrorResponse(c *gin.Context, err error, streamStarted bool) bool {
|
||||
code, message, ok := service.OpenAIUpstreamStreamReadErrorDetails(err)
|
||||
if !ok || c == nil || c.Writer == nil || service.IsResponseCommitted(c) {
|
||||
return false
|
||||
}
|
||||
if c.Writer.Written() {
|
||||
streamStarted = true
|
||||
}
|
||||
h.handleStreamingAwareErrorWithCode(
|
||||
c, http.StatusBadGateway, "upstream_error", code, message, streamStarted, true,
|
||||
)
|
||||
return true
|
||||
}
|
||||
|
||||
// ensureForwardErrorResponse 在 Forward 返回错误但尚未写响应时补写统一错误响应。
|
||||
|
||||
@@ -96,6 +96,35 @@ func TestOpenAIHandleStreamingAwareError_JSONEscaping(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIHandleStreamingAwareErrorWithCode_EmitsStableClassification(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
w := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(w)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
|
||||
h := &OpenAIGatewayHandler{}
|
||||
h.handleStreamingAwareErrorWithCode(
|
||||
c,
|
||||
http.StatusBadGateway,
|
||||
"upstream_error",
|
||||
service.OpenAIUpstreamHTTP2StreamErrorCode,
|
||||
"Upstream HTTP/2 stream failed",
|
||||
true,
|
||||
true,
|
||||
)
|
||||
|
||||
body := w.Body.String()
|
||||
require.Contains(t, body, "event: error\n")
|
||||
require.Equal(t, "upstream_error", gjson.Get(body[strings.Index(body, "{"):], "error.type").String())
|
||||
require.Equal(t, service.OpenAIUpstreamHTTP2StreamErrorCode, gjson.Get(body[strings.Index(body, "{"):], "error.code").String())
|
||||
require.NotContains(t, body, "stream ID")
|
||||
|
||||
streamErr, ok := service.GetOpsStreamError(c)
|
||||
require.True(t, ok)
|
||||
require.True(t, streamErr.CountTowardsSLA)
|
||||
require.Equal(t, http.StatusBadGateway, streamErr.IntendedStatus)
|
||||
}
|
||||
|
||||
func TestOpenAIForwardSucceededForScheduling(t *testing.T) {
|
||||
require.True(t, openAIForwardSucceededForScheduling(nil))
|
||||
require.True(t, openAIForwardSucceededForScheduling(&service.OpenAIForwardResult{}))
|
||||
@@ -1875,6 +1904,212 @@ func TestOpenAIResponsesWebSocket_FailoverOnUpstreamUsageLimitEvent(t *testing.T
|
||||
require.Equal(t, []int64{int64(9902)}, accountRepo.rateLimitedIDs)
|
||||
}
|
||||
|
||||
func TestOpenAIResponsesWebSocket_FirstOutputTimeoutWithoutDownstreamReusesClientForOneFailover(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
firstHitCh := make(chan []byte, 1)
|
||||
secondHitCh := make(chan []byte, 1)
|
||||
var firstConnections atomic.Int32
|
||||
var secondConnections atomic.Int32
|
||||
|
||||
firstUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
firstConnections.Add(1)
|
||||
conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer func() { _ = conn.CloseNow() }()
|
||||
|
||||
readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
|
||||
_, payload, readErr := conn.Read(readCtx)
|
||||
cancelRead()
|
||||
if readErr == nil {
|
||||
firstHitCh <- payload
|
||||
}
|
||||
|
||||
select {
|
||||
case <-r.Context().Done():
|
||||
case <-time.After(3 * time.Second):
|
||||
}
|
||||
}))
|
||||
defer firstUpstream.Close()
|
||||
|
||||
secondUpstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
secondConnections.Add(1)
|
||||
conn, err := coderws.Accept(w, r, &coderws.AcceptOptions{CompressionMode: coderws.CompressionContextTakeover})
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer func() { _ = conn.CloseNow() }()
|
||||
|
||||
readCtx, cancelRead := context.WithTimeout(r.Context(), 3*time.Second)
|
||||
_, payload, readErr := conn.Read(readCtx)
|
||||
cancelRead()
|
||||
if readErr == nil {
|
||||
secondHitCh <- payload
|
||||
}
|
||||
|
||||
for _, event := range []string{
|
||||
`{"type":"response.created","response":{"id":"resp_ws_timeout_b","model":"gpt-5.1"}}`,
|
||||
`{"type":"response.output_text.delta","response_id":"resp_ws_timeout_b","delta":"recovered"}`,
|
||||
`{"type":"response.completed","response":{"id":"resp_ws_timeout_b","model":"gpt-5.1","usage":{"input_tokens":1,"output_tokens":1}}}`,
|
||||
} {
|
||||
writeCtx, cancelWrite := context.WithTimeout(r.Context(), 3*time.Second)
|
||||
writeErr := conn.Write(writeCtx, coderws.MessageText, []byte(event))
|
||||
cancelWrite()
|
||||
if writeErr != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
readCtx, cancelRead = context.WithTimeout(r.Context(), 3*time.Second)
|
||||
_, _, _ = conn.Read(readCtx)
|
||||
cancelRead()
|
||||
}))
|
||||
defer secondUpstream.Close()
|
||||
|
||||
groupID := int64(4212)
|
||||
accounts := []service.Account{
|
||||
{
|
||||
ID: 9912,
|
||||
Name: "openai-ws-first-semantic-timeout",
|
||||
Platform: service.PlatformOpenAI,
|
||||
Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Priority: 1,
|
||||
Credentials: map[string]any{"api_key": "sk-first", "base_url": firstUpstream.URL},
|
||||
Extra: map[string]any{
|
||||
"openai_apikey_responses_websockets_v2_enabled": true,
|
||||
"openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
|
||||
},
|
||||
},
|
||||
{
|
||||
ID: 9913,
|
||||
Name: "openai-ws-failover-healthy",
|
||||
Platform: service.PlatformOpenAI,
|
||||
Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive,
|
||||
Schedulable: true,
|
||||
Concurrency: 1,
|
||||
Priority: 2,
|
||||
Credentials: map[string]any{"api_key": "sk-second", "base_url": secondUpstream.URL},
|
||||
Extra: map[string]any{
|
||||
"openai_apikey_responses_websockets_v2_enabled": true,
|
||||
"openai_apikey_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
cfg := &config.Config{}
|
||||
cfg.RunMode = config.RunModeSimple
|
||||
cfg.Default.RateMultiplier = 1
|
||||
cfg.Security.URLAllowlist.Enabled = false
|
||||
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
|
||||
cfg.Gateway.OpenAIFirstOutputTimeoutSeconds = 1
|
||||
cfg.Gateway.OpenAIWS.Enabled = true
|
||||
cfg.Gateway.OpenAIWS.APIKeyEnabled = true
|
||||
cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
|
||||
cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
|
||||
cfg.Gateway.OpenAIWS.DialTimeoutSeconds = 3
|
||||
cfg.Gateway.OpenAIWS.ReadTimeoutSeconds = 3
|
||||
cfg.Gateway.OpenAIWS.WriteTimeoutSeconds = 3
|
||||
cfg.Gateway.OpenAIWS.IngressInterTurnIdleTimeoutSeconds = 3
|
||||
cfg.Gateway.MaxAccountSwitches = 3
|
||||
|
||||
accountRepo := &openAIWSFailoverHandlerAccountRepoStub{accounts: accounts}
|
||||
rateLimitSvc := service.NewRateLimitService(accountRepo, nil, cfg, nil, nil)
|
||||
billingCacheSvc := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
|
||||
gatewaySvc := service.NewOpenAIGatewayService(
|
||||
accountRepo, nil, nil, nil, nil, nil, nil, cfg, nil, nil,
|
||||
service.NewBillingService(cfg, nil), rateLimitSvc, billingCacheSvc,
|
||||
nil, &service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
cache := &concurrencyCacheMock{
|
||||
acquireUserSlotFn: func(context.Context, int64, int, string) (bool, error) { return true, nil },
|
||||
acquireAccountSlotFn: func(context.Context, int64, int, string) (bool, error) {
|
||||
return true, nil
|
||||
},
|
||||
}
|
||||
h := &OpenAIGatewayHandler{
|
||||
gatewayService: gatewaySvc,
|
||||
billingCacheService: billingCacheSvc,
|
||||
apiKeyService: &service.APIKeyService{},
|
||||
concurrencyHelper: NewConcurrencyHelper(service.NewConcurrencyService(cache), SSEPingFormatNone, time.Second),
|
||||
maxAccountSwitches: 3,
|
||||
}
|
||||
|
||||
apiKey := &service.APIKey{
|
||||
ID: 1812,
|
||||
GroupID: &groupID,
|
||||
User: &service.User{ID: 1712, Status: service.StatusActive},
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, Status: service.StatusActive},
|
||||
}
|
||||
handlerDone := make(chan struct{})
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.User.ID, Concurrency: 1})
|
||||
c.Next()
|
||||
})
|
||||
router.GET("/openai/v1/responses", func(c *gin.Context) {
|
||||
h.ResponsesWebSocket(c)
|
||||
close(handlerDone)
|
||||
})
|
||||
handlerServer := httptest.NewServer(router)
|
||||
defer handlerServer.Close()
|
||||
|
||||
dialCtx, cancelDial := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
clientConn, _, err := coderws.Dial(
|
||||
dialCtx,
|
||||
"ws"+strings.TrimPrefix(handlerServer.URL, "http")+"/openai/v1/responses",
|
||||
&coderws.DialOptions{CompressionMode: coderws.CompressionContextTakeover},
|
||||
)
|
||||
cancelDial()
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = clientConn.CloseNow() }()
|
||||
|
||||
writeCtx, cancelWrite := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
err = clientConn.Write(writeCtx, coderws.MessageText, []byte(`{"type":"response.create","model":"gpt-5.1","stream":false}`))
|
||||
cancelWrite()
|
||||
require.NoError(t, err)
|
||||
|
||||
var eventTypes []string
|
||||
readCtx, cancelRead := context.WithTimeout(context.Background(), 6*time.Second)
|
||||
for {
|
||||
_, event, readErr := clientConn.Read(readCtx)
|
||||
require.NoError(t, readErr)
|
||||
eventType := gjson.GetBytes(event, "type").String()
|
||||
eventTypes = append(eventTypes, eventType)
|
||||
if eventType == "response.completed" {
|
||||
require.Equal(t, "resp_ws_timeout_b", gjson.GetBytes(event, "response.id").String())
|
||||
break
|
||||
}
|
||||
}
|
||||
cancelRead()
|
||||
require.Contains(t, eventTypes, "response.output_text.delta")
|
||||
require.NoError(t, clientConn.Close(coderws.StatusNormalClosure, "done"))
|
||||
|
||||
select {
|
||||
case <-handlerDone:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("websocket handler did not finish after healthy failover turn")
|
||||
}
|
||||
select {
|
||||
case <-firstHitCh:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("first upstream did not receive replayable request")
|
||||
}
|
||||
select {
|
||||
case <-secondHitCh:
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("second upstream did not receive replayed request")
|
||||
}
|
||||
require.Equal(t, int32(1), firstConnections.Load())
|
||||
require.Equal(t, int32(1), secondConnections.Load())
|
||||
require.NotContains(t, accountRepo.rateLimitedIDs, int64(9913), "healthy failover account must not be penalized")
|
||||
}
|
||||
|
||||
func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSUsageLogCase) openAIResponsesWSUsageLogResult {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -30,17 +30,11 @@ func TestOpenAIGatewayHandlerResponses_GrokResponsesLiteImageToolDeclarationBypa
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayHandlerResponses_ImagePermissionHardSignalsStillRejected(t *testing.T) {
|
||||
passiveNamespace := `{"model":"gpt-5.5","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}],"tool_choice":"auto","input":"write code"}`
|
||||
tests := []struct {
|
||||
name string
|
||||
platform string
|
||||
body string
|
||||
}{
|
||||
{
|
||||
name: "OpenAI keeps declaration semantics",
|
||||
platform: service.PlatformOpenAI,
|
||||
body: passiveNamespace,
|
||||
},
|
||||
{
|
||||
name: "Grok native image_generation declaration",
|
||||
platform: service.PlatformGrok,
|
||||
@@ -51,6 +45,16 @@ func TestOpenAIGatewayHandlerResponses_ImagePermissionHardSignalsStillRejected(t
|
||||
platform: service.PlatformGrok,
|
||||
body: `{"model":"grok-4.5","tools":[{"type":"namespace","name":"image_gen"}],"tool_choice":{"type":"namespace","name":"image_gen"},"input":"draw"}`,
|
||||
},
|
||||
{
|
||||
name: "OpenAI native image_generation tool",
|
||||
platform: service.PlatformOpenAI,
|
||||
body: `{"model":"gpt-5.5","tools":[{"type":"image_generation","model":"gpt-image-2"}],"input":"draw a cat"}`,
|
||||
},
|
||||
{
|
||||
name: "OpenAI image model",
|
||||
platform: service.PlatformOpenAI,
|
||||
body: `{"model":"gpt-image-2","input":"draw a cat"}`,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
@@ -63,6 +67,14 @@ func TestOpenAIGatewayHandlerResponses_ImagePermissionHardSignalsStillRejected(t
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIGatewayHandlerResponses_PassiveNamespaceDoesNotTrigger403(t *testing.T) {
|
||||
passiveNamespace := `{"model":"gpt-5.5","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen"}]}],"tool_choice":"auto","input":"write code"}`
|
||||
rec := runOpenAIResponsesImagePermissionGateTest(t, service.PlatformOpenAI, passiveNamespace)
|
||||
|
||||
require.NotEqual(t, http.StatusForbidden, rec.Code,
|
||||
"passive image_gen namespace with tool_choice=auto should not trigger 403 (#4447)")
|
||||
}
|
||||
|
||||
func runOpenAIResponsesImagePermissionGateTest(t *testing.T, platform string, body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -86,8 +86,8 @@ func (h *OpenAIGatewayHandler) Images(c *gin.Context) {
|
||||
h.errorResponse(c, http.StatusForbidden, "permission_error", service.ImageGenerationPermissionMessage())
|
||||
return
|
||||
}
|
||||
if decision := h.checkContentModeration(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, requestModel, parsed.ModerationBody()); decision != nil && decision.Blocked {
|
||||
h.errorResponse(c, contentModerationStatus(decision), contentModerationErrorCode(decision), decision.Message)
|
||||
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, requestModel, parsed.ModerationBody()); decision != nil && !decision.AllowNextStage {
|
||||
h.openAISecurityAuditError(c, decision)
|
||||
return
|
||||
}
|
||||
imageReleaseFunc, acquired := h.acquireImageGenerationSlot(c, streamStarted)
|
||||
|
||||
@@ -72,24 +72,10 @@ const (
|
||||
opsErrorLogMinQueueSize = 256
|
||||
opsErrorLogMaxQueueSize = 8192
|
||||
opsErrorLogBatchSize = 32
|
||||
opsErrorLogMaxQueueBytes = 32 * 1024 * 1024
|
||||
opsErrorLogMaxUserAgentBytes = 512
|
||||
)
|
||||
|
||||
// looksLikeSystemKey 粗筛"形似本系统 key"的输入:长度 16-128 且仅含 [a-zA-Z0-9_-]。
|
||||
// 不用前缀匹配(APIKeyPrefix 可配置)。用于反查审计表前挡掉随机扫描的乱码输入。
|
||||
func looksLikeSystemKey(key string) bool {
|
||||
if len(key) < 16 || len(key) > 128 {
|
||||
return false
|
||||
}
|
||||
for _, c := range key {
|
||||
allowed := (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') ||
|
||||
(c >= '0' && c <= '9') || c == '_' || c == '-'
|
||||
if !allowed {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// keyPrefix 返回脱敏前缀(前 n 个字符);不足 n 则原样返回。
|
||||
func keyPrefix(key string, n int) string {
|
||||
if len(key) <= n {
|
||||
@@ -98,43 +84,26 @@ func keyPrefix(key string, n int) string {
|
||||
return key[:n]
|
||||
}
|
||||
|
||||
// extractAttemptedKey 按认证中间件同样的顺序从请求头提取提交的 key 明文。
|
||||
// 与 api_key_auth.go:43-59 一致:Authorization 仅取 Bearer scheme,非 Bearer 则忽略并继续 x-api-key → x-goog-api-key。
|
||||
func extractAttemptedKey(c *gin.Context) string {
|
||||
if h := c.GetHeader("Authorization"); h != "" {
|
||||
parts := strings.SplitN(h, " ", 2)
|
||||
if len(parts) == 2 && strings.EqualFold(parts[0], "Bearer") {
|
||||
return strings.TrimSpace(parts[1])
|
||||
}
|
||||
// 非 Bearer:与中间件一致,忽略 Authorization,继续尝试其它 header(不在此 return)。
|
||||
}
|
||||
if k := c.GetHeader("x-api-key"); k != "" {
|
||||
return strings.TrimSpace(k)
|
||||
}
|
||||
if k := c.GetHeader("x-goog-api-key"); k != "" {
|
||||
return strings.TrimSpace(k)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
type opsErrorLogJob struct {
|
||||
ops *service.OpsService
|
||||
entry *service.OpsInsertErrorLogInput
|
||||
ops *service.OpsService
|
||||
entry *service.OpsInsertErrorLogInput
|
||||
queuedBytes int64
|
||||
}
|
||||
|
||||
var (
|
||||
opsErrorLogOnce sync.Once
|
||||
opsErrorLogQueue chan opsErrorLogJob
|
||||
|
||||
opsErrorLogStopOnce sync.Once
|
||||
opsErrorLogWorkersWg sync.WaitGroup
|
||||
opsErrorLogMu sync.RWMutex
|
||||
opsErrorLogStopping bool
|
||||
opsErrorLogQueueLen atomic.Int64
|
||||
opsErrorLogEnqueued atomic.Int64
|
||||
opsErrorLogDropped atomic.Int64
|
||||
opsErrorLogProcessed atomic.Int64
|
||||
opsErrorLogSanitized atomic.Int64
|
||||
opsErrorLogStopOnce sync.Once
|
||||
opsErrorLogWorkersWg sync.WaitGroup
|
||||
opsErrorLogMu sync.RWMutex
|
||||
opsErrorLogStopping bool
|
||||
opsErrorLogQueueLen atomic.Int64
|
||||
opsErrorLogQueueBytes atomic.Int64
|
||||
opsErrorLogEnqueued atomic.Int64
|
||||
opsErrorLogDropped atomic.Int64
|
||||
opsErrorLogProcessed atomic.Int64
|
||||
opsErrorLogSanitized atomic.Int64
|
||||
|
||||
opsErrorLogLastDropLogAt atomic.Int64
|
||||
|
||||
@@ -154,6 +123,7 @@ func startOpsErrorLogWorkers() {
|
||||
workerCount, queueSize := opsErrorLogConfig()
|
||||
opsErrorLogQueue = make(chan opsErrorLogJob, queueSize)
|
||||
opsErrorLogQueueLen.Store(0)
|
||||
opsErrorLogQueueBytes.Store(0)
|
||||
|
||||
opsErrorLogWorkersWg.Add(workerCount)
|
||||
for i := 0; i < workerCount; i++ {
|
||||
@@ -165,6 +135,7 @@ func startOpsErrorLogWorkers() {
|
||||
return
|
||||
}
|
||||
opsErrorLogQueueLen.Add(-1)
|
||||
opsErrorLogQueueBytes.Add(-job.queuedBytes)
|
||||
batch := make([]opsErrorLogJob, 0, opsErrorLogBatchSize)
|
||||
batch = append(batch, job)
|
||||
|
||||
@@ -184,6 +155,7 @@ func startOpsErrorLogWorkers() {
|
||||
return
|
||||
}
|
||||
opsErrorLogQueueLen.Add(-1)
|
||||
opsErrorLogQueueBytes.Add(-nextJob.queuedBytes)
|
||||
batch = append(batch, nextJob)
|
||||
case <-timer.C:
|
||||
break batchLoop
|
||||
@@ -239,6 +211,20 @@ func enqueueOpsErrorLog(ops *service.OpsService, entry *service.OpsInsertErrorLo
|
||||
if ops == nil || entry == nil {
|
||||
return
|
||||
}
|
||||
entry.UserAgent = normalizeOpsPersistentUserAgent(entry.UserAgent)
|
||||
if entry.ErrorBody != "" {
|
||||
originalBody := entry.ErrorBody
|
||||
body, truncated := service.SanitizeOpsErrorBodyForQueue(originalBody)
|
||||
entry.ErrorBody = body
|
||||
if truncated || body != originalBody {
|
||||
opsErrorLogSanitized.Add(1)
|
||||
}
|
||||
}
|
||||
if err := service.SanitizeOpsUpstreamErrorsForQueue(entry); err != nil {
|
||||
opsErrorLogDropped.Add(1)
|
||||
maybeLogOpsErrorLogDrop()
|
||||
return
|
||||
}
|
||||
select {
|
||||
case <-opsErrorLogShutdownCh:
|
||||
return
|
||||
@@ -259,18 +245,29 @@ func enqueueOpsErrorLog(ops *service.OpsService, entry *service.OpsInsertErrorLo
|
||||
if opsErrorLogStopping || opsErrorLogQueue == nil {
|
||||
return
|
||||
}
|
||||
queuedBytes := estimateOpsErrorLogJobBytes(entry)
|
||||
if !reserveOpsErrorLogQueueBytes(queuedBytes) {
|
||||
opsErrorLogDropped.Add(1)
|
||||
maybeLogOpsErrorLogDrop()
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case opsErrorLogQueue <- opsErrorLogJob{ops: ops, entry: entry}:
|
||||
opsErrorLogQueueLen.Add(1)
|
||||
case opsErrorLogQueue <- opsErrorLogJob{ops: ops, entry: entry, queuedBytes: queuedBytes}:
|
||||
opsErrorLogEnqueued.Add(1)
|
||||
default:
|
||||
opsErrorLogQueueLen.Add(-1)
|
||||
opsErrorLogQueueBytes.Add(-queuedBytes)
|
||||
// Queue is full; drop to avoid blocking request handling.
|
||||
opsErrorLogDropped.Add(1)
|
||||
maybeLogOpsErrorLogDrop()
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeOpsPersistentUserAgent(value string) string {
|
||||
return truncateString(strings.TrimSpace(strings.ToValidUTF8(value, "")), opsErrorLogMaxUserAgentBytes)
|
||||
}
|
||||
|
||||
func StopOpsErrorLogWorkers() bool {
|
||||
opsErrorLogStopOnce.Do(func() {
|
||||
opsErrorLogShutdownOnce.Do(func() {
|
||||
@@ -293,6 +290,7 @@ func stopOpsErrorLogWorkers() bool {
|
||||
|
||||
if ch == nil {
|
||||
opsErrorLogQueueLen.Store(0)
|
||||
opsErrorLogQueueBytes.Store(0)
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -305,6 +303,7 @@ func stopOpsErrorLogWorkers() bool {
|
||||
select {
|
||||
case <-done:
|
||||
opsErrorLogQueueLen.Store(0)
|
||||
opsErrorLogQueueBytes.Store(0)
|
||||
return true
|
||||
case <-time.After(opsErrorLogDrainTimeout):
|
||||
return false
|
||||
@@ -315,6 +314,14 @@ func OpsErrorLogQueueLength() int64 {
|
||||
return opsErrorLogQueueLen.Load()
|
||||
}
|
||||
|
||||
func OpsErrorLogQueueBytes() int64 {
|
||||
return opsErrorLogQueueBytes.Load()
|
||||
}
|
||||
|
||||
func OpsErrorLogQueueBytesCapacity() int64 {
|
||||
return opsErrorLogMaxQueueBytes
|
||||
}
|
||||
|
||||
func OpsErrorLogQueueCapacity() int {
|
||||
opsErrorLogMu.RLock()
|
||||
ch := opsErrorLogQueue
|
||||
@@ -355,12 +362,15 @@ func maybeLogOpsErrorLogDrop() {
|
||||
}
|
||||
|
||||
queued := opsErrorLogQueueLen.Load()
|
||||
queuedBytes := opsErrorLogQueueBytes.Load()
|
||||
queueCap := OpsErrorLogQueueCapacity()
|
||||
|
||||
log.Printf(
|
||||
"[OpsErrorLogger] queue is full; dropping logs (queued=%d cap=%d enqueued_total=%d dropped_total=%d processed_total=%d sanitized_total=%d)",
|
||||
"[OpsErrorLogger] queue is full; dropping logs (queued=%d cap=%d queued_bytes=%d bytes_cap=%d enqueued_total=%d dropped_total=%d processed_total=%d sanitized_total=%d)",
|
||||
queued,
|
||||
queueCap,
|
||||
queuedBytes,
|
||||
opsErrorLogMaxQueueBytes,
|
||||
opsErrorLogEnqueued.Load(),
|
||||
opsErrorLogDropped.Load(),
|
||||
opsErrorLogProcessed.Load(),
|
||||
@@ -368,6 +378,46 @@ func maybeLogOpsErrorLogDrop() {
|
||||
)
|
||||
}
|
||||
|
||||
func reserveOpsErrorLogQueueBytes(size int64) bool {
|
||||
if size < 1 {
|
||||
size = 1
|
||||
}
|
||||
for {
|
||||
current := opsErrorLogQueueBytes.Load()
|
||||
if current > opsErrorLogMaxQueueBytes-size {
|
||||
return false
|
||||
}
|
||||
if opsErrorLogQueueBytes.CompareAndSwap(current, current+size) {
|
||||
opsErrorLogQueueLen.Add(1)
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func estimateOpsErrorLogJobBytes(entry *service.OpsInsertErrorLogInput) int64 {
|
||||
if entry == nil {
|
||||
return 1
|
||||
}
|
||||
const fixedOverhead = 512
|
||||
size := fixedOverhead + len(entry.RequestID) + len(entry.ClientRequestID) +
|
||||
len(entry.Platform) + len(entry.Model) + len(entry.RequestPath) +
|
||||
len(entry.InboundEndpoint) + len(entry.UpstreamEndpoint) +
|
||||
len(entry.RequestedModel) + len(entry.UpstreamModel) + len(entry.UserAgent) +
|
||||
len(entry.ErrorPhase) + len(entry.ErrorType) + len(entry.Severity) +
|
||||
len(entry.ErrorMessage) + len(entry.ErrorBody) + len(entry.ErrorSource) +
|
||||
len(entry.ErrorOwner) + len(entry.APIKeyPrefix)
|
||||
if entry.UpstreamErrorMessage != nil {
|
||||
size += len(*entry.UpstreamErrorMessage)
|
||||
}
|
||||
if entry.UpstreamErrorDetail != nil {
|
||||
size += len(*entry.UpstreamErrorDetail)
|
||||
}
|
||||
if entry.UpstreamErrorsJSON != nil {
|
||||
size += len(*entry.UpstreamErrorsJSON)
|
||||
}
|
||||
return int64(size)
|
||||
}
|
||||
|
||||
func opsErrorLogConfig() (workerCount int, queueSize int) {
|
||||
workerCount = runtime.GOMAXPROCS(0) * 2
|
||||
if workerCount < opsErrorLogMinWorkerCount {
|
||||
@@ -470,9 +520,12 @@ type opsCaptureWriter struct {
|
||||
gin.ResponseWriter
|
||||
limit int
|
||||
buf bytes.Buffer
|
||||
ctx *gin.Context
|
||||
}
|
||||
|
||||
const opsCaptureWriterLimit = 64 * 1024
|
||||
const opsCaptureWriterLimit = service.OpsErrorLogQueueBodyMaxBytes
|
||||
|
||||
const opsCaptureWriterPoolMaxRetainedCapacity = service.OpsErrorLogQueueBodyMaxBytes
|
||||
|
||||
var opsCaptureWriterPool = sync.Pool{
|
||||
New: func() any {
|
||||
@@ -496,11 +549,19 @@ func releaseOpsCaptureWriter(w *opsCaptureWriter) {
|
||||
return
|
||||
}
|
||||
w.ResponseWriter = nil
|
||||
w.ctx = nil
|
||||
w.limit = opsCaptureWriterLimit
|
||||
if !shouldPoolOpsCaptureWriter(w) {
|
||||
return
|
||||
}
|
||||
w.buf.Reset()
|
||||
opsCaptureWriterPool.Put(w)
|
||||
}
|
||||
|
||||
func shouldPoolOpsCaptureWriter(w *opsCaptureWriter) bool {
|
||||
return w != nil && w.buf.Cap() <= opsCaptureWriterPoolMaxRetainedCapacity
|
||||
}
|
||||
|
||||
func (w *opsCaptureWriter) Status() int {
|
||||
if w.ResponseWriter == nil {
|
||||
return 0
|
||||
@@ -577,7 +638,7 @@ 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 {
|
||||
if w.shouldCapture() && w.Status() >= 400 && w.limit > 0 && w.buf.Len() < w.limit {
|
||||
remaining := w.limit - w.buf.Len()
|
||||
if len(b) > remaining {
|
||||
_, _ = w.buf.Write(b[:remaining])
|
||||
@@ -592,7 +653,7 @@ 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 {
|
||||
if w.shouldCapture() && w.Status() >= 400 && w.limit > 0 && w.buf.Len() < w.limit {
|
||||
remaining := w.limit - w.buf.Len()
|
||||
if len(s) > remaining {
|
||||
_, _ = w.buf.WriteString(s[:remaining])
|
||||
@@ -603,6 +664,14 @@ func (w *opsCaptureWriter) WriteString(s string) (int, error) {
|
||||
return w.ResponseWriter.WriteString(s)
|
||||
}
|
||||
|
||||
func (w *opsCaptureWriter) shouldCapture() bool {
|
||||
if w.ctx == nil {
|
||||
return true
|
||||
}
|
||||
_, rejected := middleware2.GetIngressRejectReason(w.ctx)
|
||||
return !rejected
|
||||
}
|
||||
|
||||
// OpsErrorLoggerMiddleware records error responses (status >= 400) into ops_error_logs.
|
||||
//
|
||||
// Notes:
|
||||
@@ -612,6 +681,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
originalWriter := c.Writer
|
||||
w := acquireOpsCaptureWriter(originalWriter)
|
||||
w.ctx = c
|
||||
defer func() {
|
||||
// Restore the original writer before returning so outer middlewares
|
||||
// don't observe a pooled wrapper that has been released.
|
||||
@@ -623,6 +693,10 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
c.Writer = w
|
||||
c.Next()
|
||||
|
||||
if _, rejected := middleware2.GetIngressRejectReason(c); rejected {
|
||||
return
|
||||
}
|
||||
|
||||
if ops == nil {
|
||||
return
|
||||
}
|
||||
@@ -989,7 +1063,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
IsCountTokens: isCountTokensRequest(c),
|
||||
|
||||
ErrorMessage: parsed.Message,
|
||||
// Keep the full captured error body (capture is already capped at 64KB) so the
|
||||
// Keep the captured error body (already capped at the queue-safe limit) so the
|
||||
// service layer can sanitize JSON before truncating for storage.
|
||||
ErrorBody: string(body),
|
||||
ErrorSource: errorSource,
|
||||
@@ -1002,7 +1076,7 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
|
||||
if apiKey != nil {
|
||||
entry.APIKeyID = &apiKey.ID
|
||||
// 有效(未删除)key 报错时快照前缀,key 之后被删也保留;与 INVALID_API_KEY 的 attempted_key_prefix 互斥。
|
||||
// 有效 key 报错时快照前缀,key 之后被删也保留。
|
||||
entry.APIKeyPrefix = keyPrefix(apiKey.Key, 8)
|
||||
if apiKey.User != nil {
|
||||
entry.UserID = &apiKey.User.ID
|
||||
@@ -1022,22 +1096,6 @@ func OpsErrorLoggerMiddleware(ops *service.OpsService) gin.HandlerFunc {
|
||||
entry.ClientIP = &clientIP
|
||||
}
|
||||
|
||||
// 已删除 key 归因:仅 INVALID_API_KEY 才尝试。响应已写出,此处不阻塞客户端。
|
||||
if parsed.Code == opsCodeInvalidAPIKey {
|
||||
if attemptedKey := extractAttemptedKey(c); attemptedKey != "" {
|
||||
entry.AttemptedKeyPrefix = keyPrefix(attemptedKey, 8)
|
||||
if looksLikeSystemKey(attemptedKey) {
|
||||
if res, lookupErr := ops.LookupDeletedKeyAudit(c.Request.Context(), attemptedKey); lookupErr != nil {
|
||||
log.Printf("[OpsErrorLogger] LookupDeletedKeyAudit failed: %v", lookupErr)
|
||||
} else if res != nil {
|
||||
owner := res.UserID
|
||||
entry.DeletedKeyOwnerUserID = &owner
|
||||
entry.DeletedKeyName = res.KeyName
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enqueueOpsErrorLog(ops, entry)
|
||||
}
|
||||
}
|
||||
@@ -1072,8 +1130,20 @@ func logOpsStreamError(c *gin.Context, ops *service.OpsService, wireStatus int)
|
||||
if classifyStatus <= 0 {
|
||||
classifyStatus = wireStatus
|
||||
}
|
||||
normalizedType := normalizeOpsErrorType(streamErr.ErrType, "")
|
||||
phase, isBusinessLimited, errorOwner, errorSource := classifyOpsErrorLog(c, normalizedType, streamErr.Message, "", classifyStatus)
|
||||
normalizedType := normalizeOpsErrorType(streamErr.ErrType, streamErr.Code)
|
||||
phase, isBusinessLimited, errorOwner, errorSource := classifyOpsErrorLog(c, normalizedType, streamErr.Message, streamErr.Code, classifyStatus)
|
||||
recordedStatus := wireStatus
|
||||
if streamErr.CountTowardsSLA && streamErr.IntendedStatus >= 400 {
|
||||
recordedStatus = streamErr.IntendedStatus
|
||||
}
|
||||
errorBody := ""
|
||||
if streamErr.Code != "" {
|
||||
if payload, err := json.Marshal(gin.H{"error": gin.H{
|
||||
"type": normalizedType, "code": streamErr.Code, "message": streamErr.Message,
|
||||
}}); err == nil {
|
||||
errorBody = string(payload)
|
||||
}
|
||||
}
|
||||
|
||||
apiKey := getOpsAPIKey(c)
|
||||
clientRequestID, _ := c.Request.Context().Value(ctxkey.ClientRequestID).(string)
|
||||
@@ -1140,12 +1210,12 @@ func logOpsStreamError(c *gin.Context, ops *service.OpsService, wireStatus int)
|
||||
ErrorPhase: phase,
|
||||
ErrorType: normalizedType,
|
||||
Severity: classifyOpsSeverity(normalizedType, classifyStatus),
|
||||
StatusCode: wireStatus,
|
||||
StatusCode: recordedStatus,
|
||||
IsBusinessLimited: isBusinessLimited,
|
||||
IsCountTokens: isCountTokensRequest(c),
|
||||
|
||||
ErrorMessage: streamErr.Message,
|
||||
ErrorBody: "",
|
||||
ErrorBody: errorBody,
|
||||
ErrorSource: errorSource,
|
||||
ErrorOwner: errorOwner,
|
||||
|
||||
@@ -1694,11 +1764,8 @@ func shouldSkipOpsErrorLog(ctx context.Context, ops *service.OpsService, message
|
||||
}
|
||||
|
||||
// Get advanced settings to check filter configuration
|
||||
settings, err := ops.GetOpsAdvancedSettings(ctx)
|
||||
if err != nil || settings == nil {
|
||||
// If we can't get settings, don't skip (fail open)
|
||||
return false
|
||||
}
|
||||
_ = ctx
|
||||
settings := ops.OpsAdvancedSettingsSnapshot()
|
||||
|
||||
msgLower := strings.ToLower(message)
|
||||
bodyLower := strings.ToLower(body)
|
||||
|
||||
@@ -1,39 +1,9 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func TestLooksLikeSystemKey(t *testing.T) {
|
||||
cases := []struct {
|
||||
in string
|
||||
want bool
|
||||
}{
|
||||
{"sk-abcdef0123456789", true},
|
||||
{"ABCdef_-0123456789", true},
|
||||
{"short", false},
|
||||
{"with space xxxxxxxxxx", false},
|
||||
{"汉字key1234567890", false},
|
||||
{"", false},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := looksLikeSystemKey(c.in); got != c.want {
|
||||
t.Errorf("looksLikeSystemKey(%q)=%v want %v", c.in, got, c.want)
|
||||
}
|
||||
}
|
||||
long := make([]byte, 129)
|
||||
for i := range long {
|
||||
long[i] = 'a'
|
||||
}
|
||||
if looksLikeSystemKey(string(long)) {
|
||||
t.Errorf("129-char key should be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyPrefix(t *testing.T) {
|
||||
if got := keyPrefix("sk-3f2a9c7e", 8); got != "sk-3f2a9" {
|
||||
t.Errorf("keyPrefix=%q want %q", got, "sk-3f2a9")
|
||||
@@ -42,77 +12,3 @@ func TestKeyPrefix(t *testing.T) {
|
||||
t.Errorf("short key should be returned as-is, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExtractAttemptedKey(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
headers map[string]string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "Bearer in Authorization",
|
||||
headers: map[string]string{"Authorization": "Bearer sk-testkey0123456789"},
|
||||
want: "sk-testkey0123456789",
|
||||
},
|
||||
{
|
||||
name: "Bearer case-insensitive",
|
||||
headers: map[string]string{"Authorization": "BEARER sk-testkey0123456789"},
|
||||
want: "sk-testkey0123456789",
|
||||
},
|
||||
{
|
||||
name: "x-api-key header",
|
||||
headers: map[string]string{"x-api-key": "sk-xapikey0123456789"},
|
||||
want: "sk-xapikey0123456789",
|
||||
},
|
||||
{
|
||||
name: "x-goog-api-key header",
|
||||
headers: map[string]string{"x-goog-api-key": "sk-goog0123456789"},
|
||||
want: "sk-goog0123456789",
|
||||
},
|
||||
{
|
||||
name: "Authorization takes priority over x-api-key",
|
||||
headers: map[string]string{"Authorization": "Bearer sk-auth0123456789", "x-api-key": "sk-xapi0123456789"},
|
||||
want: "sk-auth0123456789",
|
||||
},
|
||||
{
|
||||
name: "x-api-key takes priority over x-goog-api-key",
|
||||
headers: map[string]string{"x-api-key": "sk-xapi0123456789", "x-goog-api-key": "sk-goog0123456789"},
|
||||
want: "sk-xapi0123456789",
|
||||
},
|
||||
{
|
||||
name: "no key headers",
|
||||
headers: map[string]string{},
|
||||
want: "",
|
||||
},
|
||||
{
|
||||
name: "Bearer with leading/trailing spaces trimmed",
|
||||
headers: map[string]string{"Authorization": "Bearer sk-trimmed0123456789 "},
|
||||
want: "sk-trimmed0123456789",
|
||||
},
|
||||
{
|
||||
// 非 Bearer Authorization 应被忽略,继续 fall-through 到 x-api-key(与认证中间件一致)
|
||||
name: "non-Bearer Authorization falls through to x-api-key",
|
||||
headers: map[string]string{"Authorization": "junk-not-bearer", "x-api-key": "sk-realkey1234567"},
|
||||
want: "sk-realkey1234567",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
for k, v := range tc.headers {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
c.Request = req
|
||||
|
||||
got := extractAttemptedKey(c)
|
||||
if got != tc.want {
|
||||
t.Errorf("extractAttemptedKey(%v) = %q, want %q", tc.headers, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,10 +1,13 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"unicode/utf8"
|
||||
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
@@ -12,6 +15,82 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type ingressRejectSettingRepo struct {
|
||||
service.SettingRepository
|
||||
getValueCalls int
|
||||
}
|
||||
|
||||
func (r *ingressRejectSettingRepo) GetValue(context.Context, string) (string, error) {
|
||||
r.getValueCalls++
|
||||
return "", service.ErrSettingNotFound
|
||||
}
|
||||
|
||||
func (r *ingressRejectSettingRepo) GetMultiple(context.Context, []string) (map[string]string, error) {
|
||||
r.getValueCalls++
|
||||
return map[string]string{}, nil
|
||||
}
|
||||
|
||||
func (r *ingressRejectSettingRepo) Set(context.Context, string, string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type ingressRejectOpsRepo struct {
|
||||
service.OpsRepository
|
||||
insertCalls int
|
||||
}
|
||||
|
||||
func (r *ingressRejectOpsRepo) InsertErrorLog(context.Context, *service.OpsInsertErrorLogInput) (int64, error) {
|
||||
r.insertCalls++
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func (r *ingressRejectOpsRepo) BatchInsertErrorLogs(context.Context, []*service.OpsInsertErrorLogInput) (int64, error) {
|
||||
r.insertCalls++
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
func TestOpsErrorLogQueueByteBudget(t *testing.T) {
|
||||
previousBytes := opsErrorLogQueueBytes.Load()
|
||||
previousLen := opsErrorLogQueueLen.Load()
|
||||
opsErrorLogQueueBytes.Store(0)
|
||||
opsErrorLogQueueLen.Store(0)
|
||||
t.Cleanup(func() {
|
||||
opsErrorLogQueueBytes.Store(previousBytes)
|
||||
opsErrorLogQueueLen.Store(previousLen)
|
||||
})
|
||||
|
||||
if !reserveOpsErrorLogQueueBytes(opsErrorLogMaxQueueBytes - 1) {
|
||||
t.Fatal("first reservation within byte budget should succeed")
|
||||
}
|
||||
if reserveOpsErrorLogQueueBytes(2) {
|
||||
t.Fatal("reservation beyond byte budget should be rejected")
|
||||
}
|
||||
if got := OpsErrorLogQueueBytes(); got != opsErrorLogMaxQueueBytes-1 {
|
||||
t.Fatalf("queued bytes = %d, want %d", got, opsErrorLogMaxQueueBytes-1)
|
||||
}
|
||||
if got := OpsErrorLogQueueLength(); got != 1 {
|
||||
t.Fatalf("queue length = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEstimateOpsErrorLogJobBytesIncludesVariablePayloads(t *testing.T) {
|
||||
base := estimateOpsErrorLogJobBytes(&service.OpsInsertErrorLogInput{})
|
||||
message := "upstream message"
|
||||
detail := "upstream detail"
|
||||
events := `[{"error":"x"}]`
|
||||
entry := &service.OpsInsertErrorLogInput{
|
||||
ErrorBody: strings.Repeat("x", 1024),
|
||||
ErrorMessage: "client error",
|
||||
UserAgent: "test-agent",
|
||||
UpstreamErrorMessage: &message,
|
||||
UpstreamErrorDetail: &detail,
|
||||
UpstreamErrorsJSON: &events,
|
||||
}
|
||||
if got := estimateOpsErrorLogJobBytes(entry); got <= base+1024 {
|
||||
t.Fatalf("estimated bytes = %d, expected variable payloads above %d", got, base+1024)
|
||||
}
|
||||
}
|
||||
|
||||
func resetOpsErrorLoggerStateForTest(t *testing.T) {
|
||||
t.Helper()
|
||||
|
||||
@@ -119,6 +198,29 @@ func TestOpsCaptureWriterPool_ResetOnRelease(t *testing.T) {
|
||||
require.Zero(t, reused.buf.Len(), "writer should be reset before reuse")
|
||||
}
|
||||
|
||||
func TestOpsCaptureWriterPool_DropsLargeBuffers(t *testing.T) {
|
||||
w := &opsCaptureWriter{}
|
||||
w.buf.Grow(opsCaptureWriterPoolMaxRetainedCapacity + 1)
|
||||
require.False(t, shouldPoolOpsCaptureWriter(w))
|
||||
}
|
||||
|
||||
func TestEnqueueOpsErrorLog_SanitizesAndBoundsBodyBeforeQueue(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 1)
|
||||
ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
secret := strings.Repeat("s", service.OpsErrorLogQueueBodyMaxBytes)
|
||||
entry := &service.OpsInsertErrorLogInput{
|
||||
ErrorPhase: "request",
|
||||
ErrorType: "api_error",
|
||||
ErrorBody: `{"authorization":"Bearer ` + secret + `","message":"failed"}`,
|
||||
}
|
||||
|
||||
enqueueOpsErrorLog(ops, entry)
|
||||
job := <-opsErrorLogQueue
|
||||
require.LessOrEqual(t, len(job.entry.ErrorBody), service.OpsErrorLogQueueBodyMaxBytes)
|
||||
require.NotContains(t, job.entry.ErrorBody, secret)
|
||||
require.Equal(t, int64(1), OpsErrorLogSanitizedTotal())
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_DoesNotBreakOuterMiddlewares(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -149,6 +251,42 @@ func setupOpsErrorLogTestQueue(t *testing.T, size int) {
|
||||
opsErrorLogMu.Unlock()
|
||||
}
|
||||
|
||||
func TestOpsErrorLoggerMiddleware_HardSkipsIngressRejection(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 4)
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
settings := &ingressRejectSettingRepo{}
|
||||
repo := &ingressRejectOpsRepo{}
|
||||
ops := service.NewOpsService(repo, settings, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
// Construction may read unrelated runtime settings; only request-path reads matter here.
|
||||
settings.getValueCalls = 0
|
||||
|
||||
router := gin.New()
|
||||
router.Use(OpsErrorLoggerMiddleware(ops))
|
||||
router.GET("/v1/messages", func(c *gin.Context) {
|
||||
middleware2.MarkIngressRejected(c, middleware2.IngressRejectInvalidAPIKey)
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"code": "INVALID_API_KEY", "message": "Invalid API key"})
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/messages", nil)
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||||
require.JSONEq(t, `{"code":"INVALID_API_KEY","message":"Invalid API key"}`, w.Body.String())
|
||||
require.Zero(t, settings.getValueCalls, "ingress rejection must bypass monitoring settings reads")
|
||||
require.Zero(t, repo.insertCalls, "ingress rejection must bypass inserts")
|
||||
require.Zero(t, OpsErrorLogEnqueuedTotal(), "ingress rejection must not enter the error queue")
|
||||
}
|
||||
|
||||
func TestNormalizeOpsPersistentUserAgentBoundsAndPreservesUTF8(t *testing.T) {
|
||||
value := strings.Repeat("a", opsErrorLogMaxUserAgentBytes-1) + "你" + strings.Repeat("b", 32)
|
||||
got := normalizeOpsPersistentUserAgent(" " + value + " ")
|
||||
require.LessOrEqual(t, len(got), opsErrorLogMaxUserAgentBytes)
|
||||
require.True(t, utf8.ValidString(got))
|
||||
require.NotContains(t, got, "b")
|
||||
}
|
||||
|
||||
// 就地(in-band) SSE 错误挂在已固化的 HTTP 200 流上:wire 状态码为 200,
|
||||
// 常规 status>=400 采集路径不会触发。logOpsStreamError 必须据 MarkOpsStreamError
|
||||
// 补记一条错误日志,且用 IntendedStatus(429) 分级、StatusCode 仍记 wire 的 200。
|
||||
@@ -182,6 +320,36 @@ func TestLogOpsStreamError_RecordsInBandConcurrencyLimit(t *testing.T) {
|
||||
require.Equal(t, "Concurrency limit exceeded for account, please retry later", job.entry.ErrorMessage)
|
||||
}
|
||||
|
||||
func TestLogOpsStreamError_UpstreamFailureCountsTowardsSLA(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 4)
|
||||
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
c.Set(opsModelKey, "gpt-5.6-sol")
|
||||
|
||||
service.MarkOpsStreamFailure(
|
||||
c,
|
||||
"upstream_error",
|
||||
service.OpenAIUpstreamHTTP2StreamErrorCode,
|
||||
"Upstream HTTP/2 stream failed",
|
||||
http.StatusBadGateway,
|
||||
)
|
||||
|
||||
ops := service.NewOpsService(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
logOpsStreamError(c, ops, http.StatusOK)
|
||||
|
||||
job := <-opsErrorLogQueue
|
||||
require.NotNil(t, job.entry)
|
||||
require.Equal(t, http.StatusBadGateway, job.entry.StatusCode)
|
||||
require.Equal(t, "upstream_error", job.entry.ErrorType)
|
||||
require.Equal(t, "upstream", job.entry.ErrorPhase)
|
||||
require.Equal(t, "provider", job.entry.ErrorOwner)
|
||||
require.False(t, job.entry.IsBusinessLimited)
|
||||
require.Contains(t, job.entry.ErrorBody, service.OpenAIUpstreamHTTP2StreamErrorCode)
|
||||
}
|
||||
|
||||
// 未标记流内错误时 logOpsStreamError 必须是 no-op(不误记正常的 200 流)。
|
||||
func TestLogOpsStreamError_NoopWhenNotMarked(t *testing.T) {
|
||||
setupOpsErrorLogTestQueue(t, 4)
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestOpsCaptureWriterDoesNotCopyIngressRejectBody(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
context, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
writer := acquireOpsCaptureWriter(context.Writer)
|
||||
defer releaseOpsCaptureWriter(writer)
|
||||
writer.ctx = context
|
||||
context.Writer = writer
|
||||
middleware2.MarkIngressRejected(context, middleware2.IngressRejectInvalidAPIKey)
|
||||
context.Status(http.StatusUnauthorized)
|
||||
_, err := context.Writer.WriteString(`{"code":"INVALID_API_KEY","message":"Invalid API key"}`)
|
||||
require.NoError(t, err)
|
||||
require.Zero(t, writer.buf.Len())
|
||||
}
|
||||
@@ -0,0 +1,187 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/googleapi"
|
||||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
coderws "github.com/coder/websocket"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func (h *OpenAIGatewayHandler) openAISecurityAuditError(c *gin.Context, decision *securityaudit.Decision) {
|
||||
if decision == nil {
|
||||
return
|
||||
}
|
||||
if decision.Legacy != nil && decision.Legacy.Blocked {
|
||||
h.errorResponse(c, securityAuditStatus(decision), securityAuditErrorCode(decision), securityAuditMessage(decision))
|
||||
return
|
||||
}
|
||||
errType := "api_error"
|
||||
if decision.Kind == securityaudit.DecisionBlock {
|
||||
errType = "permission_error"
|
||||
}
|
||||
c.JSON(securityAuditStatus(decision), gin.H{"error": gin.H{
|
||||
"type": errType, "code": securityAuditErrorCode(decision), "message": securityAuditMessage(decision),
|
||||
}})
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) openAISecurityAuditError(c *gin.Context, decision *securityaudit.Decision) {
|
||||
if decision == nil {
|
||||
return
|
||||
}
|
||||
if decision.Legacy != nil && decision.Legacy.Blocked {
|
||||
h.chatCompletionsErrorResponse(c, securityAuditStatus(decision), securityAuditErrorCode(decision), securityAuditMessage(decision))
|
||||
return
|
||||
}
|
||||
errType := "api_error"
|
||||
if decision.Kind == securityaudit.DecisionBlock {
|
||||
errType = "permission_error"
|
||||
}
|
||||
c.JSON(securityAuditStatus(decision), gin.H{"error": gin.H{
|
||||
"type": errType, "code": securityAuditErrorCode(decision), "message": securityAuditMessage(decision),
|
||||
}})
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) responsesSecurityAuditError(c *gin.Context, decision *securityaudit.Decision) {
|
||||
if decision == nil {
|
||||
return
|
||||
}
|
||||
if decision.Legacy != nil && decision.Legacy.Blocked {
|
||||
h.responsesErrorResponse(c, securityAuditStatus(decision), securityAuditErrorCode(decision), securityAuditMessage(decision))
|
||||
return
|
||||
}
|
||||
c.JSON(securityAuditStatus(decision), gin.H{"error": gin.H{
|
||||
"type": "api_error", "code": securityAuditErrorCode(decision), "message": securityAuditMessage(decision),
|
||||
}})
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) anthropicSecurityAuditError(c *gin.Context, decision *securityaudit.Decision) {
|
||||
if decision == nil {
|
||||
return
|
||||
}
|
||||
if decision.Legacy != nil && decision.Legacy.Blocked {
|
||||
h.errorResponse(c, securityAuditStatus(decision), securityAuditErrorCode(decision), securityAuditMessage(decision))
|
||||
return
|
||||
}
|
||||
errType := "api_error"
|
||||
if decision.Kind == securityaudit.DecisionBlock {
|
||||
errType = "permission_error"
|
||||
}
|
||||
c.JSON(securityAuditStatus(decision), gin.H{"type": "error", "error": gin.H{
|
||||
"type": errType, "code": securityAuditErrorCode(decision), "message": securityAuditMessage(decision),
|
||||
}})
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) anthropicSecurityAuditError(c *gin.Context, decision *securityaudit.Decision) {
|
||||
if decision == nil {
|
||||
return
|
||||
}
|
||||
if decision.Legacy != nil && decision.Legacy.Blocked {
|
||||
h.anthropicErrorResponse(c, securityAuditStatus(decision), securityAuditErrorCode(decision), securityAuditMessage(decision))
|
||||
return
|
||||
}
|
||||
errType := "api_error"
|
||||
if decision.Kind == securityaudit.DecisionBlock {
|
||||
errType = "permission_error"
|
||||
}
|
||||
c.JSON(securityAuditStatus(decision), gin.H{"type": "error", "error": gin.H{
|
||||
"type": errType, "code": securityAuditErrorCode(decision), "message": securityAuditMessage(decision),
|
||||
}})
|
||||
}
|
||||
|
||||
func googleSecurityAuditError(c *gin.Context, decision *securityaudit.Decision) {
|
||||
if decision == nil {
|
||||
return
|
||||
}
|
||||
if decision.Legacy != nil && decision.Legacy.Blocked {
|
||||
googleError(c, securityAuditStatus(decision), securityAuditMessage(decision))
|
||||
return
|
||||
}
|
||||
status := securityAuditStatus(decision)
|
||||
googleStatus := googleapi.HTTPStatusToGoogleStatus(status)
|
||||
if status == http.StatusServiceUnavailable {
|
||||
googleStatus = "UNAVAILABLE"
|
||||
}
|
||||
requestID := ""
|
||||
if c != nil && c.Request != nil {
|
||||
requestID = contentModerationRequestID(c.Request.Context())
|
||||
}
|
||||
c.JSON(status, gin.H{"error": gin.H{
|
||||
"code": status, "message": securityAuditMessage(decision), "status": googleStatus,
|
||||
"details": []gin.H{{
|
||||
"@type": "type.googleapis.com/google.rpc.ErrorInfo",
|
||||
"reason": securityAuditErrorCode(decision), "domain": "sub2api.securityaudit",
|
||||
"metadata": gin.H{"request_id": requestID},
|
||||
}},
|
||||
}})
|
||||
}
|
||||
|
||||
func writeSecurityAuditWSError(ctx context.Context, conn *coderws.Conn, decision *securityaudit.Decision) {
|
||||
if conn == nil || decision == nil {
|
||||
return
|
||||
}
|
||||
if decision.Legacy != nil && decision.Legacy.Blocked {
|
||||
legacy := decision.Legacy
|
||||
writeContentModerationWSError(ctx, conn, (legacyContentModerationDecision{legacy}).toService())
|
||||
return
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
payload, err := json.Marshal(gin.H{
|
||||
"event_id": "evt_prompt_guard_rejected", "type": "error",
|
||||
"error": gin.H{"type": "invalid_request_error", "code": securityAuditErrorCode(decision), "message": securityAuditMessage(decision)},
|
||||
})
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
writeCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
|
||||
defer cancel()
|
||||
_ = conn.Write(writeCtx, coderws.MessageText, payload)
|
||||
}
|
||||
|
||||
type legacyContentModerationDecision struct{ value *securityaudit.LegacyDecision }
|
||||
|
||||
func (d legacyContentModerationDecision) toService() *service.ContentModerationDecision {
|
||||
if d.value == nil {
|
||||
return nil
|
||||
}
|
||||
return &service.ContentModerationDecision{Allowed: d.value.Allowed, Blocked: d.value.Blocked, Flagged: d.value.Flagged, Message: d.value.Message, StatusCode: d.value.StatusCode, Action: d.value.Action}
|
||||
}
|
||||
|
||||
func securityAuditWSCloseStatus(decision *securityaudit.Decision) coderws.StatusCode {
|
||||
if decision == nil {
|
||||
return coderws.StatusInternalError
|
||||
}
|
||||
if decision.Legacy != nil && decision.Legacy.Blocked {
|
||||
return coderws.StatusPolicyViolation
|
||||
}
|
||||
if decision.Kind == securityaudit.DecisionBlock {
|
||||
return coderws.StatusCode(4403)
|
||||
}
|
||||
return coderws.StatusTryAgainLater
|
||||
}
|
||||
|
||||
func securityAuditWSCloseReason(decision *securityaudit.Decision) string {
|
||||
if decision == nil {
|
||||
return securityaudit.ErrorCodeUnavailable
|
||||
}
|
||||
if decision.Legacy != nil && decision.Legacy.Blocked {
|
||||
message := strings.TrimSpace(decision.Legacy.Message)
|
||||
if message != "" {
|
||||
return message
|
||||
}
|
||||
return "content_policy_violation"
|
||||
}
|
||||
code := securityAuditErrorCode(decision)
|
||||
if code == "" {
|
||||
return securityaudit.ErrorCodeUnavailable
|
||||
}
|
||||
return code
|
||||
}
|
||||
@@ -0,0 +1,159 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func promptGuardDecision(kind securityaudit.DecisionKind) *securityaudit.Decision {
|
||||
decision := &securityaudit.Decision{Kind: kind, AllowNextStage: false}
|
||||
switch kind {
|
||||
case securityaudit.DecisionBlock:
|
||||
decision.HTTPStatus = http.StatusForbidden
|
||||
decision.ErrorCode = securityaudit.ErrorCodeBlocked
|
||||
decision.ClientMessage = "提示词安全审计拒绝了该请求,请调整输入后重试"
|
||||
case securityaudit.DecisionInvalid:
|
||||
decision.HTTPStatus = http.StatusServiceUnavailable
|
||||
decision.ErrorCode = securityaudit.ErrorCodeInvalidResponse
|
||||
decision.ClientMessage = "提示词安全审计暂时不可用,请稍后重试"
|
||||
default:
|
||||
decision.HTTPStatus = http.StatusServiceUnavailable
|
||||
decision.ErrorCode = securityaudit.ErrorCodeUnavailable
|
||||
decision.ClientMessage = "提示词安全审计暂时不可用,请稍后重试"
|
||||
}
|
||||
return decision
|
||||
}
|
||||
|
||||
func securityAuditErrorTestContext(t *testing.T) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
t.Helper()
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
ctx := context.WithValue(context.Background(), ctxkey.RequestID, "request-error-golden")
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/test", nil).WithContext(ctx)
|
||||
return c, recorder
|
||||
}
|
||||
|
||||
func decodeErrorJSON(t *testing.T, recorder *httptest.ResponseRecorder) map[string]any {
|
||||
t.Helper()
|
||||
var payload map[string]any
|
||||
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &payload))
|
||||
return payload
|
||||
}
|
||||
|
||||
func requireObject(t *testing.T, value any) map[string]any {
|
||||
t.Helper()
|
||||
object, ok := value.(map[string]any)
|
||||
require.True(t, ok)
|
||||
return object
|
||||
}
|
||||
|
||||
func requireArray(t *testing.T, value any) []any {
|
||||
t.Helper()
|
||||
array, ok := value.([]any)
|
||||
require.True(t, ok)
|
||||
return array
|
||||
}
|
||||
|
||||
func TestPromptGuardOpenAIAndClaudeErrorEnvelopesGolden(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, kind := range []securityaudit.DecisionKind{securityaudit.DecisionBlock, securityaudit.DecisionUnavailable, securityaudit.DecisionInvalid} {
|
||||
decision := promptGuardDecision(kind)
|
||||
t.Run("openai_"+string(kind), func(t *testing.T) {
|
||||
c, recorder := securityAuditErrorTestContext(t)
|
||||
(&OpenAIGatewayHandler{}).openAISecurityAuditError(c, decision)
|
||||
require.Equal(t, decision.HTTPStatus, recorder.Code)
|
||||
payload := decodeErrorJSON(t, recorder)
|
||||
errorObject := requireObject(t, payload["error"])
|
||||
require.Equal(t, decision.ErrorCode, errorObject["code"])
|
||||
if kind == securityaudit.DecisionBlock {
|
||||
require.Equal(t, "permission_error", errorObject["type"])
|
||||
} else {
|
||||
require.Equal(t, "api_error", errorObject["type"])
|
||||
}
|
||||
require.NotContains(t, recorder.Body.String(), "raw prompt")
|
||||
require.NotContains(t, recorder.Body.String(), "guard-one")
|
||||
})
|
||||
|
||||
t.Run("responses_"+string(kind), func(t *testing.T) {
|
||||
c, recorder := securityAuditErrorTestContext(t)
|
||||
(&GatewayHandler{}).responsesSecurityAuditError(c, decision)
|
||||
require.Equal(t, decision.HTTPStatus, recorder.Code)
|
||||
errorObject := requireObject(t, decodeErrorJSON(t, recorder)["error"])
|
||||
require.Equal(t, decision.ErrorCode, errorObject["code"])
|
||||
require.Equal(t, "api_error", errorObject["type"])
|
||||
})
|
||||
|
||||
t.Run("claude_"+string(kind), func(t *testing.T) {
|
||||
c, recorder := securityAuditErrorTestContext(t)
|
||||
(&GatewayHandler{}).anthropicSecurityAuditError(c, decision)
|
||||
require.Equal(t, decision.HTTPStatus, recorder.Code)
|
||||
payload := decodeErrorJSON(t, recorder)
|
||||
require.Equal(t, "error", payload["type"])
|
||||
errorObject := requireObject(t, payload["error"])
|
||||
require.Equal(t, decision.ErrorCode, errorObject["code"])
|
||||
if kind == securityaudit.DecisionBlock {
|
||||
require.Equal(t, "permission_error", errorObject["type"])
|
||||
} else {
|
||||
require.Equal(t, "api_error", errorObject["type"])
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPromptGuardGeminiErrorEnvelopeGolden(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, kind := range []securityaudit.DecisionKind{securityaudit.DecisionBlock, securityaudit.DecisionUnavailable, securityaudit.DecisionInvalid} {
|
||||
decision := promptGuardDecision(kind)
|
||||
c, recorder := securityAuditErrorTestContext(t)
|
||||
googleSecurityAuditError(c, decision)
|
||||
require.Equal(t, decision.HTTPStatus, recorder.Code)
|
||||
payload := decodeErrorJSON(t, recorder)
|
||||
errorObject := requireObject(t, payload["error"])
|
||||
require.Equal(t, float64(decision.HTTPStatus), errorObject["code"], "Gemini code must remain numeric")
|
||||
if decision.HTTPStatus == http.StatusForbidden {
|
||||
require.Equal(t, "PERMISSION_DENIED", errorObject["status"])
|
||||
} else {
|
||||
require.Equal(t, "UNAVAILABLE", errorObject["status"])
|
||||
}
|
||||
details := requireArray(t, errorObject["details"])
|
||||
require.Len(t, details, 1)
|
||||
errorInfo := requireObject(t, details[0])
|
||||
require.Equal(t, "type.googleapis.com/google.rpc.ErrorInfo", errorInfo["@type"])
|
||||
require.Equal(t, decision.ErrorCode, errorInfo["reason"])
|
||||
require.Equal(t, "sub2api.securityaudit", errorInfo["domain"])
|
||||
metadata := requireObject(t, errorInfo["metadata"])
|
||||
require.Equal(t, map[string]any{"request_id": "request-error-golden"}, metadata)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPromptGuardWebSocketCloseMappingGolden(t *testing.T) {
|
||||
require.Equal(t, int64(4403), int64(securityAuditWSCloseStatus(promptGuardDecision(securityaudit.DecisionBlock))))
|
||||
require.Equal(t, securityaudit.ErrorCodeBlocked, securityAuditWSCloseReason(promptGuardDecision(securityaudit.DecisionBlock)))
|
||||
require.Equal(t, int64(1013), int64(securityAuditWSCloseStatus(promptGuardDecision(securityaudit.DecisionUnavailable))))
|
||||
require.Equal(t, securityaudit.ErrorCodeUnavailable, securityAuditWSCloseReason(promptGuardDecision(securityaudit.DecisionUnavailable)))
|
||||
require.Equal(t, int64(1013), int64(securityAuditWSCloseStatus(promptGuardDecision(securityaudit.DecisionInvalid))))
|
||||
require.Equal(t, securityaudit.ErrorCodeInvalidResponse, securityAuditWSCloseReason(promptGuardDecision(securityaudit.DecisionInvalid)))
|
||||
}
|
||||
|
||||
func TestLegacyModerationErrorKeepsExistingClientPriority(t *testing.T) {
|
||||
legacy := &securityaudit.Decision{
|
||||
Kind: securityaudit.DecisionBlock, HTTPStatus: http.StatusForbidden,
|
||||
ErrorCode: "content_policy_violation", ClientMessage: "legacy exact message",
|
||||
Legacy: &securityaudit.LegacyDecision{Blocked: true, StatusCode: http.StatusForbidden, ErrorCode: "content_policy_violation", Message: "legacy exact message"},
|
||||
Prompt: &securityaudit.PromptDecision{Kind: securityaudit.DecisionBlock, ErrorCode: securityaudit.ErrorCodeBlocked},
|
||||
}
|
||||
c, recorder := securityAuditErrorTestContext(t)
|
||||
(&GatewayHandler{}).openAISecurityAuditError(c, legacy)
|
||||
require.Equal(t, http.StatusForbidden, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), "legacy exact message")
|
||||
require.Contains(t, recorder.Body.String(), "content_policy_violation")
|
||||
require.NotContains(t, recorder.Body.String(), securityaudit.ErrorCodeBlocked)
|
||||
}
|
||||
@@ -0,0 +1,153 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const securityAuditCompletedContextKey = "sub2api.security_audit.completed"
|
||||
|
||||
// cachesSecurityAuditCompletion reports whether a successful audit may be
|
||||
// reused for the rest of the gin request. WebSocket turns share one Context
|
||||
// across many response.create frames and must be audited independently.
|
||||
func cachesSecurityAuditCompletion(stage string) bool {
|
||||
switch strings.TrimSpace(stage) {
|
||||
case "", "http":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func (h *GatewayHandler) checkSecurityAudit(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol, model string, body []byte) *securityaudit.Decision {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
return runSecurityAudit(c, reqLog, h.securityAuditCoordinator, h.contentModerationService, apiKey, subject, protocol, model, body, "http")
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) checkSecurityAudit(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol, model string, body []byte) *securityaudit.Decision {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
return runSecurityAudit(c, reqLog, h.securityAuditCoordinator, h.contentModerationService, apiKey, subject, protocol, model, body, "http")
|
||||
}
|
||||
|
||||
func (h *OpenAIGatewayHandler) checkSecurityAuditStage(c *gin.Context, reqLog *zap.Logger, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol, model string, body []byte, stage string) *securityaudit.Decision {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
return runSecurityAudit(c, reqLog, h.securityAuditCoordinator, h.contentModerationService, apiKey, subject, protocol, model, body, stage)
|
||||
}
|
||||
|
||||
func runSecurityAudit(c *gin.Context, reqLog *zap.Logger, coordinator *securityaudit.Coordinator, legacy *service.ContentModerationService, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol, model string, body []byte, stage string) *securityaudit.Decision {
|
||||
if c == nil || c.Request == nil {
|
||||
return nil
|
||||
}
|
||||
cacheCompletion := cachesSecurityAuditCompletion(stage)
|
||||
if cacheCompletion {
|
||||
if completed, exists := c.Get(securityAuditCompletedContextKey); exists && completed == true {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
if coordinator == nil {
|
||||
legacyDecision := runContentModeration(c, reqLog, legacy, apiKey, subject, protocol, model, body)
|
||||
if legacyDecision == nil {
|
||||
return nil
|
||||
}
|
||||
decision := securityaudit.Decision{Kind: securityaudit.DecisionAllow, HTTPStatus: http.StatusOK, AllowNextStage: true}
|
||||
decision.Legacy = &securityaudit.LegacyDecision{
|
||||
Allowed: legacyDecision.Allowed, Blocked: legacyDecision.Blocked, Flagged: legacyDecision.Flagged,
|
||||
Message: legacyDecision.Message, StatusCode: legacyDecision.StatusCode,
|
||||
ErrorCode: "content_policy_violation", Action: legacyDecision.Action,
|
||||
}
|
||||
if legacyDecision.Blocked {
|
||||
decision.Kind, decision.HTTPStatus, decision.ErrorCode, decision.ClientMessage, decision.AllowNextStage = securityaudit.DecisionBlock, contentModerationStatus(legacyDecision), "content_policy_violation", legacyDecision.Message, false
|
||||
}
|
||||
if decision.AllowNextStage && cacheCompletion {
|
||||
c.Set(securityAuditCompletedContextKey, true)
|
||||
}
|
||||
return &decision
|
||||
}
|
||||
request := buildSecurityAuditRequest(c, apiKey, subject, protocol, model, body, stage)
|
||||
if reqLog != nil {
|
||||
reqLog.Info("security_audit.gateway_check_start",
|
||||
zap.String("request_id", request.RequestID), zap.Int64("user_id", request.UserID),
|
||||
zap.Int64("api_key_id", request.APIKeyID), zap.Int64p("group_id", request.GroupID),
|
||||
zap.String("endpoint", request.Endpoint), zap.String("provider", request.Provider),
|
||||
zap.String("protocol", request.Protocol), zap.String("model", request.Model), zap.String("stage", request.Stage),
|
||||
zap.Int("body_bytes", len(body)))
|
||||
}
|
||||
decision := coordinator.Check(c.Request.Context(), request)
|
||||
if decision.AllowNextStage && cacheCompletion {
|
||||
c.Set(securityAuditCompletedContextKey, true)
|
||||
}
|
||||
if reqLog != nil {
|
||||
reqLog.Info("security_audit.gateway_check_done",
|
||||
zap.String("request_id", request.RequestID), zap.String("decision", string(decision.Kind)),
|
||||
zap.String("error_code", decision.ErrorCode), zap.Bool("allow_next_stage", decision.AllowNextStage),
|
||||
zap.String("stage", request.Stage))
|
||||
}
|
||||
return &decision
|
||||
}
|
||||
|
||||
func buildSecurityAuditRequest(c *gin.Context, apiKey *service.APIKey, subject middleware2.AuthSubject, protocol, model string, body []byte, stage string) securityaudit.Request {
|
||||
legacy := buildContentModerationInput(c, apiKey, subject, protocol, model, body)
|
||||
request := securityaudit.Request{
|
||||
RequestID: legacy.RequestID, UserID: legacy.UserID, UserEmail: legacy.UserEmail,
|
||||
APIKeyID: legacy.APIKeyID, APIKeyName: legacy.APIKeyName, GroupID: cloneSecurityAuditGroupID(legacy.GroupID),
|
||||
GroupName: legacy.GroupName, Provider: legacy.Provider, Endpoint: legacy.Endpoint,
|
||||
Protocol: legacy.Protocol, Model: legacy.Model, Body: body, Stage: strings.TrimSpace(stage),
|
||||
}
|
||||
if apiKey != nil && apiKey.User != nil {
|
||||
request.Username = apiKey.User.Username
|
||||
if request.UserEmail == "" {
|
||||
request.UserEmail = apiKey.User.Email
|
||||
}
|
||||
}
|
||||
if request.Stage == "" {
|
||||
request.Stage = "http"
|
||||
}
|
||||
return request
|
||||
}
|
||||
|
||||
func securityAuditStatus(decision *securityaudit.Decision) int {
|
||||
if decision == nil || decision.HTTPStatus < 400 || decision.HTTPStatus > 599 {
|
||||
return http.StatusForbidden
|
||||
}
|
||||
return decision.HTTPStatus
|
||||
}
|
||||
|
||||
func securityAuditErrorCode(decision *securityaudit.Decision) string {
|
||||
if decision == nil || strings.TrimSpace(decision.ErrorCode) == "" {
|
||||
return "content_policy_violation"
|
||||
}
|
||||
return decision.ErrorCode
|
||||
}
|
||||
|
||||
func securityAuditMessage(decision *securityaudit.Decision) string {
|
||||
if decision == nil {
|
||||
return "Request blocked by content policy"
|
||||
}
|
||||
if decision.Legacy != nil && decision.Legacy.Blocked && strings.TrimSpace(decision.Legacy.Message) != "" {
|
||||
return decision.Legacy.Message
|
||||
}
|
||||
if strings.TrimSpace(decision.ClientMessage) != "" {
|
||||
return decision.ClientMessage
|
||||
}
|
||||
return "Request blocked by content policy"
|
||||
}
|
||||
|
||||
func cloneSecurityAuditGroupID(value *int64) *int64 {
|
||||
if value == nil {
|
||||
return nil
|
||||
}
|
||||
cloned := *value
|
||||
return &cloned
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestCachesSecurityAuditCompletionSkipsWebSocketStages(t *testing.T) {
|
||||
require.True(t, cachesSecurityAuditCompletion("http"))
|
||||
require.True(t, cachesSecurityAuditCompletion(""))
|
||||
require.False(t, cachesSecurityAuditCompletion("first_turn"))
|
||||
require.False(t, cachesSecurityAuditCompletion("subsequent_turn"))
|
||||
}
|
||||
|
||||
func TestRunSecurityAuditDoesNotSkipSubsequentWebSocketTurns(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
engine := &turnCountingEngine{mode: securityaudit.ModeAsync}
|
||||
coordinator := securityaudit.NewCoordinator(nil, engine)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
|
||||
|
||||
subject := middleware2.AuthSubject{UserID: 7, Concurrency: 1}
|
||||
first := runSecurityAudit(c, nil, coordinator, nil, nil, subject, "openai_responses", "gpt-test",
|
||||
[]byte(`{"type":"response.create","response":{"input":"benign"}}`), "first_turn")
|
||||
require.NotNil(t, first)
|
||||
require.True(t, first.AllowNextStage)
|
||||
require.Equal(t, int64(1), engine.enqueues.Load())
|
||||
_, cached := c.Get(securityAuditCompletedContextKey)
|
||||
require.False(t, cached, "WebSocket stages must not set the HTTP completion cache")
|
||||
|
||||
// Even if an HTTP path previously cached completion on this Context, WS turns
|
||||
// must still audit every response.create payload.
|
||||
c.Set(securityAuditCompletedContextKey, true)
|
||||
|
||||
second := runSecurityAudit(c, nil, coordinator, nil, nil, subject, "openai_responses", "gpt-test",
|
||||
[]byte(`{"type":"response.create","response":{"input":"malicious follow-up"}}`), "subsequent_turn")
|
||||
require.NotNil(t, second)
|
||||
require.Equal(t, int64(2), engine.enqueues.Load(), "subsequent WebSocket turns must be audited again")
|
||||
}
|
||||
|
||||
type turnCountingEngine struct {
|
||||
mode securityaudit.Mode
|
||||
enqueues atomic.Int64
|
||||
}
|
||||
|
||||
func (e *turnCountingEngine) EffectiveMode() securityaudit.Mode { return e.mode }
|
||||
func (e *turnCountingEngine) Enqueue(context.Context, securityaudit.Request) error {
|
||||
e.enqueues.Add(1)
|
||||
return nil
|
||||
}
|
||||
func (e *turnCountingEngine) Evaluate(context.Context, securityaudit.Request) (*securityaudit.PromptDecision, error) {
|
||||
return &securityaudit.PromptDecision{Kind: securityaudit.DecisionAllow, AllowNextStage: true}, nil
|
||||
}
|
||||
@@ -0,0 +1,208 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||||
middleware2 "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"
|
||||
)
|
||||
|
||||
type handlerPromptEngine struct {
|
||||
mu sync.Mutex
|
||||
|
||||
mode securityaudit.Mode
|
||||
decision *securityaudit.PromptDecision
|
||||
err error
|
||||
evaluated int
|
||||
enqueued int
|
||||
requests []securityaudit.Request
|
||||
}
|
||||
|
||||
func (e *handlerPromptEngine) EffectiveMode() securityaudit.Mode { return e.mode }
|
||||
func (e *handlerPromptEngine) Enqueue(_ context.Context, req securityaudit.Request) error {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
e.enqueued++
|
||||
e.requests = append(e.requests, req.Clone())
|
||||
return e.err
|
||||
}
|
||||
func (e *handlerPromptEngine) Evaluate(_ context.Context, req securityaudit.Request) (*securityaudit.PromptDecision, error) {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
e.evaluated++
|
||||
e.requests = append(e.requests, req.Clone())
|
||||
return e.decision, e.err
|
||||
}
|
||||
func (e *handlerPromptEngine) snapshot() (evaluated, enqueued int, requests []securityaudit.Request) {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
requests = make([]securityaudit.Request, len(e.requests))
|
||||
copy(requests, e.requests)
|
||||
return e.evaluated, e.enqueued, requests
|
||||
}
|
||||
|
||||
func securityAuditMediaTestMiddleware(c *gin.Context) {
|
||||
groupID := int64(3)
|
||||
user := &service.User{ID: 7, Username: "media-user", Email: "media@example.test"}
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
ID: 9, UserID: 7, User: user, Name: "media-key", GroupID: &groupID,
|
||||
Group: &service.Group{ID: groupID, Name: "media-group", Platform: service.PlatformOpenAI, AllowImageGeneration: true},
|
||||
})
|
||||
c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 7, Concurrency: 2})
|
||||
c.Next()
|
||||
}
|
||||
|
||||
func blockingHandlerPromptEngine() *handlerPromptEngine {
|
||||
return &handlerPromptEngine{mode: securityaudit.ModeBlocking, decision: &securityaudit.PromptDecision{
|
||||
Kind: securityaudit.DecisionBlock, ErrorCode: securityaudit.ErrorCodeBlocked, AllowNextStage: false,
|
||||
}}
|
||||
}
|
||||
|
||||
func TestAsyncImagePromptGuardRunsBeforeTaskCreation(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
store := &asyncImageMemoryStore{tasks: map[string]*service.ImageTaskRecord{}}
|
||||
tasks := service.NewImageTaskServiceWithUploader(store, nil, time.Hour, time.Minute)
|
||||
engine := blockingHandlerPromptEngine()
|
||||
openAI := &OpenAIGatewayHandler{securityAuditCoordinator: securityaudit.NewCoordinator(nil, engine)}
|
||||
h := &AsyncImageHandler{tasks: tasks, openAI: openAI}
|
||||
executions := 0
|
||||
h.execute = func(string, *gin.Context) { executions++ }
|
||||
|
||||
router := gin.New()
|
||||
router.Use(securityAuditMediaTestMiddleware)
|
||||
router.POST("/v1/images/generations/async", h.Submit)
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/images/generations/async", strings.NewReader(`{"model":"gpt-image-2","prompt":"blocked async prompt"}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
require.Equal(t, http.StatusForbidden, recorder.Code)
|
||||
require.Contains(t, recorder.Body.String(), securityaudit.ErrorCodeBlocked)
|
||||
require.Empty(t, store.tasks, "no asynchronous task may exist after a blocking decision")
|
||||
require.Zero(t, executions)
|
||||
evaluated, _, requests := engine.snapshot()
|
||||
require.Equal(t, 1, evaluated)
|
||||
require.Len(t, requests, 1)
|
||||
require.Contains(t, string(requests[0].Body), "blocked async prompt")
|
||||
}
|
||||
|
||||
func TestAsyncImageSuccessfulPrecheckIsNotRepeatedByDetachedExecution(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
store := &asyncImageMemoryStore{tasks: map[string]*service.ImageTaskRecord{}}
|
||||
tasks := service.NewImageTaskServiceWithUploader(store, nil, time.Hour, time.Minute)
|
||||
engine := &handlerPromptEngine{mode: securityaudit.ModeBlocking, decision: &securityaudit.PromptDecision{Kind: securityaudit.DecisionAllow, AllowNextStage: true}}
|
||||
openAI := &OpenAIGatewayHandler{securityAuditCoordinator: securityaudit.NewCoordinator(nil, engine)}
|
||||
h := &AsyncImageHandler{tasks: tasks, openAI: openAI}
|
||||
var executionMu sync.Mutex
|
||||
repeatedDecision := false
|
||||
h.execute = func(_ string, c *gin.Context) {
|
||||
apiKey, _ := middleware2.GetAPIKeyFromContext(c)
|
||||
subject, _ := middleware2.GetAuthSubjectFromContext(c)
|
||||
decision := openAI.checkSecurityAudit(c, nil, apiKey, subject, service.ContentModerationProtocolOpenAIImages, "gpt-image-2", []byte(`{"prompt":"must not rescan"}`))
|
||||
executionMu.Lock()
|
||||
repeatedDecision = decision != nil
|
||||
executionMu.Unlock()
|
||||
c.JSON(http.StatusOK, gin.H{"created": 1, "data": []any{}})
|
||||
}
|
||||
|
||||
router := gin.New()
|
||||
router.Use(securityAuditMediaTestMiddleware)
|
||||
router.POST("/v1/images/generations/async", h.Submit)
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/images/generations/async", strings.NewReader(`{"model":"gpt-image-2","prompt":"allowed async prompt"}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
require.Equal(t, http.StatusAccepted, recorder.Code)
|
||||
require.Eventually(t, func() bool {
|
||||
store.mu.RLock()
|
||||
defer store.mu.RUnlock()
|
||||
for _, task := range store.tasks {
|
||||
if task.Status == service.ImageTaskStatusCompleted {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}, time.Second, 10*time.Millisecond)
|
||||
evaluated, _, _ := engine.snapshot()
|
||||
require.Equal(t, 1, evaluated)
|
||||
executionMu.Lock()
|
||||
require.False(t, repeatedDecision)
|
||||
executionMu.Unlock()
|
||||
}
|
||||
|
||||
func TestBatchImagePromptGuardRunsBeforePersistenceOrBilling(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
engine := blockingHandlerPromptEngine()
|
||||
openAI := &OpenAIGatewayHandler{securityAuditCoordinator: securityaudit.NewCoordinator(nil, engine)}
|
||||
h := &BatchImageHandler{openAI: openAI}
|
||||
router := gin.New()
|
||||
router.Use(securityAuditMediaTestMiddleware)
|
||||
router.POST("/v1/images/batches", h.Submit)
|
||||
body := map[string]any{
|
||||
"model": "gemini-image-test",
|
||||
"items": []map[string]any{{
|
||||
"custom_id": "one", "prompt": "blocked batch prompt",
|
||||
"reference_images": []map[string]any{{"mime_type": "image/png", "data": []byte("BINARY_CANARY")}},
|
||||
}},
|
||||
}
|
||||
raw, err := json.Marshal(body)
|
||||
require.NoError(t, err)
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1/images/batches", strings.NewReader(string(raw)))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
require.NotPanics(t, func() { router.ServeHTTP(recorder, request) }, "nil service would panic if Submit were reached")
|
||||
require.Equal(t, http.StatusForbidden, recorder.Code)
|
||||
evaluated, _, requests := engine.snapshot()
|
||||
require.Equal(t, 1, evaluated)
|
||||
require.Len(t, requests, 1)
|
||||
require.Contains(t, string(requests[0].Body), "blocked batch prompt")
|
||||
require.NotContains(t, string(requests[0].Body), "BINARY_CANARY")
|
||||
require.NotContains(t, string(requests[0].Body), "QklOQVJZX0NBTkFSWQ==")
|
||||
}
|
||||
|
||||
func TestSecurityAuditBlockingFailuresLeaveAllDownstreamCountersAtZero(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, kind := range []securityaudit.DecisionKind{securityaudit.DecisionBlock, securityaudit.DecisionUnavailable, securityaudit.DecisionInvalid} {
|
||||
t.Run(string(kind), func(t *testing.T) {
|
||||
promptDecision := promptGuardDecision(kind)
|
||||
engine := &handlerPromptEngine{mode: securityaudit.ModeBlocking, decision: &securityaudit.PromptDecision{
|
||||
Kind: kind, ErrorCode: promptDecision.ErrorCode, AllowNextStage: false,
|
||||
}}
|
||||
coordinator := securityaudit.NewCoordinator(nil, engine)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"gpt-test","messages":[{"role":"user","content":"guard me"}]}`))
|
||||
groupID := int64(3)
|
||||
apiKey := &service.APIKey{ID: 9, UserID: 7, GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI}}
|
||||
subject := middleware2.AuthSubject{UserID: 7, Concurrency: 2}
|
||||
decision := runSecurityAudit(c, nil, coordinator, nil, apiKey, subject, service.ContentModerationProtocolOpenAIChat, "gpt-test", []byte(`{"messages":[{"role":"user","content":"guard me"}]}`), "http")
|
||||
require.NotNil(t, decision)
|
||||
require.False(t, decision.AllowNextStage)
|
||||
require.False(t, recorder.Result().Header.Get("Content-Type") != "", "Guard evaluation itself must not start SSE/HTTP output")
|
||||
|
||||
accountSelections, billingChecks, billingPreconsumes, upstreamDispatches := 0, 0, 0, 0
|
||||
if decision.AllowNextStage {
|
||||
accountSelections++
|
||||
billingChecks++
|
||||
billingPreconsumes++
|
||||
upstreamDispatches++
|
||||
}
|
||||
require.Zero(t, accountSelections)
|
||||
require.Zero(t, billingChecks)
|
||||
require.Zero(t, billingPreconsumes)
|
||||
require.Zero(t, upstreamDispatches)
|
||||
(&OpenAIGatewayHandler{}).openAISecurityAuditError(c, decision)
|
||||
require.Equal(t, promptDecision.HTTPStatus, recorder.Code)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"go/ast"
|
||||
"go/parser"
|
||||
"go/token"
|
||||
"os"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type promptAuditOrderCase struct {
|
||||
file string
|
||||
function string
|
||||
auditToken string
|
||||
}
|
||||
|
||||
func TestPromptAuditGatePrecedesAccountBillingAndUpstreamSideEffects(t *testing.T) {
|
||||
tests := []promptAuditOrderCase{
|
||||
{file: "gateway_handler.go", function: "Messages", auditToken: "checkSecurityAudit"},
|
||||
{file: "gateway_handler_chat_completions.go", function: "ChatCompletions", auditToken: "checkSecurityAudit"},
|
||||
{file: "gateway_handler_responses.go", function: "Responses", auditToken: "checkSecurityAudit"},
|
||||
{file: "gemini_v1beta_handler.go", function: "GeminiV1BetaModels", auditToken: "checkSecurityAudit"},
|
||||
{file: "openai_gateway_handler.go", function: "Responses", auditToken: "checkSecurityAudit"},
|
||||
{file: "openai_gateway_handler.go", function: "Messages", auditToken: "checkSecurityAudit"},
|
||||
{file: "openai_chat_completions.go", function: "ChatCompletions", auditToken: "checkSecurityAudit"},
|
||||
{file: "openai_images.go", function: "Images", auditToken: "checkSecurityAudit"},
|
||||
{file: "grok_media.go", function: "handleGrokMedia", auditToken: "checkSecurityAudit"},
|
||||
{file: "openai_embeddings.go", function: "Embeddings", auditToken: "checkSecurityAudit"},
|
||||
{file: "openai_alpha_search.go", function: "AlphaSearch", auditToken: "checkSecurityAudit"},
|
||||
{file: "image_task_handler.go", function: "Submit", auditToken: "checkSecurityAuditBeforeSubmit"},
|
||||
{file: "batch_image_handler.go", function: "Submit", auditToken: "checkSecurityAuditBeforeSubmit"},
|
||||
}
|
||||
sideEffectTokens := []string{
|
||||
"CheckBillingEligibility(", "SelectAccount", ".Forward", "acquireResponsesUserSlot(",
|
||||
"AcquireUserSlot", "TryAcquireUserSlot", "acquireImageGenerationSlot(",
|
||||
"h.tasks.Create(", "h.service.Submit(",
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.file+"/"+tt.function, func(t *testing.T) {
|
||||
functionSource := stripGoComments(goFunctionSource(t, tt.file, tt.function))
|
||||
auditIndex := strings.Index(functionSource, tt.auditToken)
|
||||
require.NotEqual(t, -1, auditIndex, "missing Prompt Audit gate")
|
||||
foundSideEffect := false
|
||||
for _, sideEffect := range sideEffectTokens {
|
||||
index := strings.Index(functionSource, sideEffect)
|
||||
if index < 0 {
|
||||
continue
|
||||
}
|
||||
foundSideEffect = true
|
||||
require.Lessf(t, auditIndex, index, "%s must run before %s", tt.auditToken, sideEffect)
|
||||
}
|
||||
require.True(t, foundSideEffect, "coverage case must contain a downstream side effect")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func stripGoComments(source string) string {
|
||||
source = regexp.MustCompile(`(?s)/\*.*?\*/`).ReplaceAllString(source, "")
|
||||
return regexp.MustCompile(`(?m)//.*$`).ReplaceAllString(source, "")
|
||||
}
|
||||
|
||||
func goFunctionSource(t *testing.T, filename, functionName string) string {
|
||||
t.Helper()
|
||||
raw, err := os.ReadFile(filename)
|
||||
require.NoError(t, err)
|
||||
files := token.NewFileSet()
|
||||
parsed, err := parser.ParseFile(files, filename, raw, 0)
|
||||
require.NoError(t, err)
|
||||
for _, declaration := range parsed.Decls {
|
||||
function, ok := declaration.(*ast.FuncDecl)
|
||||
if !ok || function.Name.Name != functionName || function.Body == nil {
|
||||
continue
|
||||
}
|
||||
start := files.Position(function.Pos()).Offset
|
||||
end := files.Position(function.End()).Offset
|
||||
require.Greater(t, end, start)
|
||||
return string(raw[start:end])
|
||||
}
|
||||
t.Fatalf("function %s not found in %s", functionName, filename)
|
||||
return ""
|
||||
}
|
||||
@@ -93,6 +93,9 @@ func (s *userHandlerRepoStub) BatchSetConcurrency(context.Context, []int64, int)
|
||||
func (s *userHandlerRepoStub) BatchAddConcurrency(context.Context, []int64, int) (int, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (s *userHandlerRepoStub) BatchUpdateLimits(context.Context, []int64, *int, *int) (int, error) {
|
||||
return 0, nil
|
||||
}
|
||||
func (s *userHandlerRepoStub) ExistsByEmail(context.Context, string) (bool, error) { return false, nil }
|
||||
func (s *userHandlerRepoStub) RemoveGroupFromAllowedGroups(context.Context, int64) (int64, error) {
|
||||
return 0, nil
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/admin"
|
||||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/google/wire"
|
||||
@@ -38,6 +40,7 @@ func ProvideAdminHandlers(
|
||||
channelMonitorHandler *admin.ChannelMonitorHandler,
|
||||
channelMonitorTemplateHandler *admin.ChannelMonitorRequestTemplateHandler,
|
||||
contentModerationHandler *admin.ContentModerationHandler,
|
||||
promptAuditHandler *securityaudit.PromptAdminHandler,
|
||||
paymentHandler *admin.PaymentHandler,
|
||||
affiliateHandler *admin.AffiliateHandler,
|
||||
complianceHandler *admin.ComplianceHandler,
|
||||
@@ -75,6 +78,7 @@ func ProvideAdminHandlers(
|
||||
ChannelMonitor: channelMonitorHandler,
|
||||
ChannelMonitorTemplate: channelMonitorTemplateHandler,
|
||||
ContentModeration: contentModerationHandler,
|
||||
PromptAudit: promptAuditHandler,
|
||||
Payment: paymentHandler,
|
||||
Affiliate: affiliateHandler,
|
||||
Compliance: complianceHandler,
|
||||
@@ -82,6 +86,62 @@ func ProvideAdminHandlers(
|
||||
}
|
||||
}
|
||||
|
||||
func ProvideGatewayHandler(
|
||||
gatewayService *service.GatewayService,
|
||||
openAIGatewayService *service.OpenAIGatewayService,
|
||||
geminiCompatService *service.GeminiMessagesCompatService,
|
||||
antigravityGatewayService *service.AntigravityGatewayService,
|
||||
userService *service.UserService,
|
||||
concurrencyService *service.ConcurrencyService,
|
||||
billingCacheService *service.BillingCacheService,
|
||||
usageService *service.UsageService,
|
||||
apiKeyService *service.APIKeyService,
|
||||
usageRecordWorkerPool *service.UsageRecordWorkerPool,
|
||||
errorPassthroughService *service.ErrorPassthroughService,
|
||||
contentModerationService *service.ContentModerationService,
|
||||
userMsgQueueService *service.UserMessageQueueService,
|
||||
cfg *config.Config,
|
||||
settingService *service.SettingService,
|
||||
coordinator *securityaudit.Coordinator,
|
||||
) *GatewayHandler {
|
||||
h := NewGatewayHandler(gatewayService, openAIGatewayService, geminiCompatService, antigravityGatewayService,
|
||||
userService, concurrencyService, billingCacheService, usageService, apiKeyService, usageRecordWorkerPool,
|
||||
errorPassthroughService, contentModerationService, userMsgQueueService, cfg, settingService)
|
||||
h.securityAuditCoordinator = coordinator
|
||||
return h
|
||||
}
|
||||
|
||||
func ProvideOpenAIGatewayHandler(
|
||||
gatewayService *service.OpenAIGatewayService,
|
||||
concurrencyService *service.ConcurrencyService,
|
||||
billingCacheService *service.BillingCacheService,
|
||||
apiKeyService *service.APIKeyService,
|
||||
usageRecordWorkerPool *service.UsageRecordWorkerPool,
|
||||
errorPassthroughService *service.ErrorPassthroughService,
|
||||
contentModerationService *service.ContentModerationService,
|
||||
opsService *service.OpsService,
|
||||
grokQuotaService *service.GrokQuotaService,
|
||||
cfg *config.Config,
|
||||
coordinator *securityaudit.Coordinator,
|
||||
) *OpenAIGatewayHandler {
|
||||
h := NewOpenAIGatewayHandler(gatewayService, concurrencyService, billingCacheService, apiKeyService,
|
||||
usageRecordWorkerPool, errorPassthroughService, contentModerationService, opsService, cfg)
|
||||
h.securityAuditCoordinator = coordinator
|
||||
h.grokMediaEligibilityProber = grokQuotaService
|
||||
return h
|
||||
}
|
||||
|
||||
func ProvideBatchImageHandler(
|
||||
batchService *service.BatchImagePublicService,
|
||||
download *service.BatchImageDownloadService,
|
||||
cleanup *service.BatchImageCleanupService,
|
||||
openAI *OpenAIGatewayHandler,
|
||||
) *BatchImageHandler {
|
||||
h := NewBatchImageHandler(batchService, download, cleanup)
|
||||
h.openAI = openAI
|
||||
return h
|
||||
}
|
||||
|
||||
// ProvideSystemHandler creates admin.SystemHandler with UpdateService
|
||||
func ProvideSystemHandler(updateService *service.UpdateService, lockService *service.SystemOperationLockService) *admin.SystemHandler {
|
||||
return admin.NewSystemHandler(updateService, lockService)
|
||||
@@ -95,9 +155,10 @@ func ProvideSettingHandler(settingService *service.SettingService, buildInfo Bui
|
||||
}
|
||||
|
||||
// ProvideAdminSettingHandler creates admin.SettingHandler with notification template APIs.
|
||||
func ProvideAdminSettingHandler(settingService *service.SettingService, emailService *service.EmailService, turnstileService *service.TurnstileService, opsService *service.OpsService, paymentConfigService *service.PaymentConfigService, paymentService *service.PaymentService, userAttributeService *service.UserAttributeService, notificationEmailService *service.NotificationEmailService) *admin.SettingHandler {
|
||||
func ProvideAdminSettingHandler(settingService *service.SettingService, emailService *service.EmailService, turnstileService *service.TurnstileService, opsService *service.OpsService, paymentConfigService *service.PaymentConfigService, paymentService *service.PaymentService, userAttributeService *service.UserAttributeService, notificationEmailService *service.NotificationEmailService, totpService *service.TotpService, userService *service.UserService) *admin.SettingHandler {
|
||||
h := admin.NewSettingHandler(settingService, emailService, turnstileService, opsService, paymentConfigService, paymentService, userAttributeService)
|
||||
h.SetNotificationEmailService(notificationEmailService)
|
||||
h.SetStepUpDeps(totpService, userService)
|
||||
return h
|
||||
}
|
||||
|
||||
@@ -157,15 +218,15 @@ var ProviderSet = wire.NewSet(
|
||||
NewSubscriptionHandler,
|
||||
NewAnnouncementHandler,
|
||||
NewChannelMonitorUserHandler,
|
||||
NewGatewayHandler,
|
||||
NewOpenAIGatewayHandler,
|
||||
ProvideGatewayHandler,
|
||||
ProvideOpenAIGatewayHandler,
|
||||
NewTotpHandler,
|
||||
ProvideSettingHandler,
|
||||
NewPaymentHandler,
|
||||
NewPaymentWebhookHandler,
|
||||
NewAvailableChannelHandler,
|
||||
NewAsyncImageHandler,
|
||||
NewBatchImageHandler,
|
||||
ProvideBatchImageHandler,
|
||||
|
||||
// Admin handlers
|
||||
admin.NewDashboardHandler,
|
||||
|
||||
@@ -152,11 +152,25 @@ type AnthropicEventToResponsesState struct {
|
||||
|
||||
// For message output: accumulate text parts
|
||||
ContentIndex int
|
||||
// TextAccum accumulates the current text part so that output_text.done and
|
||||
// content_part.done can carry the full text (deltas carry increments only).
|
||||
TextAccum string
|
||||
|
||||
// For function_call: track per-output info
|
||||
CurrentCallID string
|
||||
CurrentName string
|
||||
|
||||
// Content of the currently open item, folded into Outputs when it closes.
|
||||
CurrentContent []ResponsesContentPart // message
|
||||
CurrentArgs string // function_call
|
||||
CurrentSummary string // reasoning
|
||||
|
||||
// Outputs accumulates every closed output item so that response.completed
|
||||
// can carry the full output list. The OpenAI SDK's get_final_response()
|
||||
// parses the terminal event's response directly; without this, clients see
|
||||
// an empty output_text.
|
||||
Outputs []ResponsesOutput
|
||||
|
||||
// Usage from message_start / message_delta. InputTokens here follows
|
||||
// Anthropic semantics (excludes cached tokens); they are added back when
|
||||
// emitting the OpenAI Responses usage.
|
||||
@@ -293,6 +307,22 @@ func anthToResHandleContentBlockStart(evt *AnthropicStreamEvent, state *Anthropi
|
||||
}))
|
||||
}
|
||||
|
||||
// response.content_part.added must precede the output_text.delta events
|
||||
// for that part. The message item is added with content: [], and the
|
||||
// OpenAI SDK's accumulating stream helper (client.responses.stream) only
|
||||
// appends a content part when it sees content_part.added. Without it the
|
||||
// following output_text.delta indexes output.content[content_index] and
|
||||
// raises IndexError. Raw event iteration
|
||||
// (responses.create(stream=True)) does not accumulate, which is why this
|
||||
// went unnoticed.
|
||||
events = append(events, makeResponsesEvent(state, "response.content_part.added", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
ContentIndex: state.ContentIndex,
|
||||
ItemID: state.CurrentItemID,
|
||||
Part: &ResponsesContentPart{Type: "output_text", Text: ""},
|
||||
}))
|
||||
state.TextAccum = ""
|
||||
|
||||
case "tool_use":
|
||||
// Close previous item if any
|
||||
events = append(events, closeCurrentResponsesItem(state)...)
|
||||
@@ -327,6 +357,7 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi
|
||||
if evt.Delta.Text == "" {
|
||||
return nil
|
||||
}
|
||||
state.TextAccum += evt.Delta.Text
|
||||
return []ResponsesStreamEvent{makeResponsesEvent(state, "response.output_text.delta", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
ContentIndex: state.ContentIndex,
|
||||
@@ -338,6 +369,7 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi
|
||||
if evt.Delta.Thinking == "" {
|
||||
return nil
|
||||
}
|
||||
state.CurrentSummary += evt.Delta.Thinking
|
||||
return []ResponsesStreamEvent{makeResponsesEvent(state, "response.reasoning_summary_text.delta", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
SummaryIndex: 0,
|
||||
@@ -349,6 +381,7 @@ func anthToResHandleContentBlockDelta(evt *AnthropicStreamEvent, state *Anthropi
|
||||
if evt.Delta.PartialJSON == "" {
|
||||
return nil
|
||||
}
|
||||
state.CurrentArgs += evt.Delta.PartialJSON
|
||||
return []ResponsesStreamEvent{makeResponsesEvent(state, "response.function_call_arguments.delta", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
Delta: evt.Delta.PartialJSON,
|
||||
@@ -393,12 +426,24 @@ func anthToResHandleContentBlockStop(evt *AnthropicStreamEvent, state *Anthropic
|
||||
return events
|
||||
|
||||
case "message":
|
||||
// Emit output_text.done (text block is done, but message item stays open for potential more blocks)
|
||||
// Text block is done: emit output_text.done then content_part.done (the
|
||||
// order OpenAI uses), both carrying the part's full text. The message
|
||||
// item itself stays open since more blocks may follow.
|
||||
text := state.TextAccum
|
||||
state.TextAccum = ""
|
||||
state.CurrentContent = append(state.CurrentContent, ResponsesContentPart{Type: "output_text", Text: text})
|
||||
return []ResponsesStreamEvent{
|
||||
makeResponsesEvent(state, "response.output_text.done", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
ContentIndex: state.ContentIndex,
|
||||
ItemID: state.CurrentItemID,
|
||||
Text: text,
|
||||
}),
|
||||
makeResponsesEvent(state, "response.content_part.done", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex,
|
||||
ContentIndex: state.ContentIndex,
|
||||
ItemID: state.CurrentItemID,
|
||||
Part: &ResponsesContentPart{Type: "output_text", Text: text},
|
||||
}),
|
||||
}
|
||||
}
|
||||
@@ -454,24 +499,48 @@ func closeCurrentResponsesItem(state *AnthropicEventToResponsesState) []Response
|
||||
return nil
|
||||
}
|
||||
|
||||
itemType := state.CurrentItemType
|
||||
itemID := state.CurrentItemID
|
||||
// Assemble the full item: both output_item.done and response.completed must
|
||||
// carry its content. Emitting only {type,id,status} makes SDK-side
|
||||
// accumulation produce an empty output.
|
||||
item := ResponsesOutput{
|
||||
Type: state.CurrentItemType,
|
||||
ID: state.CurrentItemID,
|
||||
Status: "completed",
|
||||
}
|
||||
switch state.CurrentItemType {
|
||||
case "message":
|
||||
item.Role = "assistant"
|
||||
item.Content = state.CurrentContent
|
||||
case "function_call":
|
||||
item.CallID = state.CurrentCallID
|
||||
item.Name = state.CurrentName
|
||||
args := state.CurrentArgs
|
||||
if args == "" {
|
||||
args = "{}"
|
||||
}
|
||||
item.Arguments = args
|
||||
case "reasoning":
|
||||
if state.CurrentSummary != "" {
|
||||
item.Summary = []ResponsesSummary{{Type: "summary_text", Text: state.CurrentSummary}}
|
||||
}
|
||||
}
|
||||
state.Outputs = append(state.Outputs, item)
|
||||
|
||||
// Reset
|
||||
state.CurrentItemType = ""
|
||||
state.CurrentItemID = ""
|
||||
state.CurrentCallID = ""
|
||||
state.CurrentName = ""
|
||||
state.CurrentContent = nil
|
||||
state.CurrentArgs = ""
|
||||
state.CurrentSummary = ""
|
||||
state.TextAccum = ""
|
||||
state.OutputIndex++
|
||||
state.ContentIndex = 0
|
||||
|
||||
return []ResponsesStreamEvent{makeResponsesEvent(state, "response.output_item.done", &ResponsesStreamEvent{
|
||||
OutputIndex: state.OutputIndex - 1, // Use the index before increment
|
||||
Item: &ResponsesOutput{
|
||||
Type: itemType,
|
||||
ID: itemID,
|
||||
Status: "completed",
|
||||
},
|
||||
Item: &item,
|
||||
})}
|
||||
}
|
||||
|
||||
@@ -519,6 +588,14 @@ func makeResponsesCompletedEvent(
|
||||
eventType = "response.incomplete"
|
||||
}
|
||||
|
||||
// Carry the output items accumulated over the stream. The SDK's
|
||||
// get_final_response() reads them straight from the terminal event, so an
|
||||
// empty list leaves clients with an empty result.
|
||||
outputs := state.Outputs
|
||||
if outputs == nil {
|
||||
outputs = []ResponsesOutput{}
|
||||
}
|
||||
|
||||
return ResponsesStreamEvent{
|
||||
Type: eventType,
|
||||
SequenceNumber: seq,
|
||||
@@ -527,7 +604,7 @@ func makeResponsesCompletedEvent(
|
||||
Object: "response",
|
||||
Model: state.Model,
|
||||
Status: status,
|
||||
Output: []ResponsesOutput{},
|
||||
Output: outputs,
|
||||
Usage: usage,
|
||||
IncompleteDetails: incompleteDetails,
|
||||
},
|
||||
|
||||
@@ -0,0 +1,187 @@
|
||||
package apicompat
|
||||
|
||||
import "testing"
|
||||
|
||||
// TestAnthropicEventToResponses_TextEmitsContentPart pins that a message text
|
||||
// stream emits response.content_part.added, and that it precedes the first
|
||||
// output_text.delta for that part.
|
||||
//
|
||||
// Why: the OpenAI SDK's accumulating stream helper (client.responses.stream)
|
||||
// only appends a content part to the message item when it sees
|
||||
// content_part.added. The item is added with content: [], so a missing event
|
||||
// makes the following output_text.delta index output.content[content_index] and
|
||||
// raise IndexError. Raw event iteration does not accumulate, so a regression
|
||||
// here is easy to miss.
|
||||
func TestAnthropicEventToResponses_TextEmitsContentPart(t *testing.T) {
|
||||
state := NewAnthropicEventToResponsesState()
|
||||
state.Model = "claude-sonnet-4-5"
|
||||
|
||||
var types []string
|
||||
feed := func(evt *AnthropicStreamEvent) {
|
||||
for _, out := range AnthropicEventToResponsesEvents(evt, state) {
|
||||
types = append(types, out.Type)
|
||||
}
|
||||
}
|
||||
|
||||
idx := 0
|
||||
feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1", Model: "claude-sonnet-4-5"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{Type: "text"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "Hel"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "lo"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx})
|
||||
feed(&AnthropicStreamEvent{Type: "message_stop"})
|
||||
|
||||
posOf := func(target string) int {
|
||||
for i, ty := range types {
|
||||
if ty == target {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
partAdded := posOf("response.content_part.added")
|
||||
firstDelta := posOf("response.output_text.delta")
|
||||
|
||||
if partAdded < 0 {
|
||||
t.Fatalf("response.content_part.added was not emitted; got %v", types)
|
||||
}
|
||||
if firstDelta < 0 {
|
||||
t.Fatalf("response.output_text.delta was not emitted; got %v", types)
|
||||
}
|
||||
if partAdded > firstDelta {
|
||||
t.Errorf("content_part.added must precede the first output_text.delta; got %v", types)
|
||||
}
|
||||
if posOf("response.content_part.done") < 0 {
|
||||
t.Errorf("response.content_part.done was not emitted; got %v", types)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicEventToResponses_DoneEventsCarryFullText pins that done events
|
||||
// carry the part's full text (deltas carry increments only).
|
||||
func TestAnthropicEventToResponses_DoneEventsCarryFullText(t *testing.T) {
|
||||
state := NewAnthropicEventToResponsesState()
|
||||
state.Model = "claude-sonnet-4-5"
|
||||
|
||||
var events []ResponsesStreamEvent
|
||||
feed := func(evt *AnthropicStreamEvent) {
|
||||
events = append(events, AnthropicEventToResponsesEvents(evt, state)...)
|
||||
}
|
||||
|
||||
idx := 0
|
||||
feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{Type: "text"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "Hello "}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "world"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx})
|
||||
|
||||
const want = "Hello world"
|
||||
var sawTextDone, sawPartDone bool
|
||||
for _, e := range events {
|
||||
switch e.Type {
|
||||
case "response.output_text.done":
|
||||
sawTextDone = true
|
||||
if e.Text != want {
|
||||
t.Errorf("output_text.done text = %q, want %q", e.Text, want)
|
||||
}
|
||||
case "response.content_part.done":
|
||||
sawPartDone = true
|
||||
if e.Part == nil || e.Part.Text != want {
|
||||
t.Errorf("content_part.done part = %+v, want text %q", e.Part, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !sawTextDone || !sawPartDone {
|
||||
t.Errorf("missing done events: output_text.done=%v content_part.done=%v", sawTextDone, sawPartDone)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicEventToResponses_CompletedCarriesOutput pins that
|
||||
// response.completed carries the full output list. The SDK's
|
||||
// get_final_response() and tracing integrations parse the terminal event's
|
||||
// response directly; an empty output leaves them with nothing (the text still
|
||||
// renders from deltas, which is why this is invisible when only watching the
|
||||
// stream).
|
||||
func TestAnthropicEventToResponses_CompletedCarriesOutput(t *testing.T) {
|
||||
state := NewAnthropicEventToResponsesState()
|
||||
state.Model = "claude-sonnet-4-5"
|
||||
|
||||
var events []ResponsesStreamEvent
|
||||
feed := func(evt *AnthropicStreamEvent) {
|
||||
events = append(events, AnthropicEventToResponsesEvents(evt, state)...)
|
||||
}
|
||||
|
||||
idx := 0
|
||||
feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{Type: "text"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{Type: "text_delta", Text: "4826"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx})
|
||||
feed(&AnthropicStreamEvent{Type: "message_stop"})
|
||||
|
||||
var completed *ResponsesStreamEvent
|
||||
for i := range events {
|
||||
if events[i].Type == "response.completed" {
|
||||
completed = &events[i]
|
||||
}
|
||||
}
|
||||
if completed == nil || completed.Response == nil {
|
||||
t.Fatalf("response.completed was not emitted")
|
||||
}
|
||||
if len(completed.Response.Output) == 0 {
|
||||
t.Fatalf("response.completed carries an empty output; clients would see no result")
|
||||
}
|
||||
msg := completed.Response.Output[0]
|
||||
if msg.Type != "message" || len(msg.Content) == 0 {
|
||||
t.Fatalf("output[0] = %+v, want a message with content", msg)
|
||||
}
|
||||
if msg.Content[0].Text != "4826" {
|
||||
t.Errorf("output[0].content[0].text = %q, want %q", msg.Content[0].Text, "4826")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicEventToResponses_ToolCallCompletedCarriesArguments pins that a
|
||||
// function call's accumulated arguments survive into output_item.done and
|
||||
// response.completed.
|
||||
func TestAnthropicEventToResponses_ToolCallCompletedCarriesArguments(t *testing.T) {
|
||||
state := NewAnthropicEventToResponsesState()
|
||||
state.Model = "claude-sonnet-4-5"
|
||||
|
||||
var events []ResponsesStreamEvent
|
||||
feed := func(evt *AnthropicStreamEvent) {
|
||||
events = append(events, AnthropicEventToResponsesEvents(evt, state)...)
|
||||
}
|
||||
|
||||
idx := 0
|
||||
feed(&AnthropicStreamEvent{Type: "message_start", Message: &AnthropicResponse{ID: "msg_1"}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_start", Index: &idx, ContentBlock: &AnthropicContentBlock{
|
||||
Type: "tool_use", ID: "toolu_1", Name: "get_weather",
|
||||
}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{
|
||||
Type: "input_json_delta", PartialJSON: `{"city":`,
|
||||
}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_delta", Index: &idx, Delta: &AnthropicDelta{
|
||||
Type: "input_json_delta", PartialJSON: `"SH"}`,
|
||||
}})
|
||||
feed(&AnthropicStreamEvent{Type: "content_block_stop", Index: &idx})
|
||||
feed(&AnthropicStreamEvent{Type: "message_stop"})
|
||||
|
||||
var completed *ResponsesStreamEvent
|
||||
for i := range events {
|
||||
if events[i].Type == "response.completed" {
|
||||
completed = &events[i]
|
||||
}
|
||||
}
|
||||
if completed == nil || completed.Response == nil || len(completed.Response.Output) == 0 {
|
||||
t.Fatalf("response.completed carries no output")
|
||||
}
|
||||
fc := completed.Response.Output[0]
|
||||
if fc.Type != "function_call" {
|
||||
t.Fatalf("output[0].type = %q, want function_call", fc.Type)
|
||||
}
|
||||
if fc.Arguments != `{"city":"SH"}` {
|
||||
t.Errorf("arguments = %q, want %q", fc.Arguments, `{"city":"SH"}`)
|
||||
}
|
||||
if fc.Name != "get_weather" {
|
||||
t.Errorf("name = %q, want get_weather", fc.Name)
|
||||
}
|
||||
}
|
||||
@@ -8,40 +8,10 @@ import (
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// GetClientIP 从 Gin Context 中提取客户端真实 IP 地址。
|
||||
// 按以下优先级检查 Header:
|
||||
// 1. CF-Connecting-IP (Cloudflare)
|
||||
// 2. X-Real-IP (Nginx)
|
||||
// 3. X-Forwarded-For (取第一个非私有 IP)
|
||||
// 4. c.ClientIP() (Gin 内置方法)
|
||||
// GetClientIP resolves a client address only through Gin's configured trusted
|
||||
// proxy chain. Forwarding headers from a direct or untrusted peer are ignored.
|
||||
func GetClientIP(c *gin.Context) string {
|
||||
// 1. Cloudflare
|
||||
if ip := c.GetHeader("CF-Connecting-IP"); ip != "" {
|
||||
return normalizeIP(ip)
|
||||
}
|
||||
|
||||
// 2. Nginx X-Real-IP
|
||||
if ip := c.GetHeader("X-Real-IP"); ip != "" {
|
||||
return normalizeIP(ip)
|
||||
}
|
||||
|
||||
// 3. X-Forwarded-For (多个 IP 时取第一个公网 IP)
|
||||
if xff := c.GetHeader("X-Forwarded-For"); xff != "" {
|
||||
ips := strings.Split(xff, ",")
|
||||
for _, ip := range ips {
|
||||
ip = strings.TrimSpace(ip)
|
||||
if ip != "" && !isPrivateIP(ip) {
|
||||
return normalizeIP(ip)
|
||||
}
|
||||
}
|
||||
// 如果都是私有 IP,返回第一个
|
||||
if len(ips) > 0 {
|
||||
return normalizeIP(strings.TrimSpace(ips[0]))
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Gin 内置方法
|
||||
return normalizeIP(c.ClientIP())
|
||||
return GetTrustedClientIP(c)
|
||||
}
|
||||
|
||||
// GetTrustedClientIP 从 Gin 的可信代理解析链提取客户端 IP。
|
||||
@@ -54,6 +24,13 @@ func GetTrustedClientIP(c *gin.Context) string {
|
||||
return normalizeIP(c.ClientIP())
|
||||
}
|
||||
|
||||
// GetSecurityClientIP returns the address resolved through Gin's configured
|
||||
// trusted-proxy chain. The legacy toggle is retained for configuration/API
|
||||
// compatibility, but never makes raw forwarding headers trustworthy by itself.
|
||||
func GetSecurityClientIP(c *gin.Context, _ bool) string {
|
||||
return GetTrustedClientIP(c)
|
||||
}
|
||||
|
||||
// normalizeIP 规范化 IP 地址,去除端口号和空格。
|
||||
func normalizeIP(ip string) string {
|
||||
ip = strings.TrimSpace(ip)
|
||||
@@ -64,9 +41,6 @@ func normalizeIP(ip string) string {
|
||||
return ip
|
||||
}
|
||||
|
||||
// privateNets 预编译私有 IP CIDR 块,避免每次调用 isPrivateIP 时重复解析
|
||||
var privateNets []*net.IPNet
|
||||
|
||||
// CompiledIPRules 表示预编译的 IP 匹配规则。
|
||||
// PatternCount 记录原始规则数量,用于保留“规则存在但全无效”时的行为语义。
|
||||
type CompiledIPRules struct {
|
||||
@@ -75,23 +49,6 @@ type CompiledIPRules struct {
|
||||
PatternCount int
|
||||
}
|
||||
|
||||
func init() {
|
||||
for _, cidr := range []string{
|
||||
"10.0.0.0/8",
|
||||
"172.16.0.0/12",
|
||||
"192.168.0.0/16",
|
||||
"127.0.0.0/8",
|
||||
"::1/128",
|
||||
"fc00::/7",
|
||||
} {
|
||||
_, block, err := net.ParseCIDR(cidr)
|
||||
if err != nil {
|
||||
panic("invalid CIDR: " + cidr)
|
||||
}
|
||||
privateNets = append(privateNets, block)
|
||||
}
|
||||
}
|
||||
|
||||
// CompileIPRules 将 IP/CIDR 字符串规则预编译为可复用结构。
|
||||
// 非法规则会被忽略,但 PatternCount 会保留原始规则条数。
|
||||
func CompileIPRules(patterns []string) *CompiledIPRules {
|
||||
@@ -139,20 +96,6 @@ func matchesCompiledRules(parsedIP net.IP, rules *CompiledIPRules) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// isPrivateIP 检查 IP 是否为私有地址。
|
||||
func isPrivateIP(ipStr string) bool {
|
||||
ip := net.ParseIP(ipStr)
|
||||
if ip == nil {
|
||||
return false
|
||||
}
|
||||
for _, block := range privateNets {
|
||||
if block.Contains(ip) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// MatchesPattern 检查 IP 是否匹配指定的模式(支持单个 IP 或 CIDR)。
|
||||
// pattern 可以是:
|
||||
// - 单个 IP: "192.168.1.100"
|
||||
|
||||
@@ -10,48 +10,6 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestIsPrivateIP(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
ip string
|
||||
expected bool
|
||||
}{
|
||||
// 私有 IPv4
|
||||
{"10.x 私有地址", "10.0.0.1", true},
|
||||
{"10.x 私有地址段末", "10.255.255.255", true},
|
||||
{"172.16.x 私有地址", "172.16.0.1", true},
|
||||
{"172.31.x 私有地址", "172.31.255.255", true},
|
||||
{"192.168.x 私有地址", "192.168.1.1", true},
|
||||
{"127.0.0.1 本地回环", "127.0.0.1", true},
|
||||
{"127.x 回环段", "127.255.255.255", true},
|
||||
|
||||
// 公网 IPv4
|
||||
{"8.8.8.8 公网 DNS", "8.8.8.8", false},
|
||||
{"1.1.1.1 公网", "1.1.1.1", false},
|
||||
{"172.15.255.255 非私有", "172.15.255.255", false},
|
||||
{"172.32.0.0 非私有", "172.32.0.0", false},
|
||||
{"11.0.0.1 公网", "11.0.0.1", false},
|
||||
|
||||
// IPv6
|
||||
{"::1 IPv6 回环", "::1", true},
|
||||
{"fc00:: IPv6 私有", "fc00::1", true},
|
||||
{"fd00:: IPv6 私有", "fd00::1", true},
|
||||
{"2001:db8::1 IPv6 公网", "2001:db8::1", false},
|
||||
|
||||
// 无效输入
|
||||
{"空字符串", "", false},
|
||||
{"非法字符串", "not-an-ip", false},
|
||||
{"不完整 IP", "192.168", false},
|
||||
}
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := isPrivateIP(tc.ip)
|
||||
require.Equal(t, tc.expected, got, "isPrivateIP(%q)", tc.ip)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetTrustedClientIPUsesGinClientIP(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
@@ -94,3 +52,48 @@ func TestCheckIPRestrictionWithCompiledRules_InvalidWhitelistStillDenies(t *test
|
||||
require.False(t, allowed)
|
||||
require.Equal(t, "access denied", reason)
|
||||
}
|
||||
|
||||
func TestGetSecurityClientIPNeverTrustsHeadersFromUntrustedPeer(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
trustForwarded bool
|
||||
want string
|
||||
}{
|
||||
{name: "legacy toggle disabled", trustForwarded: false, want: "9.9.9.9"},
|
||||
{name: "legacy toggle enabled", trustForwarded: true, want: "9.9.9.9"},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies(nil))
|
||||
r.GET("/t", func(c *gin.Context) {
|
||||
c.String(200, GetSecurityClientIP(c, tc.trustForwarded))
|
||||
})
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Real-IP", "1.2.3.4")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, 200, w.Code)
|
||||
require.Equal(t, tc.want, w.Body.String())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetSecurityClientIPUsesConfiguredTrustedProxy(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
r := gin.New()
|
||||
require.NoError(t, r.SetTrustedProxies([]string{"9.9.9.9"}))
|
||||
r.GET("/t", func(c *gin.Context) { c.String(200, GetSecurityClientIP(c, true)) })
|
||||
|
||||
w := httptest.NewRecorder()
|
||||
req := httptest.NewRequest("GET", "/t", nil)
|
||||
req.RemoteAddr = "9.9.9.9:12345"
|
||||
req.Header.Set("X-Forwarded-For", "1.2.3.4")
|
||||
r.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, "1.2.3.4", w.Body.String())
|
||||
}
|
||||
|
||||
@@ -26,6 +26,10 @@ const (
|
||||
LevelWarn = zapcore.WarnLevel
|
||||
LevelError = zapcore.ErrorLevel
|
||||
LevelFatal = zapcore.FatalLevel
|
||||
|
||||
// OpsSystemLogSkipField keeps an event in the standard logger while
|
||||
// preventing the database-backed Ops system-log sink from indexing it.
|
||||
OpsSystemLogSkipField = "ops_system_log_skip"
|
||||
)
|
||||
|
||||
type Sink interface {
|
||||
|
||||
@@ -76,6 +76,8 @@ type BillingSummary struct {
|
||||
UsedPercent *float64 `json:"used_percent,omitempty"`
|
||||
Plan string `json:"plan,omitempty"` // SuperGrok | SuperGrok Heavy | ""
|
||||
StatusCode int `json:"status_code,omitempty"`
|
||||
WeeklyStatusCode int `json:"weekly_status_code,omitempty"`
|
||||
MonthlyStatusCode int `json:"monthly_status_code,omitempty"`
|
||||
Source string `json:"source,omitempty"`
|
||||
FetchedAt string `json:"fetched_at,omitempty"`
|
||||
UpdatedAt string `json:"updated_at,omitempty"`
|
||||
|
||||
@@ -41,7 +41,9 @@ const (
|
||||
|
||||
var (
|
||||
oauthEndpointAllowedHosts = []string{"x.ai", "*.x.ai"}
|
||||
baseURLAllowedHosts = []string{"api.x.ai", "cli-chat-proxy.grok.com"}
|
||||
// *.api.x.ai 覆盖 xAI 区域端点(us-east-1/us-west-2/eu-west-1 等),
|
||||
// 运营方可在端点间手动切换以规避单点不可用。
|
||||
baseURLAllowedHosts = []string{"api.x.ai", "*.api.x.ai", "cli-chat-proxy.grok.com"}
|
||||
)
|
||||
|
||||
// OAuthSession stores one PKCE OAuth flow.
|
||||
@@ -332,10 +334,17 @@ func normalizeKnownBaseURLPath(raw string) (string, error) {
|
||||
return strings.TrimRight(parsed.String(), "/"), nil
|
||||
}
|
||||
|
||||
// IsOfficialBaseURLHost 报告 host 是否属于官方 API / CLI 网关主机。
|
||||
// IsOfficialBaseURLHost 报告 host 是否属于官方 API / 区域 API / CLI 网关主机。
|
||||
func IsOfficialBaseURLHost(host string) bool {
|
||||
host = strings.ToLower(strings.TrimSpace(host))
|
||||
for _, allowed := range baseURLAllowedHosts {
|
||||
if strings.HasPrefix(allowed, "*.") {
|
||||
suffix := strings.TrimPrefix(allowed, "*.")
|
||||
if host == suffix || strings.HasSuffix(host, "."+suffix) {
|
||||
return true
|
||||
}
|
||||
continue
|
||||
}
|
||||
if host == allowed {
|
||||
return true
|
||||
}
|
||||
@@ -343,7 +352,18 @@ func IsOfficialBaseURLHost(host string) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// IsOfficialBaseURL 报告 raw 是否指向官方主机(api.x.ai 或 CLI 网关),
|
||||
// IsParseableBaseURL 报告 raw 是否能解析出 host。
|
||||
// 供读取路径判定存量脏数据:无法解析的值应回落默认端点,而不是把流量发往未定义目标。
|
||||
func IsParseableBaseURL(raw string) bool {
|
||||
trimmed := strings.TrimSpace(raw)
|
||||
if trimmed == "" {
|
||||
return false
|
||||
}
|
||||
parsed, err := url.Parse(trimmed)
|
||||
return err == nil && parsed.Host != ""
|
||||
}
|
||||
|
||||
// IsOfficialBaseURL 报告 raw 是否指向官方主机(api.x.ai / *.api.x.ai 区域端点 / CLI 网关),
|
||||
// 容忍存量凭证中的历史变体(大小写、显式 443 端口、百分号编码 path 等)。
|
||||
// 无法解析的值一并视为官方,调用方据此回落默认端点而不是把流量发往未定义目标。
|
||||
func IsOfficialBaseURL(raw string) bool {
|
||||
|
||||
@@ -211,12 +211,32 @@ func TestIsOfficialBaseURL(t *testing.T) {
|
||||
"https://relay.example.test/xai/v1",
|
||||
"http://relay.example.test/v1",
|
||||
"https://grok.com.evil.example.test/v1",
|
||||
"https://api.x.ai.evil.example.test/v1", // 后缀伪装不属于 *.api.x.ai
|
||||
}
|
||||
for _, raw := range custom {
|
||||
require.False(t, IsOfficialBaseURL(raw), "expected custom: %q", raw)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegionalAPIEndpointsAreOfficialAndTrusted(t *testing.T) {
|
||||
regional := []string{
|
||||
"https://us-east-1.api.x.ai/v1",
|
||||
"https://us-west-2.api.x.ai/v1",
|
||||
"https://eu-west-1.api.x.ai/v1",
|
||||
}
|
||||
for _, raw := range regional {
|
||||
require.True(t, IsOfficialBaseURL(raw), "expected official: %q", raw)
|
||||
|
||||
validated, err := ValidateTrustedBaseURL(raw)
|
||||
require.NoError(t, err, "trusted validation should accept regional endpoint %q", raw)
|
||||
require.Equal(t, raw, validated)
|
||||
}
|
||||
|
||||
// 区域端点作为官方主机同样强制 /v1 path
|
||||
_, err := ValidateTrustedBaseURL("https://us-east-1.api.x.ai/other")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestValidateBaseURLsRejectEmptyQueryDelimiter(t *testing.T) {
|
||||
_, err := ValidateBaseURL("https://grok.example.test/v1?")
|
||||
require.Error(t, err)
|
||||
|
||||
@@ -893,6 +893,20 @@ func (r *accountRepository) ListOpsAccountsForStats(ctx context.Context, platfor
|
||||
func accountListOrder(params pagination.PaginationParams) []func(*entsql.Selector) {
|
||||
sortBy := strings.ToLower(strings.TrimSpace(params.SortBy))
|
||||
sortOrder := params.NormalizedSortOrder(pagination.SortOrderAsc)
|
||||
if sortBy == "upstream_billing_rate" {
|
||||
direction := "ASC"
|
||||
tieOrder := entsql.Asc
|
||||
if sortOrder == pagination.SortOrderDesc {
|
||||
direction = "DESC"
|
||||
tieOrder = entsql.Desc
|
||||
}
|
||||
return []func(*entsql.Selector){func(s *entsql.Selector) {
|
||||
extra := s.C(dbaccount.FieldExtra)
|
||||
expression := upstreamBillingRateSortExpression(extra)
|
||||
s.OrderExpr(entsql.Expr(expression + " " + direction + " NULLS LAST"))
|
||||
s.OrderBy(tieOrder(s.C(dbaccount.FieldID)))
|
||||
}}
|
||||
}
|
||||
|
||||
field := dbaccount.FieldName
|
||||
defaultOrder := true
|
||||
@@ -934,6 +948,40 @@ func accountListOrder(params pagination.PaginationParams) []func(*entsql.Selecto
|
||||
return []func(*entsql.Selector){dbent.Asc(field), dbent.Asc(dbaccount.FieldID)}
|
||||
}
|
||||
|
||||
func upstreamBillingRateSortExpression(extra string) string {
|
||||
status := extra + " #>> '{upstream_billing_probe,status}'"
|
||||
effectiveJSON := extra + " #> '{upstream_billing_probe,data,effective_rate_multiplier}'"
|
||||
effective := extra + " #>> '{upstream_billing_probe,data,effective_rate_multiplier}'"
|
||||
resolvedJSON := extra + " #> '{upstream_billing_probe,data,resolved_rate_multiplier}'"
|
||||
resolved := extra + " #>> '{upstream_billing_probe,data,resolved_rate_multiplier}'"
|
||||
peakEnabledJSON := extra + " #> '{upstream_billing_probe,data,peak_rate_enabled}'"
|
||||
peakEnabled := extra + " #>> '{upstream_billing_probe,data,peak_rate_enabled}'"
|
||||
peakStart := extra + " #>> '{upstream_billing_probe,data,peak_start}'"
|
||||
peakEnd := extra + " #>> '{upstream_billing_probe,data,peak_end}'"
|
||||
peakMultiplierJSON := extra + " #> '{upstream_billing_probe,data,peak_rate_multiplier}'"
|
||||
peakMultiplier := extra + " #>> '{upstream_billing_probe,data,peak_rate_multiplier}'"
|
||||
peakMultiplierValue := "(CASE WHEN jsonb_typeof(" + peakMultiplierJSON + ") = 'number' THEN (" + peakMultiplier + ")::numeric END)"
|
||||
billingScope := extra + " #>> '{upstream_billing_probe,data,billing_scope}'"
|
||||
timezone := extra + " #>> '{upstream_billing_probe,data,timezone}'"
|
||||
validClock := "'^([01][0-9]|2[0-3]):[0-5][0-9]$'"
|
||||
startMinute := "(CASE WHEN " + peakStart + " ~ " + validClock + " THEN split_part(" + peakStart + ", ':', 1)::numeric * 60 + split_part(" + peakStart + ", ':', 2)::numeric END)"
|
||||
endMinute := "(CASE WHEN " + peakEnd + " ~ " + validClock + " THEN split_part(" + peakEnd + ", ':', 1)::numeric * 60 + split_part(" + peakEnd + ", ':', 2)::numeric END)"
|
||||
localMinute := "(EXTRACT(HOUR FROM (CURRENT_TIMESTAMP AT TIME ZONE (" + timezone + "))) * 60 + EXTRACT(MINUTE FROM (CURRENT_TIMESTAMP AT TIME ZONE (" + timezone + "))))"
|
||||
validPeakWindow := peakStart + " ~ " + validClock + " AND " +
|
||||
peakEnd + " ~ " + validClock + " AND " +
|
||||
startMinute + " < " + endMinute
|
||||
validPeakConfig := validPeakWindow + " AND " + peakMultiplierValue + " >= 0 AND " +
|
||||
"EXISTS (SELECT 1 FROM pg_timezone_names WHERE name = " + timezone + ")"
|
||||
dynamicRate := "CASE WHEN " + peakEnabled + " = 'false' THEN (" + resolved + ")::numeric WHEN " + peakEnabled + " = 'true' AND " + validPeakConfig +
|
||||
" THEN (" + resolved + ")::numeric * CASE WHEN " + localMinute + " >= " + startMinute + " AND " + localMinute + " < " + endMinute +
|
||||
" THEN " + peakMultiplierValue + " ELSE 1 END ELSE NULL END"
|
||||
legacySnapshot := "jsonb_typeof(" + resolvedJSON + ") IS NULL AND jsonb_typeof(" + peakEnabledJSON + ") IS NULL"
|
||||
|
||||
return "CASE WHEN " + status + " IN ('ok', 'failed') AND (jsonb_typeof(" + resolvedJSON + ") = 'number' OR jsonb_typeof(" + effectiveJSON + ") = 'number') THEN CASE WHEN jsonb_typeof(" +
|
||||
resolvedJSON + ") = 'number' AND jsonb_typeof(" + peakEnabledJSON + ") = 'boolean' THEN CASE WHEN " + billingScope + " = 'token' THEN " + dynamicRate + " ELSE NULL END WHEN " + legacySnapshot +
|
||||
" AND jsonb_typeof(" + effectiveJSON + ") = 'number' THEN (" + effective + ")::numeric END END"
|
||||
}
|
||||
|
||||
func (r *accountRepository) ListByGroup(ctx context.Context, groupID int64) ([]service.Account, error) {
|
||||
accounts, err := r.queryAccountsByGroup(ctx, groupID, accountGroupQueryOptions{
|
||||
status: service.StatusActive,
|
||||
@@ -1877,6 +1925,46 @@ func (r *accountRepository) ListSchedulableByGroupIDAndPlatforms(ctx context.Con
|
||||
})
|
||||
}
|
||||
|
||||
// ListModelAvailabilityCandidates returns the persistently configured account
|
||||
// pool used to decide whether a model is supported. Unlike scheduling queries,
|
||||
// it intentionally ignores transient runtime state (rate limits, overload,
|
||||
// temporary unschedulability, and expiry windows).
|
||||
func (r *accountRepository) ListModelAvailabilityCandidates(
|
||||
ctx context.Context,
|
||||
groupID *int64,
|
||||
platforms []string,
|
||||
includeGrouped bool,
|
||||
) ([]service.Account, error) {
|
||||
if len(platforms) == 0 {
|
||||
return []service.Account{}, nil
|
||||
}
|
||||
if groupID != nil {
|
||||
return r.queryAccountsByGroup(ctx, *groupID, accountGroupQueryOptions{
|
||||
status: service.StatusActive,
|
||||
schedulable: true,
|
||||
ignoreTransientState: true,
|
||||
platforms: platforms,
|
||||
})
|
||||
}
|
||||
|
||||
preds := []dbpredicate.Account{
|
||||
dbaccount.StatusEQ(service.StatusActive),
|
||||
dbaccount.SchedulableEQ(true),
|
||||
dbaccount.PlatformIn(platforms...),
|
||||
}
|
||||
if !includeGrouped {
|
||||
preds = append(preds, dbaccount.Not(dbaccount.HasAccountGroups()))
|
||||
}
|
||||
accounts, err := r.client.Account.Query().
|
||||
Where(preds...).
|
||||
Order(dbent.Asc(dbaccount.FieldPriority)).
|
||||
All(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return r.accountsToService(ctx, accounts)
|
||||
}
|
||||
|
||||
func (r *accountRepository) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error {
|
||||
now := time.Now()
|
||||
_, err := r.client.Account.Update().
|
||||
@@ -2575,6 +2663,12 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
args = append(args, *updates.Schedulable)
|
||||
idx++
|
||||
}
|
||||
if updates.ProbeEnabled != nil {
|
||||
if updates.Extra == nil {
|
||||
updates.Extra = make(map[string]any)
|
||||
}
|
||||
updates.Extra[service.UpstreamBillingProbeEnabledExtraKey] = *updates.ProbeEnabled
|
||||
}
|
||||
// JSONB 需要合并而非覆盖,使用 raw SQL 保持旧行为。
|
||||
if len(updates.Credentials) > 0 {
|
||||
payload, err := json.Marshal(updates.Credentials)
|
||||
@@ -2605,8 +2699,14 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
|
||||
setClauses = append(setClauses, "updated_at = NOW()")
|
||||
|
||||
query := "UPDATE accounts SET " + joinClauses(setClauses, ", ") + " WHERE id = ANY($" + itoa(idx) + ") AND deleted_at IS NULL"
|
||||
whereClause := " WHERE id = ANY($" + itoa(idx) + ") AND deleted_at IS NULL"
|
||||
args = append(args, pq.Array(ids))
|
||||
idx++
|
||||
if updates.ProbeEnabled != nil {
|
||||
whereClause += " AND platform = $" + itoa(idx) + " AND type = $" + itoa(idx+1)
|
||||
args = append(args, service.PlatformOpenAI, service.AccountTypeAPIKey)
|
||||
}
|
||||
query := "UPDATE accounts SET " + joinClauses(setClauses, ", ") + whereClause
|
||||
|
||||
baseCtx := ctx
|
||||
contextTx := dbent.TxFromContext(ctx)
|
||||
@@ -2635,6 +2735,20 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if updates.ProbeEnabled != nil {
|
||||
expectedRows := int64(0)
|
||||
seenIDs := make(map[int64]struct{}, len(ids))
|
||||
for _, id := range ids {
|
||||
if _, seen := seenIDs[id]; seen {
|
||||
continue
|
||||
}
|
||||
seenIDs[id] = struct{}{}
|
||||
expectedRows++
|
||||
}
|
||||
if rows != expectedRows {
|
||||
return 0, service.ErrUpstreamBillingProbeAccountInvalid
|
||||
}
|
||||
}
|
||||
if rows > 0 {
|
||||
payload := map[string]any{"account_ids": ids}
|
||||
if err := enqueueSchedulerOutbox(ctx, exec, service.SchedulerOutboxEventAccountBulkChanged, nil, nil, payload); err != nil {
|
||||
@@ -2662,9 +2776,10 @@ func (r *accountRepository) BulkUpdate(ctx context.Context, ids []int64, updates
|
||||
}
|
||||
|
||||
type accountGroupQueryOptions struct {
|
||||
status string
|
||||
schedulable bool
|
||||
platforms []string // 允许的多个平台,空切片表示不进行平台过滤
|
||||
status string
|
||||
schedulable bool
|
||||
ignoreTransientState bool
|
||||
platforms []string // 允许的多个平台,空切片表示不进行平台过滤
|
||||
}
|
||||
|
||||
func (r *accountRepository) queryAccountsByGroup(ctx context.Context, groupID int64, opts accountGroupQueryOptions) ([]service.Account, error) {
|
||||
@@ -2681,14 +2796,16 @@ func (r *accountRepository) queryAccountsByGroup(ctx context.Context, groupID in
|
||||
preds = append(preds, dbaccount.PlatformIn(opts.platforms...))
|
||||
}
|
||||
if opts.schedulable {
|
||||
now := time.Now()
|
||||
preds = append(preds,
|
||||
dbaccount.SchedulableEQ(true),
|
||||
tempUnschedulablePredicate(),
|
||||
notExpiredPredicate(now),
|
||||
dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
|
||||
dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
|
||||
)
|
||||
preds = append(preds, dbaccount.SchedulableEQ(true))
|
||||
if !opts.ignoreTransientState {
|
||||
now := time.Now()
|
||||
preds = append(preds,
|
||||
tempUnschedulablePredicate(),
|
||||
notExpiredPredicate(now),
|
||||
dbaccount.Or(dbaccount.OverloadUntilIsNil(), dbaccount.OverloadUntilLTE(now)),
|
||||
dbaccount.Or(dbaccount.RateLimitResetAtIsNil(), dbaccount.RateLimitResetAtLTE(now)),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
if len(preds) > 0 {
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
_ "github.com/Wei-Shaw/sub2api/ent/runtime"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"entgo.io/ent/dialect"
|
||||
entsql "entgo.io/ent/dialect/sql"
|
||||
)
|
||||
|
||||
func TestListModelAvailabilityCandidates_GroupQueryIgnoresTransientState(t *testing.T) {
|
||||
var capturedSQL string
|
||||
db, mock, err := sqlmock.New(sqlmock.QueryMatcherOption(captureEntQueryMatcher{actual: &capturedSQL}))
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
|
||||
driver := entsql.OpenDB(dialect.Postgres, db)
|
||||
client := dbent.NewClient(dbent.Driver(driver))
|
||||
t.Cleanup(func() { _ = client.Close() })
|
||||
repo := newAccountRepositoryWithSQL(client, db, nil)
|
||||
|
||||
mock.ExpectQuery("model availability candidates").
|
||||
WillReturnRows(sqlmock.NewRows([]string{"id"}))
|
||||
groupID := int64(42)
|
||||
accounts, err := repo.ListModelAvailabilityCandidates(
|
||||
context.Background(),
|
||||
&groupID,
|
||||
[]string{service.PlatformAnthropic},
|
||||
false,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, accounts)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
|
||||
normalized := normalizeSQLWhitespace(capturedSQL)
|
||||
_, whereClause, found := strings.Cut(normalized, " WHERE ")
|
||||
require.True(t, found, "expected WHERE clause in query: %s", normalized)
|
||||
whereClause, _, _ = strings.Cut(whereClause, " ORDER BY ")
|
||||
for _, configuredPredicate := range []string{"group_id", "status", "schedulable", "platform"} {
|
||||
require.Contains(t, whereClause, configuredPredicate)
|
||||
}
|
||||
for _, transientPredicate := range []string{
|
||||
"rate_limit_reset_at",
|
||||
"overload_until",
|
||||
"temp_unschedulable_until",
|
||||
"expires_at",
|
||||
"auto_pause_on_expired",
|
||||
} {
|
||||
require.NotContains(t, whereClause, transientPredicate, "configured-state diagnosis must not filter transient predicate %q", transientPredicate)
|
||||
}
|
||||
}
|
||||
@@ -3,6 +3,9 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
@@ -33,3 +36,105 @@ func (s *AccountRepoSuite) TestListWithFilters_SortByPriorityDesc() {
|
||||
s.Require().Equal("high-priority", accounts[0].Name)
|
||||
s.Require().Equal("low-priority", accounts[1].Name)
|
||||
}
|
||||
|
||||
func (s *AccountRepoSuite) TestListWithFilters_SortByUpstreamBillingRateWithMissingLast() {
|
||||
makeAccount := func(name, status string, rate any) {
|
||||
extra := map[string]any{}
|
||||
if rate != nil {
|
||||
extra[service.UpstreamBillingProbeExtraKey] = map[string]any{
|
||||
"status": status,
|
||||
"data": map[string]any{"effective_rate_multiplier": rate},
|
||||
}
|
||||
}
|
||||
mustCreateAccount(s.T(), s.client, &service.Account{Name: name, Extra: extra})
|
||||
}
|
||||
makeAccount("high-rate", service.UpstreamBillingProbeStatusOK, 0.8)
|
||||
makeAccount("low-rate", service.UpstreamBillingProbeStatusOK, 0.03)
|
||||
makeAccount("missing-rate", "", nil)
|
||||
makeAccount("unsupported-with-retained-rate", service.UpstreamBillingProbeStatusUnsupported, 0.01)
|
||||
|
||||
for _, tc := range []struct {
|
||||
order string
|
||||
want []string
|
||||
}{
|
||||
{order: "asc", want: []string{"low-rate", "high-rate", "missing-rate", "unsupported-with-retained-rate"}},
|
||||
{order: "desc", want: []string{"high-rate", "low-rate", "unsupported-with-retained-rate", "missing-rate"}},
|
||||
} {
|
||||
accounts, _, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{
|
||||
Page: 1, PageSize: 10, SortBy: "upstream_billing_rate", SortOrder: tc.order,
|
||||
}, "", "", "", "", 0, "")
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(accounts, 4)
|
||||
for i, name := range tc.want {
|
||||
s.Require().Equal(name, accounts[i].Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AccountRepoSuite) TestListWithFilters_SortByCurrentUpstreamBillingRateDuringPeak() {
|
||||
now := time.Now()
|
||||
locations := []string{"UTC", "Asia/Shanghai", "America/New_York", "Europe/London"}
|
||||
var timezone string
|
||||
var minute int
|
||||
for _, name := range locations {
|
||||
location, err := time.LoadLocation(name)
|
||||
s.Require().NoError(err)
|
||||
local := now.In(location)
|
||||
candidate := local.Hour()*60 + local.Minute()
|
||||
if candidate >= 2 && candidate <= 1436 {
|
||||
timezone = name
|
||||
minute = candidate
|
||||
break
|
||||
}
|
||||
}
|
||||
s.Require().NotEmpty(timezone)
|
||||
|
||||
peakStart := fmt.Sprintf("%02d:%02d", (minute-2)/60, (minute-2)%60)
|
||||
peakEnd := fmt.Sprintf("%02d:%02d", (minute+3)/60, (minute+3)%60)
|
||||
mustCreateAccount(s.T(), s.client, &service.Account{
|
||||
Name: "current-peak-rate",
|
||||
Extra: map[string]any{
|
||||
service.UpstreamBillingProbeExtraKey: map[string]any{
|
||||
"status": service.UpstreamBillingProbeStatusOK,
|
||||
"data": map[string]any{
|
||||
"billing_scope": "token",
|
||||
"resolved_rate_multiplier": 1.0,
|
||||
"effective_rate_multiplier": 1.0,
|
||||
"peak_rate_enabled": true,
|
||||
"peak_start": peakStart,
|
||||
"peak_end": peakEnd,
|
||||
"peak_rate_multiplier": 10.0,
|
||||
"timezone": timezone,
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
mustCreateAccount(s.T(), s.client, &service.Account{
|
||||
Name: "current-off-peak-rate",
|
||||
Extra: map[string]any{
|
||||
service.UpstreamBillingProbeExtraKey: map[string]any{
|
||||
"status": service.UpstreamBillingProbeStatusOK,
|
||||
"data": map[string]any{
|
||||
"effective_rate_multiplier": 5.0,
|
||||
},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
for _, tc := range []struct {
|
||||
order string
|
||||
want []string
|
||||
}{
|
||||
{order: "asc", want: []string{"current-off-peak-rate", "current-peak-rate"}},
|
||||
{order: "desc", want: []string{"current-peak-rate", "current-off-peak-rate"}},
|
||||
} {
|
||||
accounts, _, err := s.repo.ListWithFilters(s.ctx, pagination.PaginationParams{
|
||||
Page: 1, PageSize: 10, SortBy: "upstream_billing_rate", SortOrder: tc.order,
|
||||
}, "", "", "", "", 0, "")
|
||||
s.Require().NoError(err)
|
||||
s.Require().Len(accounts, 2)
|
||||
for i, name := range tc.want {
|
||||
s.Require().Equal(name, accounts[i].Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -163,6 +163,46 @@ func TestBulkUpdateNilProbeRemovesKeyInsteadOfWritingJSONNull(t *testing.T) {
|
||||
require.Contains(t, normalizeSQLWhitespace(exec.execQueries[0]), "- 'upstream_billing_probe'")
|
||||
}
|
||||
|
||||
func TestBulkUpdateDisablingProbeRemovesSnapshot(t *testing.T) {
|
||||
exec := &recordingSQLExecutor{result: rowsAffectedResult(1)}
|
||||
repo := newAccountRepositoryWithSQL(nil, exec, nil)
|
||||
|
||||
_, err := repo.BulkUpdate(context.Background(), []int64{27}, service.AccountBulkUpdate{
|
||||
Extra: map[string]any{service.UpstreamBillingProbeEnabledExtraKey: false},
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotEmpty(t, exec.execQueries)
|
||||
require.Contains(t, normalizeSQLWhitespace(exec.execQueries[0]), "- 'upstream_billing_probe'")
|
||||
payload, ok := exec.execArgs[0][0].([]byte)
|
||||
require.True(t, ok)
|
||||
require.Equal(t, `{"upstream_billing_probe_enabled":false}`, string(payload))
|
||||
}
|
||||
|
||||
func TestBulkUpdateProbeEligibilityMismatchRollsBack(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
client := dbent.NewClient(dbent.Driver(entsql.OpenDB(dialect.Postgres, db)))
|
||||
t.Cleanup(func() { _ = client.Close() })
|
||||
|
||||
enabled := true
|
||||
mock.ExpectBegin()
|
||||
mock.ExpectExec(`(?s)UPDATE accounts SET extra = .* WHERE id = ANY\(\$2\) AND deleted_at IS NULL AND platform = \$3 AND type = \$4`).
|
||||
WithArgs(sqlmock.AnyArg(), `{27,28}`, service.PlatformOpenAI, service.AccountTypeAPIKey).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
mock.ExpectRollback()
|
||||
|
||||
repo := newAccountRepositoryWithSQL(client, db, nil)
|
||||
rows, err := repo.BulkUpdate(context.Background(), []int64{27, 28}, service.AccountBulkUpdate{
|
||||
ProbeEnabled: &enabled,
|
||||
})
|
||||
|
||||
require.ErrorIs(t, err, service.ErrUpstreamBillingProbeAccountInvalid)
|
||||
require.Zero(t, rows)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestUpdateCredentialsAtomicallyClearsProbeForOpenAIAPIKeyIdentityChange(t *testing.T) {
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
|
||||
@@ -110,28 +110,25 @@ func (c *apiKeyCache) SubscribeAuthCacheInvalidation(ctx context.Context, handle
|
||||
return fmt.Errorf("subscribe to auth cache invalidation: %w", err)
|
||||
}
|
||||
|
||||
go func() {
|
||||
defer func() {
|
||||
if err := pubsub.Close(); err != nil {
|
||||
log.Printf("Warning: failed to close auth cache invalidation pubsub: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
ch := pubsub.Channel()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case msg, ok := <-ch:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if msg != nil {
|
||||
handler(msg.Payload)
|
||||
}
|
||||
}
|
||||
defer func() {
|
||||
if err := pubsub.Close(); err != nil {
|
||||
log.Printf("Warning: failed to close auth cache invalidation pubsub: %v", err)
|
||||
}
|
||||
}()
|
||||
service.NotifyAuthCacheSubscriptionReady(ctx)
|
||||
|
||||
return nil
|
||||
ch := pubsub.Channel()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case msg, ok := <-ch:
|
||||
if !ok {
|
||||
return errors.New("auth cache invalidation pubsub channel closed")
|
||||
}
|
||||
if msg != nil {
|
||||
handler(msg.Payload)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user