Compare commits

...
118 Commits
Author SHA1 Message Date
Wesley Liddick 19149ca196 Merge pull request #4558 from fengshao1227/fix/antigravity-plan-type-preserve
fix(antigravity): 保留付费 tier 的 PlanType,IneligibleTiers 仅标记异常状态
2026-07-18 21:48:44 +08:00
Wesley Liddick 831812b39d Merge pull request #4556 from superman2003/fix/grok-free-probe-encrypted-recovery
fix(grok): stabilize Free probes and encrypted reasoning recovery
2026-07-18 21:48:31 +08:00
Wesley Liddick d98ac48eb9 Merge pull request #4543 from heathermhuang/codex/secure-protected-video-content-4498
fix(grok): securely proxy protected video content
2026-07-18 21:46:31 +08:00
li 72b29a7d4e fix(antigravity): 保留付费 tier 的 PlanType,IneligibleTiers 仅标记异常状态
Fixes #4519
2026-07-18 21:43:54 +08:00
shaw 14608dc6d4 Merge origin/main into codex/secure-protected-video-content-4498
Reconcile with #4539 (grok media account model mapping), now on main:
- handler/grok_media.go non-failover error path keeps both changes —
  #4539's grokMediaScheduleModel(account, routingModel, nil) schedule
  attribution and this branch's IsResponseCommitted guard
- auto-merged sections verified: routing/classify use #4539's routingModel,
  video lookup owner-binding and no-failover semantics intact, ForwardGrokMedia
  keeps mapping block (skipped for lookup endpoints via RequiresRequestBody),
  empty-image failover, and video-status URL rewrite in order
2026-07-18 21:38:17 +08:00
Wesley Liddick daa4989124 Merge pull request #4539 from heathermhuang/codex/fix-grok-media-model-mapping-4503
fix(grok): apply account model mapping to media
2026-07-18 21:34:10 +08:00
shaw 8a95a46a69 Merge origin/main into codex/secure-protected-video-content-4498
Resolve conflicts with main:
- service/grok_media.go: keep both post-response blocks — #4497's empty
  image-output failover (main) runs first for image endpoints, then this
  branch's video-status content-URL rewrite; the endpoint conditions are
  mutually exclusive
- handler/openai_gateway_credential_failover_loop_test.go: mark the stub
  OAuth accounts media-eligible via the grok_media_eligible extra override,
  because this branch moved grok media failover coverage to the generation
  endpoint, which is now gated by #4540's paid-eligibility probe on main
2026-07-18 21:31:28 +08:00
shaw 9d498c2474 Merge origin/main into codex/fix-grok-media-model-mapping-4503
Resolve conflicts with main:
- handler/grok_media_test.go: keep both new tests (schedule model test from
  this branch, eligibility gating tests from #4540)
- service/openai_gateway_grok_test.go: keep both new tests; update the image
  cases of the mapping table test to return a non-empty image payload because
  #4497 (already on main) now converts empty image responses into an upstream
  failover error
2026-07-18 21:25:43 +08:00
Wesley Liddick a2f802d409 Merge pull request #4541 from wucm667/fix/issue-4532-renew-expired-subscription
fix(subscription): renew expired admin assignments
2026-07-18 21:14:23 +08:00
Wesley Liddick b01196a7a8 Merge pull request #4553 from wp-a/fix/openai-ws-turn-lifecycle
[codex] enforce websocket passthrough turn lifecycle
2026-07-18 21:10:48 +08:00
superman2003 e14fb2b6ff fix(grok): recover invalid encrypted content once 2026-07-18 21:05:50 +08:00
superman2003 dd7a2b22f0 fix(grok): align Free probes with health checks 2026-07-18 21:05:50 +08:00
shaw 4e8ea7d568 fix(test): 池模式临时规则测试适配 #4547 模型级隔离语义
#4496 与 #4547 文本无冲突但语义相撞:#4496 的测试断言规则命中后账号级
跨模型封锁,#4547(issue 4527 第4点)将已知模型的临时不可调度改为按
模型隔离(SetModelRateLimit,不再触发账号级 runtime block)。

按 #4547 的语义更新断言:命中模型 gpt-5.4 记模型级封锁、gpt-5.5 不受
影响、不再调用账号级 SetTempUnschedulable;池模式规则仍生效(停止同
账号重试),issue 4470 的诉求不受影响。未知模型的账号级兜底已由
TestOpenAITempUnschedulable_UnknownModelKeepsAccountRuntimeBlock 覆盖。
2026-07-18 21:02:23 +08:00
Wesley Liddick a3f2b8fd87 Merge pull request #4540 from heathermhuang/codex/investigate-grok-oauth-test-4525
fix(grok): retry CLI chat permission denial
2026-07-18 20:50:07 +08:00
Wesley Liddick 23cdd20606 Merge pull request #4537 from heathermhuang/codex/refresh-anthropic-monitor-1953
fix(monitor): refresh Anthropic text-block extraction
2026-07-18 20:49:40 +08:00
Wesley Liddick a62b821b5f Merge pull request #4478 from yardbirds0/codex/fix-upstream-billing-probe-refresh
fix: 完善上游 Sub2API 倍率探测刷新、展示与账号配置
2026-07-18 20:49:14 +08:00
Wesley Liddick 774ff5d8c8 Merge pull request #4515 from BenjaminAaron196/feat/filter-noise-rejected-requests
(fix) 过滤入口拒绝日志并强化鉴权安全边界
2026-07-18 20:46:50 +08:00
Wesley Liddick d2667393b3 Merge pull request #4522 from wucm667/docs/issue-4518-http-bridge-prerequisite
docs: clarify OpenAI WS mode router prerequisite
2026-07-18 20:43:01 +08:00
Wesley Liddick e002fbb349 Merge pull request #4508 from wucm667/fix/model-not-found-transient-misclassification
fix: 临时账号耗尽时保留 503 错误分类
2026-07-18 20:41:35 +08:00
Wesley Liddick 005bc5c8d4 Merge pull request #4497 from heathermhuang/agent/fix-grok-media-fallback-4471
fix(grok): fail closed for ineligible OAuth media
2026-07-18 20:41:24 +08:00
Wesley Liddick bc4bd118d8 Merge pull request #4547 from heathermhuang/codex/fix-model-scoped-temp-cooldown-4527
fix(routing): isolate temporary cooldowns by model
2026-07-18 20:41:12 +08:00
Wesley Liddick f3925db11b Merge pull request #4496 from StarryKira/agent/fix-4470-pool-temp-unschedulable
fix: honor temp unschedulable rules in pool mode
2026-07-18 20:40:12 +08:00
Wesley Liddick f9a467a4c0 Merge pull request #4520 from StarryKira/codex/fix-4487
fix: report Chat Completions stream transport failures
2026-07-18 20:39:49 +08:00
Wesley Liddick 8ce9288a6a Merge pull request #4489 from fengshao1227/fix/grok-free-cache-tool-injection
fix(grok): 纯客户端函数工具不再注入原生搜索工具
2026-07-18 20:39:37 +08:00
Wesley Liddick c1e702be5e Merge pull request #4505 from cyhhao/fix/claude-1m-model-suffix
fix(gateway): normalize Claude Code 1m model suffix
2026-07-18 20:39:27 +08:00
Wesley Liddick bc6b69289c Merge pull request #4468 from docooler/fix/responses-stream-content-part
fix(apicompat): emit content_part events and full output in Responses stream
2026-07-18 20:39:15 +08:00
Wesley Liddick 5a0492bf88 Merge pull request #4517 from weiness/fix/custom-branding-flash
fix: prevent custom branding flash on initial load
2026-07-18 20:39:00 +08:00
Wesley Liddick fdc9d92fdb Merge pull request #4528 from feitianbubu/fix/plan-validity-label-unit
fix(i18n): 套餐有效期表头与表单标签去掉写死的"天",与动态单位一致
2026-07-18 20:38:47 +08:00
Wesley Liddick 8cdd372b7e Merge pull request #4507 from coo1white/fix-dockerfile-cross-compile
fix(docker): cross-compile the image instead of running Go under QEMU
2026-07-18 20:38:37 +08:00
Wesley Liddick 080a52121a Merge pull request #4506 from coo1white/fix-compose-redis-command
fix(deploy): make the redis command flags take effect
2026-07-18 20:38:27 +08:00
王鹏 f0e0b7e6d8 fix(openai-ws): enforce passthrough turn lifecycle 2026-07-18 19:17:53 +08:00
Heatherm Huang 4c845ca697 test(grok): sanitize security fixtures 2026-07-18 17:50:53 +08:00
Heatherm Huang 95d27f2dce fix(grok): fetch validated signed video content 2026-07-18 17:49:09 +08:00
Heatherm Huang f2e7a55703 fix(grok): proxy signed video status URLs 2026-07-18 17:21:53 +08:00
Heatherm Huang 1ed9f59599 test(grok): keep media failover coverage on generation 2026-07-18 15:54:43 +08:00
Heatherm Huang b4ad7fd363 test(grok): configure multipart media mapping 2026-07-18 15:49:21 +08:00
Heatherm Huang ef2a22be70 fix(routing): isolate temporary cooldowns by model 2026-07-18 15:35:30 +08:00
Heatherm Huang c831bb979f fix(grok): scope video content to request owner 2026-07-18 14:56:15 +08:00
Heatherm Huang 3edcbfd120 fix(grok): keep media content proxy URLs same-origin 2026-07-18 14:43:08 +08:00
wucm667 1db10dc559 fix(subscription): renew expired admin assignments 2026-07-18 14:42:22 +08:00
Vz7797 3f6b5c7bd7 fix(grok): proxy protected video content through upstream account
(cherry picked from commit a01177e294)
2026-07-18 14:37:25 +08:00
Heatherm Huang 2946281a43 fix(grok): retry CLI chat permission denial 2026-07-18 14:32:56 +08:00
Heatherm Huang 335edde9c8 fix(grok): apply account model mapping to media 2026-07-18 14:21:47 +08:00
mark-ly-wang 1a855d3e0e fix(monitor): extract anthropic text blocks
(cherry picked from commit 76cf1ee834)
2026-07-18 13:53:36 +08:00
Wesley Liddick b1a6b80267 Merge pull request #4526 from Wei-Shaw/feat/security-switches-default-off
feat(security): 敏感操作 step-up 2FA 开关化,安全开关默认关闭
2026-07-18 11:32:17 +08:00
feitianbubu c3f759d13a fix(i18n): 套餐有效期表头与表单标签去掉写死的"天",与动态单位一致 2026-07-18 10:59:11 +08:00
shaw 539bfc8bad feat(security): 敏感操作 step-up 2FA 开关化,安全开关默认关闭
新增系统设置 step_up_enabled(默认关闭),把敏感操作 2FA 门控做成可开关;
同时将会话 IP/UA 绑定默认值从开启改为关闭,避免用户因 IP 变动登录后掉线。

## 新增功能
- 敏感操作 step-up 2FA 总开关 step_up_enabled(默认关闭):关闭时账号/代理导出、
  备份创建/下载、S3 配置修改、提升管理员等操作恢复门控引入前的直接放行行为;
  开启后要求当前会话在 15 分钟内完成过 TOTP step-up 验证。

## 优化改进
- 会话 IP/UA 绑定默认改为关闭(功能保留,可在设置页按需开启)。
- 开启 step-up 开关需操作者本人已启用 TOTP(防自锁);关闭开关本身作为敏感操作,
  需通过 step-up 验证(防止攻击者拿到会话后先关闸再导出/备份)。
- 两个安全开关请求字段改为可空指针(省略=保持现值),避免旧客户端全量保存时
  静默重置安全开关。
- 备份恢复(整库覆盖可回滚安全设置)纳入 step-up 门控。
- 审计摘要 diffSettings 补记 step_up_enabled / session_binding_enabled 变更。

## Bug 修复
- 修复 BackupView 恢复操作 409(恢复进行中)判断未适配 apiClient 扁平化错误对象。
2026-07-18 10:46:42 +08:00
wucm667 8b75dd5576 docs: clarify OpenAI WS mode router prerequisite 2026-07-18 08:37:15 +08:00
haruka bd0f2d6405 fix: report chat stream transport failures 2026-07-18 03:24:04 +08:00
benjamin 71afe24710 fix: 修复CI检查问题 2026-07-18 00:42:16 +08:00
weiness 500f241692 fix: prevent custom branding flash on initial load 2026-07-18 00:24:06 +08:00
benjamin b92bbf0299 fix: 过滤入口拒绝日志并强化鉴权边界 2026-07-18 00:11:18 +08:00
wucm667 2594950993 fix: classify transient account exhaustion as 503 instead of 404 2026-07-17 22:10:18 +08:00
Nick 9386ce8a13 fix(docker): cross-compile the image instead of running Go under QEMU
Building the image for linux/amd64 on an arm64 host (Apple Silicon)
kept failing at 'go mod download' with 'unexpected EOF'.

Two changes:

1. The builder stages had no --platform, so the Node and Go toolchains
   ran under QEMU emulation of the target arch. Pin frontend-builder
   and backend-builder to $BUILDPLATFORM and cross-compile: add
   TARGETOS/TARGETARCH args and set them on 'go build'. The binary is
   CGO_ENABLED=0, so this is a clean pure-Go cross-compile - much
   faster, and the emulated networking that dropped module fetches
   with EOF is gone. The frontend output is JS (arch-neutral), so it
   is safe to build on the host arch too.

2. Add go module and build cache mounts, so a retry after a network
   blip goes on from where it stopped instead of starting over.

The runtime image and app behavior are unchanged.
2026-07-17 21:05:42 +07:00
Nick be74deae73 fix(deploy): make the redis command flags take effect
The redis command is one quoted script given to the inner `sh -c`.
Docker compose keeps the newlines inside the quoted string, so
`redis-server` on the first line ran as a complete command with no
flags at all, and --save / --appendonly / --appendfsync after it were
silently never applied (`redis-cli CONFIG GET appendonly` said "no").

Trailing `\` line continuations fold the script back into one command.
Checked with redis:7-alpine: appendonly is now "yes" and save is
"60 1".
2026-07-17 21:04:47 +07:00
cyh 2264a33085 fix(gateway): normalize Claude Code 1m model suffix 2026-07-17 21:59:57 +08:00
Tian Lee d2585097f3 fix: 完善上游 Sub2API 计费倍率探测与账号展示 2026-07-17 20:59:29 +08:00
Heatherm Huang e86063155f fix(grok): gate OAuth media on paid eligibility 2026-07-17 19:15:16 +08:00
haruka 1f7b7b9132 fix: honor temp unschedulable rules in pool mode 2026-07-17 18:12:33 +08:00
li 7043982ae3 fix(grok): 纯客户端函数工具不再注入原生搜索工具
Fixes #4486
2026-07-17 16:59:40 +08:00
github-actions[bot] 57914967cb chore: sync VERSION to 0.1.160 [skip ci] 2026-07-17 08:48:10 +00:00
Wesley Liddick 8bfbc5ca99 Merge pull request #4485 from Sub2API-Devs/dev
feat(security-audit): 新增 OpenAI 兼容提示词审计能力与安全审计控制台
2026-07-17 16:15:26 +08:00
Wesley Liddick cb40288fb8 Merge pull request #4479 from fengshao1227/fix/backup-s3-stepup-totp
fix(frontend): saveS3Config 接入 step-up TOTP 门控
2026-07-17 16:14:49 +08:00
Wesley Liddick 33766c299f Merge pull request #4481 from fengshao1227/fix/passive-image-namespace-capability
fix(gateway): 被动 image_gen namespace 不再强制要求 Responses capability
2026-07-17 16:14:35 +08:00
Wesley Liddick 13468b3fa5 Merge pull request #4484 from heathermhuang/codex/fix-grok-scheduler-media-cache
fix(grok): preserve media eligibility in scheduler cache
2026-07-17 16:13:31 +08:00
mt21625457 e69e46ba02 feat: allow one-click audit event filter deletion 2026-07-17 15:50:31 +08:00
Heatherm Huang eaf0691752 fix(grok): preserve media eligibility in scheduler cache 2026-07-17 15:45:19 +08:00
li 5fd030a01b test: 更新 image permission 测试适配 explicit intent 检查
被动 namespace 不再触发 403,新增 PassiveNamespaceDoesNotTrigger403
测试和 OpenAI native/image model 的正向用例。
2026-07-17 15:29:14 +08:00
li fc3c5a1e0c fix: 权限检查和并发槽也使用显式检查,修复被动 namespace 触发 403
Fixes #4447
2026-07-17 15:15:54 +08:00
li e375623abf fix(gateway): 被动 image_gen namespace 不再强制要求 Responses capability
Fixes #4476
2026-07-17 15:07:41 +08:00
li 288d3a1cce fix(frontend): saveS3Config 接入 step-up TOTP 门控
Fixes #4445
2026-07-17 14:52:54 +08:00
mt21625457 ac685ccaf5 feat(security-audit): persist full prompts on audit events and polish event review UI
- Add prompt_audit_events.full_prompt (migration 182) so admins can review
  the exact unredacted prompt that triggered a finding; blocking mode writes
  it from the snapshot, async mode reconstructs it from the Redis scan
  payload so jobs rows stay redaction-only
- Event detail API returns full_prompt (list endpoint stays lean); text is
  NUL-stripped and capped at 65536 runes
- Detail dialog shows the full prompt in a scrollable pane with fallback to
  the legacy redacted preview; page copy updated to match the new behavior
- Rework filter deletion into a dedicated dialog with time-range presets and
  criteria-change preview invalidation; localize decision/risk/category
  labels across the events workspace
- Fix pre-existing i18n message-compile spec by declaring the
  @intlify/message-compiler dev dependency
2026-07-17 14:38:10 +08:00
shaw bc50c9d013 chore: update sponsors 2026-07-17 14:05:59 +08:00
Wesley Liddick dda20ae9a6 Merge pull request #4474 from heathermhuang/codex/fix-grok-media-image-url-4420
v0.1.159 fix(grok): normalize media payloads and quarantine ineligible accounts
2026-07-17 13:49:09 +08:00
Heatherm HuangandOpenAI Codex c34dcd3700 docs: document Grok media eligibility
Co-Authored-By: OpenAI Codex <noreply@openai.com>
2026-07-17 12:18:46 +08:00
mt21625457 7ed4e7e5e3 fix(security-audit): isolate latest user prompt chunk 2026-07-17 12:12:41 +08:00
Heatherm Huang d2ac6b2c97 fix(grok): harden media quarantine edge cases 2026-07-17 12:08:50 +08:00
Heatherm Huang 8f9296140f Merge remote-tracking branch 'origin/main' into codex/fix-grok-media-image-url-4420 2026-07-17 11:51:46 +08:00
Heatherm Huang 8bc0a277f2 fix(grok): quarantine ineligible media accounts 2026-07-17 11:51:16 +08:00
Heatherm Huang 456c619379 fix(grok): normalize media reference image payloads 2026-07-17 11:51:08 +08:00
mt21625457andCursor 18e698bed6 feat(security-audit): allow admin-managed audit node targets and polish pool UI
Let admins configure private/intranet Guard endpoints without destination-class blocking, and fix prompt-audit switch layout so thumbs and labels no longer overlap.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-17 11:45:14 +08:00
docooler 26b5e87457 fix(apicompat): emit content_part events and full output in Responses stream
The Anthropic→Responses streaming converter omits two things the OpenAI
Responses wire format requires, which breaks clients that reconstruct the
response from the event stream (rather than just reading deltas).

1. response.content_part.added is never emitted.

   A message item is opened 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 next
   output_text.delta indexes output.content[content_index] and raises:

     File "openai/lib/streaming/responses/_responses.py", line 352,
       in accumulate_event
         content = output.content[event.content_index]
     IndexError: list index out of range

   Raw iteration (responses.create(stream=True)) does not accumulate and is
   unaffected, which is why this went unnoticed.

2. response.completed carries Output: []ResponsesOutput{}.

   get_final_response() and tracing integrations parse the terminal event's
   response directly, so callers see an empty output_text even though the
   deltas streamed correctly. This one is invisible when only watching the
   stream render.

Also carries the full text on output_text.done / content_part.done (deltas
carry increments only, done events carry the whole part) and fills in
content/arguments/summary on output_item.done, which had the same
empty-payload issue.

Reproduced against a live Anthropic-platform group with the openai Python
SDK 2.46.0; Arize Phoenix's playground hits the same path. Verified before
(IndexError) and after (full text via get_final_response()).

Adds regression tests covering event ordering, done-event payloads, and the
terminal event's output for both text and tool calls.
2026-07-17 03:35:26 +00:00
github-actions[bot] c2c19a7cbe chore: sync VERSION to 0.1.159 [skip ci] 2026-07-17 02:00:12 +00:00
Wesley Liddick 2a75d7d238 Merge pull request #4462 from Wei-Shaw/fix/unify-security-client-ip
fix(security): unify audit log & session binding client IP with API key ACL trust toggle
2026-07-17 09:42:43 +08:00
Wesley Liddick 1c3f2810af Merge pull request #4461 from yardbirds0/feat/account-homepage-link
feat: 支持从 API Key 账号名称跳转上游站点主页
2026-07-17 09:34:50 +08:00
shawandClaude 7c48f9a85f fix(security): unify audit log & session binding client IP with API key ACL trust toggle
Behind a reverse proxy (e.g. nginx with X-Real-IP), admin audit logs and
session IP/UA binding always recorded 127.0.0.1 because they hardcoded
the gin trusted_proxies chain, while API key IP restriction already
honored the "trust forwarded client IP" system setting.

- add ip.GetSecurityClientIP(c, trustForwarded) as the single source of
  truth for security-sensitive client IP selection; API key auth
  middlewares (main + google) refactored onto it with zero behavior change
- SessionBindingContext(cfg) now resolves the client IP via the same
  toggle and injects it into the request context; token issuance,
  binding enforcement and its mismatch audit record all read the
  injected value, so issue/verify can never diverge
- audit log middleware and audit-log clear trace record the same
  security client IP (middleware.SecurityClientIP), falling back to the
  trusted proxy chain when the injection is absent
- settings UI hint (zh/en) documents the broadened toggle scope and the
  one-time re-login after toggling while session binding is enabled

With the toggle off (default) behavior is byte-for-byte unchanged.

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-17 09:33:18 +08:00
shawandClaude 980a46ffb0 fix(i18n): add missing OpenAI account entry labels
- admin.accounts.oauth.openai.mobileRefreshTokenAuth was referenced by
  OAuthAuthorizationFlow.vue since 9f8cffe88 (Mobile RT entry) but never
  added to zh/en locales, rendering the raw key in the add-account wizard
- admin.accounts.oauth.openai.accessTokenAuth has the same latent issue
  since 26060e702 (Sora AT import); currently hidden but fixed alongside

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-17 09:32:48 +08:00
yardbirds0 824f763a77 feat: 支持从 API Key 账号名称跳转上游站点 2026-07-17 09:15:32 +08:00
Wesley Liddick 9b3c96dc68 Merge pull request #4451 from rurucxk/fix/grok-free-responses-function-tool-cache
fix(grok): enable free prompt cache for Responses function tools
2026-07-17 09:04:47 +08:00
Wesley Liddick 52f2bc69fb Merge pull request #4450 from wucm667/fix/issue-4442-lazy-stripe-loader
fix(frontend): lazy-load Stripe payment dependency
2026-07-17 09:04:33 +08:00
Wesley Liddick 067cf8b46e Merge pull request #4459 from Wei-Shaw/fix/openai-alpha-search-apikey-scheduling
fix(openai): restore APIKey account scheduling for alpha/search
2026-07-17 09:04:17 +08:00
mt21625457andCursor df9d9e2e40 fix(security-audit): harden role scan, startup, probe, and localhost dial
Scan client-injected assistant/tool/model turns, fail closed when config cannot
be trusted after startup or stale invalidation, reuse probe tokens only for the
same base URL, and restrict localhost dials to loopback addresses.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-17 09:00:14 +08:00
mt21625457andCursor 0f7f8a317e fix(security-audit): close prompt-audit bypass and privacy gaps
Stop WebSocket follow-up turns from reusing a request-wide audit cache, scan
client-controlled instruction fields, fail closed on stale weaker configs, and
tighten preview/SSRF controls including persisted request stage.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-17 08:46:57 +08:00
shaw d2b080e88b fix(openai): restore APIKey account scheduling for alpha/search
PR #4394 (776f3f0de) 在新增 alpha_search 调度能力时加入了 OAuth-only
门控,把 APIKey 账号从 /alpha/search 候选池整体剔除;但转发层自
52071d391 起就支持 APIKey 走 {base_url}/v1/alpha/search,门控注释中
「API key 会被发往 chatgpt.com 导致 401」与实际路由不符。结果是纯
APIKey 分组自 v0.1.157 起独立搜索在选号阶段即失败(无可用账号)。

修复:
- SupportsOpenAIEndpointCapability 的 alpha_search 门控放行
  AccountTypeAPIKey(OAuth/APIKey 均可,Grok 等仍拒绝);显式能力集
  语义不变,chat_completions 继续隐含放行 alpha_search。
- ForwardAlphaSearch 对 APIKey 账号的 404/405 按端点级 failover 处理
  (换号、不写账号错误状态),避免混合分组里请求死在不支持该端点的
  APIKey 上游;OAuth 账号 404 透传行为保持不变。
- 更新固化旧门控行为的用例,并在调度层新增 APIKey 放行回归锁。
2026-07-17 08:44:23 +08:00
mt21625457 d11bdb13f5 feat(security-audit): add OpenAI-compatible prompt auditing 2026-07-17 00:39:39 +08:00
rurucxk 29e39b96e5 fix(grok): enable free prompt cache for Responses function tools
Reuse applyGrokFreeMessagesFunctionToolCacheRoute on native /v1/responses
and the Grok WS HTTP bridge so Free OAuth requests with client function
tools get the same mixed-tools cache route as the Messages bridge
(append/convert web_search and x_search).

Also dedupe: Grok Build already declares function tools named web_search,
so naive append caused "Duplicate tool names: web_search". Convert those
function entries to native tool types and skip duplicates.

Only Free OAuth accounts (isKnownGrokFreeAccount); paid/unknown unchanged.
2026-07-17 00:23:19 +08:00
wucm667 08e994ad86 fix(frontend): lazy-load Stripe payment dependency 2026-07-16 23:21:58 +08:00
github-actions[bot] bc2244c83f chore: sync VERSION to 0.1.158 [skip ci] 2026-07-16 12:37:21 +00:00
Wesley Liddick 26abd19a28 Merge pull request #4433 from heathermhuang/codex/fix-grok-regressions-4412-4421
fix(grok): repair OAuth media and compatibility regressions
2026-07-16 20:16:46 +08:00
shaw c29b50394f Merge main (fix batch-limits test typecheck breaking CI) 2026-07-16 20:05:48 +08:00
shaw 3d23cc399f fix(admin): align batch-limits test with expanded NewUserHandler signature
PR #4425 was authored before #4429 widened NewUserHandler with the
step-up TOTP and user services, and merged without a rebase, breaking
typecheck on main.
2026-07-16 20:05:31 +08:00
shaw c1fd1cb44e Merge origin/main into codex/fix-grok-regressions-4412-4421
Reconcile the OAuth media route with the manual endpoint-switch redesign
(7f5d067af): media leaves for api.x.ai only when text traffic resolves to
the CLI gateway host; manually selected official/regional/custom endpoints
keep serving media as-is.
2026-07-16 19:57:48 +08:00
Wesley Liddick 7e13b6d039 Merge pull request #4425 from AdrianZhaoDev/agent/admin-users-batch-limits
feat(admin): batch update user concurrency and RPM
2026-07-16 19:33:32 +08:00
Wesley Liddick c993490a82 Merge pull request #4434 from yan9651688/feat/group-one-click-copy
feat(group): add safe one-click duplication
2026-07-16 19:33:18 +08:00
Wesley Liddick 4d91f39474 Merge pull request #4431 from linyqh/codex/fix-codex-image-handoff
fix(openai): 修复 Codex 生图网关侧交接问题
2026-07-16 19:32:33 +08:00
Wesley Liddick af6bd44d2a Merge pull request #4435 from superman2003/fix/grok-codex-wsv2-template
fix(grok): enable Codex Responses WebSocket v2 in setup template
2026-07-16 19:32:14 +08:00
Wesley Liddick 541ca3bc27 Merge pull request #4430 from feitianbubu/fix/proxy-fallback-label
fix(i18n): 代理"失败回退"改名"到期回退",与实际行为一致
2026-07-16 19:32:04 +08:00
shaw 7f5d067af2 feat(grok): 支持上游端点手动切换与快捷端点,修复 SSO 建号自定义地址被覆盖
官方端点(api.x.ai / cli-chat-proxy.grok.com)偶发不可用,运营方需要在
端点间手动切换。旧语义把 OAuth 账号存储的官方 host 一律视同"未定制"并
回落默认 CLI 网关:填了官方地址保存成功却不生效、重新编辑开关回到关闭、
再次保存直接删值,形成"修改不生效"的静默循环。

后端:
- GetGrokBaseURL OAuth 分支改为"存了什么用什么":官方 API / 区域 API /
  第三方转发地址一律按填写值转发与探测,仅空值或无法解析的脏数据回落
  默认 CLI 网关;删除官方变体运行时迁移逻辑
- *.api.x.ai 区域端点(us-east-1/us-west-2/eu-west-1 等)纳入可信 host,
  OAuth 使用时不受运营方 URL 白名单限制,官方 host 仍强制 /v1 path
- 修复 SSO 批量建号 MergeCredentials 方向缺陷:BuildAccountCredentials
  恒写官方 base_url,会覆盖导入请求指定的自定义转发地址;抽出
  grokSSOImportCredentials 显式保留请求值(与 RefreshAccountToken 对齐)

前端:
- isCustomGrokBaseUrl 仅默认 CLI 网关 host 视同未定制:api.x.ai 与区域
  端点保存后正常回显(开关开启 + 显示地址),不再被静默吞掉
- 新增 GrokBaseUrlPresets 快捷端点组件(Grok Build CLI / 官方 API /
  us-east-1 / us-west-2 / eu-west-1),接入编辑(OAuth 自定义区 + apikey
  Base URL)、新增(同前)与批量编辑(所选平台全为 grok 时显示,点击
  自动勾选 base_url);仅快速填充,输入框仍可自由填写任意第三方地址
2026-07-16 19:10:21 +08:00
superman2003 174ea22eeb fix(grok): enable Codex Responses WebSocket v2 2026-07-16 18:30:59 +08:00
yan9651688 9fc006546c Make repeated group setup safer
Admins often recreate groups with the same pricing, routing, and account membership. A server-side duplicate creates an inactive copy for review, preserves eligible account priorities, and recovers ambiguous retries without creating extra groups.

Constraint: Group has no neutral JSON metadata field for durable operation recovery
Constraint: Model routing references account IDs, so copied configuration requires matching bindings
Rejected: Rebuild from the list response | it omits configuration and account priority details
Rejected: Store operation identity in business configuration | it would pollute real group settings
Confidence: high
Scope-risk: moderate
Reversibility: clean
Directive: Keep duplicated groups inactive until an administrator reviews the copied configuration
Tested: Go unit and full tests, go vet, integration-tag compile, frontend Vitest, lint, typecheck, production build, and Playwright duplicate flow
Not-tested: PostgreSQL container integration locally because Docker is unavailable; CI will execute the database-backed suite
2026-07-16 18:18:28 +08:00
shaw b960ec1980 chore: update sponsors 2026-07-16 17:45:13 +08:00
github-actions[bot] 60732a2e8c chore: sync VERSION to 0.1.157 [skip ci] 2026-07-16 09:12:44 +00:00
viccyandOmX 9d1064badd 澄清 Codex 图片桥接与本地执行器边界
配置界面现在区分非 Responses Lite 的 hosted 桥接与客户端本地 image_gen,并准确描述账号 strip 策略。API Key Mode 同时提示完全重启 Codex 后新建 task,以重建客户端工具注册表。

Constraint: 兼容模式继续保持默认,Responses Lite 的本地工具由客户端与账号策略决定
Rejected: 写入 features.image_generation | 当前 Codex 已默认启用且仓库既有配置刻意省略该项
Confidence: high
Scope-risk: narrow
Reversibility: clean
Directive: 图片桥接文案必须区分 hosted 注入、客户端声明和 image-only 路由
Tested: 前端 lint、typecheck、关键套件 93 项、定向套件 43 项
Not-tested: Codex Desktop 真实工具注册表重建

Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-07-16 17:11:17 +08:00
viccyandOmX 2b5944fd6c 防止 WebSocket 生图结果停留在进行中状态
HTTP/SSE 已会在终态事件中修正带结果的图片 item,但 WebSocket 下行路径缺少同一处理。现在四条文本帧出口复用现有归一化,并原样保留 result 数据。

Constraint: 仅修改终态事件中带非空 result 的 image_generation_call
Rejected: 解码并改写二进制帧 | 二进制帧不属于 JSON 事件契约
Confidence: high
Scope-risk: narrow
Reversibility: clean
Directive: 保持 HTTP、SSE 与 WebSocket 的图片终态规则一致
Tested: 后端 unit、integration、go vet
Not-tested: 真实 Codex WebSocket 生图会话

Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-07-16 17:11:00 +08:00
viccyandOmX 43e95dfa45 避免无效模型清单破坏 Codex 能力发现
自定义 API-key 上游可能将普通 OpenAI 模型列表作为 2xx 响应返回。现在只校验稳定 envelope,并把结构错误归类为可重试,以便冷缓存时换号且不污染缓存。

Constraint: Codex manifest schema 会随客户端演进,仅校验顶层 models 数组
Rejected: 在网关合成完整 manifest | 容易与上游 schema 漂移
Confidence: high
Scope-risk: narrow
Reversibility: clean
Directive: 不要在此处校验 models 元素的易变字段
Tested: 后端 unit、integration、go vet
Not-tested: 真实第三方 API-key 上游的在线 failover

Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-07-16 17:10:44 +08:00
feitianbubu ca82fd6e98 fix(i18n): rename proxy fallback label to expiry fallback 2026-07-16 17:09:23 +08:00
Heatherm Huang 115116e8bf fix(grok): repair OAuth routing regressions 2026-07-16 16:52:25 +08:00
zhaozewu 7947619cc3 feat(admin): batch update user limits 2026-07-16 15:37:35 +08:00
450 changed files with 38207 additions and 1853 deletions
+19 -5
View File
@@ -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 \
+31 -2
View File
@@ -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
+12
View File
@@ -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>
## 项目概述
+12
View File
@@ -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
View File
@@ -1 +1 @@
0.1.156
0.1.160
+8
View File
@@ -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() {
+41 -3
View File
@@ -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()
+65 -14
View File
@@ -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()
+5
View File
@@ -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
View File
@@ -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(", ")
+10
View File
@@ -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()
+80
View File
@@ -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))
+29
View File
@@ -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
+6
View File
@@ -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)
}
+12 -3
View File
@@ -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
View File
@@ -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
+27 -23
View File
@@ -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()
+10
View File
@@ -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")),
}
}
+75 -8
View File
@@ -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")
}
+62
View File
@@ -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,
+72 -2
View File
@@ -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
+1
View File
@@ -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"`
+76 -5
View File
@@ -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
}
+107 -15
View File
@@ -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) != ""
}
+113
View File
@@ -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
View File
@@ -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")
+6 -5
View File
@@ -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)
+2 -2
View File
@@ -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)
+146 -79
View File
@@ -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
+65 -4
View File
@@ -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)
}
}
+10 -67
View File
@@ -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"
+45 -42
View File
@@ -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())
}
+4
View File
@@ -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 {
+2
View File
@@ -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"`
+23 -3
View File
@@ -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 {
+20
View File
@@ -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)
+129 -12
View File
@@ -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)
+18 -21
View File
@@ -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