Compare commits

...
Author SHA1 Message Date
shaw c38c5beef8 fix(i18n): CN 余额单元格误引 grokBalance 键致渲染原始 key
CNProviderBalanceCell 无快照占位标签引用 admin.accounts.grokBalance,
但该键实际嵌套在 admin.accounts.usageWindow 下,未命中时 DeepSeek/Kimi
payg 账号的用量窗口列直接显示原始键名。

- cnProviders 块新增自有占位键 balance(zh: '余额 --' / en: 'Balance --'),
  单元格改引该键
- 用 esbuild 实际加载模块对 zh/en 全部 admin 语言包做键路径深度比对:
  除本处外 CN 相关键(静态 14 个 + 动态拼接 accountMode/apiProtocol 系列)
  双语全部对齐;唯一存量漂移是 claudeMaxSimulation 块(zh 嵌套于
  modelRouting 下/en 为 groups 直属),全源码零消费方属死键,不在本次
  范围内处理
2026-08-17 17:44:33 +08:00
shaw 10c8b70203 fix(cn-providers): 修复 CN 分组五项功能缺陷(调度闸门/计费/断开漏记/count_tokens/403)
对已合并 PR #5666 + 分组入口放行后的全量功能审计发现的 P0/P1 修复,
全部先经代码与厂商文档实证再实施:

1. /v1/messages 调度闸门(P0):sanitizeGroupMessagesDispatchFields 对非
   openai 平台恒置 AllowMessagesDispatch=false,而闸门豁免名单只有 grok,
   CN 分组经正常途径创建后恒 403——原生 Anthropic 直通(Claude Code 主用例)
   完全不可达。修复:闸门对 CN 与 grok 同语义豁免;count_tokens 处的内联
   裸检查统一走同一 helper;ResolveMessagesDispatchModel 对 CN 早退,避免
   openai 专属的 gpt-5.x 默认映射发给 CN 上游。

2. 计费候选链(P0):候选链兜底含客户端原始模型名,配合 getFallbackPricing
   的 claude→Sonnet 统一兜底,映射的 CN 模型无价时 CN 流量会按 Claude 原价
   (数倍~数十倍)静默误计,且 usage 日志显示 claude-* 名无从察觉。修复:
   CN 账号的 claude-* 候选仅在显式分组/渠道定价时放行;候选全滤空时按
   ErrModelPricingUnavailable 走零成本+告警落账(顺带修复原空候选错误会
   丢弃整条 usage 记录的次生问题)。

3. 断开/中断漏记(P0,#5148 对齐,惠及 openai 平台):messages/responses/
   chat_completions 三个 handler 的错误路径此前在 err!=nil 时丢弃携带的
   部分 result——客户端断开排水后的完整 usage 被丢,payg 上游照常计费而
   平台漏记(anthropic 网关早有同修复,openai 网关缺失)。修复:错误路径
   result 非空时照常提交 usage;failover 错误恒 result=nil 无重复计费。
   Responses×anthropic 流式转换器同时改为断开后继续排水至流自然结束
   (末尾 message_delta 的 output_tokens 不再丢),finalize 帧补工具名反转
   与客户端工具还原、仅在客户端仍连接时写出。

4. count_tokens(P1,证据修正):经实证三家 Anthropic 兼容层均无
   /v1/messages/count_tokens(DeepSeek 官方文档无此端点且注明
   anthropic-version 被忽略;OpenModel 标注该端点 Anthropic only),
   anthropic 协议转发上游=常态 404,且错误处置缺模型上下文会把不计费的
   探测放大成整账号停调。修复:CN 全协议一律本地 tiktoken 估算(与 Grok
   同方案),删除上游转发死代码。

5. 403 处置(P1):CN 此前落入通用 handleAuthError,单次 HTML 403(CDN/
   代理拦截页)即永久禁用,且 403 在 failover 集里会逐账号重放连环禁用
   整组。修复:CN 与 openai 同口径——HTML 豁免 + 3 次累计 + 临时冷却。

新增回归测试 8 项:闸门豁免(含 openai 仍受控断言)、CN 调度映射空返回、
候选过滤三态、空候选零成本落账、断开排水 usage 完整性、HTML-403 零处罚、
结构化 403 首次临时停调。handler/service 全包测试通过。
2026-08-17 17:28:53 +08:00
shaw 7cdca9e495 feat(groups): 放行 kimi/zhipu/deepseek 平台分组创建入口
PR #5666 引入 CN 平台后,路由/调度/前端类型均已支持 CN 平台分组,但分组
创建入口两头缺失:后端 Create/UpdateGroupRequest 的 platform oneof 白名单
与前端 GroupsView 平台选项都没有三平台,导致 CN 账号「无可用分组」、整条
流量链路不通(composite 不能作为替代:CN 不可为 composite 路由目标)。

- group_handler.go: 两处 oneof 加 kimi/zhipu/deepseek;composite 路由目标
  白名单有意不动(DetectModelPlatform/isConcreteRequestPlatform 均无 CN 分支)
- GroupsView: platformOptions/platformFilterOptions 补三项;两处徽章配色链
  按 platformColors.ts 色系补 CN 分支
- i18n: admin.groups.platforms 补 kimi/zhipu/deepseek 键(zh/en),缺键时
  分组徽章/GroupRPM/RateMultipliers 弹窗/ChannelsView 会渲染原始 key
- GroupBadge: badgeClass/labelClass 补 CN 配色
- 新增表驱动测试:9 平台 Create/Update 全放行、非法值(别名/大小写/空格)
  全拒绝、composite target 对 CN 保持拒绝的守卫
2026-08-17 17:28:51 +08:00
Wesley Liddick e330c243a8 Merge pull request #5666 from Randark-JMT/feat/cn-providers-kimi-zhipu-deepseek
feat: 国产供应商多协议支持(Kimi/Zhipu/DeepSeek 原生 Anthropic 直通 + DeepSeek Responses)与配额/余额监控
2026-08-17 14:55:20 +08:00
Randark e728545382 style: gofmt 迁移测试注释(CI golangci-lint gofmt) 2026-08-17 04:26:22 +00:00
Randark 4b667ccd45 fix(review): 处理 PR #5666 四个阻断项(B1-B4)
- B1: 移除根目录 docker-compose.yml(个人镜像测试产物混入)
- B2: 新增迁移 224 放宽 user_platform_quotas CHECK 至 8 平台,补
  BulkInsertInitial 国产平台集成测试与迁移内容断言测试
- B3: CC/Responses×anthropic 四个上游读循环接入间隔超时泵
  (gateway.stream_data_interval_timeout,默认 180s):上游挂住 SSE
  不发数据也不断连时结束排水/组装、关闭 resp.Body 归还连接池位,
  排水超时仍按已累计 usage 返回(与 messages 主路径同语义)
- B4: 配额/余额探测发起前过 cnValidateProbeURL(与网关转发同一套
  security.url_allowlist 策略),被拒时零出站、API key 不离开本机
2026-08-17 04:07:06 +00:00
Wesley Liddick 396a9d1130 Merge pull request #5703 from InCerryGit/fix/clarify-openai-fast-policy-ui
fix(frontend): clarify OpenAI Fast/Flex policy rules
2026-08-17 09:40:53 +08:00
Wesley Liddick 51016bb03a Merge pull request #5690 from luckydududu/fix/dockerfile-golang-image-follows-gomod
fix(docker): Go 构建镜像跟上 go.mod 的 1.26.6(否则任一 Dockerfile 构建都直接失败)
2026-08-17 09:30:48 +08:00
InCerryGit 22fc0cdbf5 fix(frontend): clarify OpenAI Fast/Flex policy rules 2026-08-16 23:49:21 +08:00
Lucky 11e1e22882 fix(docker): bump Go builder image to 1.26.6 to match go.mod
89d826be2 raised backend/go.mod to `go 1.26.6` and updated the three CI
workflows' version assertions, but left the Go builder image in all three
Dockerfiles pinned at 1.26.5. Since the official golang images set
GOTOOLCHAIN=local, the toolchain is not auto-downloaded and any image build
fails hard at `go mod download`.

CI does not catch this: the workflows build with actions/setup-go, not with
these Dockerfiles.

Also extend the Go-upgrade checklist in DEV_GUIDE.md, which listed only the
CI files -- that omission is why the Dockerfiles were missed.
2026-08-16 06:17:43 +00:00
github-actions[bot] baeac1f3de chore: sync VERSION to 0.1.177 [skip ci] 2026-08-15 13:40:21 +00:00
Randark 901a0439f1 feat: 国产供应商一等支持(Kimi/Zhipu/DeepSeek 多协议 + 配额/余额监控)
后端:
- 协议凭证维度 credentials[api_protocol] ∈ chat_completions(默认)/anthropic/responses(deepseek)
- /v1/messages 零转换直通原生 Anthropic 端点(kimi/zhipu/deepseek),CC/Responses
  入站交叉组合走 apicompat 双向转换链(responses/chat_completions anthropic-native 转发器)
- count_tokens:anthropic 协议透传原生端点;其余 CN 协议本地 tiktoken 估算
- Coding Plan 额度探测(5h/weekly 滚动窗口)+ payg 余额探测(kimi/deepseek),
  deepseek 双币种 CNY+USD 明细,任一币种达标不停调
- 周期任务 [CNBalance] 并发探测 + 预算随工作量放大;响应式 429 冷却到最早窗口
  重置点;余额不足可恢复临时停调;智谱 CREDIT_LIMIT 不污染窗口解析
- CC→anthropic 流式客户端断开后继续排水上游保住 usage 计量

前端:
- 创建/编辑弹窗 account_mode + api_protocol + base_url 联动预设(含 watcher 竞态防护)
- 用量单元格:kimi/zhipu coding 显示 5h/weekly 窗口,kimi/deepseek payg 显示余额,
  多币种并列展示;探测失败保留快照;挂载自动探测 5min 去抖
- 调度阈值设置面板补 kimi/zhipu 平台(对齐后端 AllowedSchedulingThresholdPlatforms)
2026-08-15 10:37:51 +00:00
Wesley Liddick 073e92d171 Merge pull request #5668 from Wei-Shaw/fix/codex-turn-state-and-fingerprint-optin
fix(openai): Codex 回合状态回传、v2 压缩探测与指纹收敛改为 opt-in
2026-08-15 17:10:18 +08:00
shaw 4d9fedee20 fix(test): check type assertions in turn-state provenance tests
errcheck runs with check-type-assertions enabled, so the single-value
assertions on the provenance map values fail lint.
2026-08-15 16:44:30 +08:00
shaw fce41e318f fix(openai): make Codex fingerprint convergence opt-in and cover passthrough
Default codex_fingerprint_mode to off. v0.1.175 treated a missing key as
"session", so upgrading silently rewrote installation/session/thread/turn/
window identifiers for every existing OAuth account that had never configured
this field. The quota regressions in #5555, #5556 and #5582 line up with that
version boundary, with A/B reports that rolling back to v0.1.173 restores
quota. Convergence is now explicit opt-in (#5610).

Only accounts that never set the field change behaviour; explicit off /
device / session / full keep working exactly as configured. That required
flipping the persistence condition in all three account modals from
"!== 'session'" to "!== 'off'": the old rule deleted the key when it equalled
the default, which after the flip would have silently discarded an
administrator's explicit opt-in to session.

Also extend convergence to the passthrough path, which previously left client
identifiers untouched:

- resolve the ids once in forwardOpenAIPassthrough and rewrite
  client_metadata on the raw bytes (gjson extract + sjson splice) because
  passthrough is a hot path that must not fully unmarshal multi-MB bodies;
  a shared core keeps the raw and map variants from drifting
- both request builders apply the staged ids at the same relative position
  (after session isolation, before identity enforcement) so headers and body
  share one id set and turn_id stays consistent
- stage the ids unconditionally, including nil: a failover from a converged
  account to an off account must not leave the previous account's ids behind
2026-08-15 16:35:26 +08:00
shaw 8ae6d8f67e fix(openai): send session-level beta features and probe native compaction v2
OpenAI sunset the legacy unary /responses/compact endpoint (404, #5598,
#5624), so the account "compact probe" in the admin UI kept failing even for
healthy accounts, and the beta-feature negotiation header was only attached
to compaction turns.

Beta features (codex-rs session/mod.rs build_model_client_beta_features_header
+ client.rs build_responses_headers): the header is a session-level constant
attached to every /responses request, the WS handshake and /responses/compact.
Enumerating FEATURES shows no Experimental feature is enabled by default, so a
default install sends exactly "remote_compaction_v2". Mirror that:

- OAuth requests without a client-declared header get the default shape, so we
  no longer produce a "header only on compaction turns" pattern real Codex
  never emits (#5586 chains that strip the header)
- a client-declared header is preserved as-is: non-empty without v2 means the
  user disabled the feature and the gateway must not rewrite that
- native v2 turns (compaction_trigger in body) always ensure v2 is present
- non-OAuth upstreams keep the compaction-turn-only behaviour
- the WS injection sits outside the client-header copy block so prewarm and
  turn handshakes cannot land in different pool compatibility buckets

Compact probe now exercises native v2 (streaming /responses +
compaction_trigger) instead of the dead endpoint. Success requires an actual
compaction output item — scanning output_item.done/added, the terminal
response.output[] and the whole-JSON fallback — so a 2xx that silently drops
the trigger is reported as unsupported (the "got 0 items" class, #5478,
#5648). Probe identity is now UUID-shaped and applies the account's
convergence, matching real traffic on the same endpoint.
2026-08-15 16:35:08 +08:00
shaw 8219dcfc87 fix(openai): relay x-codex-turn-state and guard cross-account echo
Codex captures x-codex-turn-state from /responses SSE, /responses/compact
JSON and the WS handshake (codex-api sse/responses.rs, endpoint/compact.rs),
then echoes it back on later requests of the same turn. The HTTP path dropped
it because the header is not in the generic response allowlist, while the WS
path already relayed it — an inconsistency that broke the protocol chain.

Relay it explicitly at every commit point instead of widening the global
allowlist (which would leak it into Anthropic/Gemini responses):

- streaming, non-streaming and SSE-to-JSON handlers relay it, clearing any
  value left over by a previous failover attempt when upstream sends none
- under the first-output guard the header is only staged; provenance is
  recorded when applyAttemptResponseHeaders actually writes it, because a
  first-output timeout discards the staged headers and the client never
  receives that blob
- record (api key + client session) -> minting account, TTL-bounded with an
  opportunistic sweep, and strip echoes known to come from another account
  before they go upstream. Stripping only: injection is the Claude bridge's
  job. Same-account or unknown provenance passes through unchanged.
2026-08-15 16:34:47 +08:00
Wesley Liddick 1d3b9665c8 Merge pull request #5641 from InCerryGit/fix/issue-5624-remote-compaction-v2
fix(openai): preserve remote compaction v2 responses endpoint
2026-08-15 13:46:31 +08:00
Wesley Liddick c204d33b09 Merge pull request #5649 from lyen1688/feat/group-usage-daily-rollups
feat: 优化分组用量统计
2026-08-15 09:00:16 +08:00
lyen1688 45dcce0e49 修复:隔离分组用量仓储测试时区
分组用量仓储测试原先固定恢复上海时区,覆盖集成测试包初始化的 UTC,并导致后续今日统计多算。统一保存并恢复测试进入前的时区,避免 unit 和 integration 用例相互污染。
2026-08-14 23:48:03 +08:00
lyen1688 89d826be29 修复:解决分组用量 PR 的 CI 失败
恢复分组时区测试进入前的全局时区,避免污染高峰计费和支付租约测试。同步升级 Go 1.26.6 与 nanoid 3.3.18,清除后端和前端安全门禁。
2026-08-14 23:25:31 +08:00
lyen1688 cb7b03795d feat: 优化分组用量统计 2026-08-14 22:34:43 +08:00
InCerryGit a8b9ea22b7 fix(openai): separate native and legacy compaction routing 2026-08-14 18:18:53 +08:00
InCerryGit 9662cff2e7 fix(openai): preserve remote compaction v2 responses endpoint 2026-08-14 16:23:52 +08:00
Wesley Liddick fbfdcef818 Merge pull request #5573 from IanShaw027/fix/grok-long-context-and-media-fallback
fix(grok): 长上下文阶梯不再被 OpenAI 账号开关否决,媒体 ID 不继承文本价
2026-08-13 10:32:07 +08:00
github-actions[bot] 0e82efe489 chore: sync VERSION to 0.1.176 [skip ci] 2026-08-13 01:46:47 +00:00
Wesley Liddick e803e3851c Merge pull request #5559 from seng1e/fix/scheduled-backup-leader-lock
fix(backup): 定时备份加 leader 锁,避免多实例重复备份
2026-08-13 09:24:11 +08:00
Wesley Liddick 5912ae960c Merge pull request #5543 from luckydududu/fix/invalidate-channel-cache-on-group-platform-change
fix(group): 改分组 platform 后失效渠道缓存(否则最长 10 分钟按旧平台计价)
2026-08-13 09:24:00 +08:00
Wesley Liddick e77b4b260e Merge pull request #5504 from fengshao1227/fix/responses-probe-inconclusive-keeps-unknown
fix(openai): 探测响应未跑完时不再落标为「上游不支持 Responses」
2026-08-13 09:23:50 +08:00
IanShaw027 e215c98c2c fix: 账号页自动刷新偏好改为模块初始化时恢复
onMounted 再读一次会覆盖已在 setup 阶段恢复的偏好。
2026-08-13 09:22:41 +08:00
IanShaw027 e29b93a1fb fix: 未知 Grok 文本兜底排除 image/video/audio 等媒体族
带版本号的媒体 ID(grok-2-image-1212、grok-5-video)会被
grok-N* 白名单当成文本,错吃 4.5 token 价。vision 多模态
仍按 token 计;HasIdentifiedTokenPricing 也不放行这些 ID。
2026-08-13 09:22:38 +08:00
IanShaw027 fd82dfd52d fix: Grok 长上下文只跟分组开关,不受 OpenAI 账号开关否决
openai_long_context_billing_enabled 是 OpenAI 账号设置,Grok
账号无法打开。按账号硬传 false 会让分组开关失效,官方 200k
阶梯和渠道多档区间全部塌到第一档。非 OpenAI 不再传账号门闩。
2026-08-13 09:22:34 +08:00
Wesley Liddick 4853662f99 Merge pull request #5540 from luckydududu/fix/channel-pricing-model-name-normalization
fix(channel): 定价冲突检测与定价缓存键的归一化对齐(#4754 修复残留)
2026-08-13 09:12:41 +08:00
Wesley Liddick 0fa577a19c Merge pull request #5571 from IanShaw027/feat/grok-jwt-tier-and-4.6
feat(grok): 以 JWT tier 识别订阅档位并接入 grok-4.6,修正徽章滞后、未知模型零计费与搜索/Voice 边界
2026-08-13 09:08:35 +08:00
IanShaw027 b61e4bcc45 fix: 对齐 JWT 常量 gofmt,并补全 available groups 契约字段
golangci-lint 要求 const 块按类型对齐;分组 DTO 新增
long_context_pricing_enabled 后,契约夹具未同步导致比较失败。
2026-08-13 08:55:56 +08:00
IanShaw027 0ae151a238 fix: 徽章改读 grok_usage_snapshot,增量刷新比较 Grok 快照
后端写入 grok_usage_snapshot,列表却读无 writer 的 grok_quota_snapshot;
增量刷新不比 Grok 快照,档位更新后行不替换。

- 正式快照优先 usage,quota 仅兜底;计划字段只收非空字符串
- buildGrokUsageRefreshKey 接入自动刷新
- 注明网关 Extra 写入依赖每请求独立 Account 副本
2026-08-13 08:38:10 +08:00
IanShaw027 b830bc14d6 fix: 长上下文默认保持开启,并与 OpenAI 账号开关取交集
迁移默认 false 会让存量分组静默丢掉 ≥200k 阶梯与渠道多档价。
OpenAI 路径传入 Group 后,账号级 openai_long_context_billing_enabled 被顶掉。

- 列默认 TRUE,并对存量行回填
- 分组价卡不再手填 token 区间,勾选走官方/预设阶梯
- 分组开关与账号开关 AND;价卡 JSON 损坏时打 warn
2026-08-13 08:38:10 +08:00
IanShaw027 8c4c3c09ce fix: 未登记的 Grok 文本模型回退 grok-4.5 价卡
新文本 ID 在闭集价卡上 fail-closed,请求成功而用量记 0。

- grok / grok-N* / build / composer 未登记别名继承 4.5 价
- imagine/voice/web-search/speech 等非文本族不回退
2026-08-13 08:38:10 +08:00
IanShaw027 678eb22a40 fix: Realtime 仅在观察到音频后计费,并修正标志位求值顺序
原先 elapsed>0 就出账,握手失败也会扣费;随后用 audioObserved,
但又在同一个 return 里先 Load 再收 errCh,标志位恒为 false,会话全部漏计。

- 先等中继结束再读 audioObserved
- 无音频或零时长不出账;每次连接独立 request id
2026-08-13 08:38:10 +08:00
IanShaw027 c4d883b8da feat: Chat 与 Responses 往返保留 x_search,并补 sources 抽取
Chat Completions 的 {"type":"x_search"} 会被 apicompat 丢掉,独立
/x_search 又没带 include 与结构化提示,主路径容易返回空结果仍计费。

- Chat↔Responses 保留 x_search 过滤字段与 tool_choice
- declared 只注册实际存活的 x_search,web_search 选择项仍丢弃
- 上游补 include x_search_call.action.sources 与结构化输出提示
2026-08-13 08:37:57 +08:00
IanShaw027 0de6d7e9ba feat: 新增独立 /x_search,走原生 x_search 并沿用搜索计费
Grok 分组原先只有 /web_search,无法带 handle/日期过滤,也无法走
xAI 的 x_search tool。

- POST /x_search(仅 Grok 分组),复用 web_search 的审计、failover 与按次计费
- 上游 Responses 强制 x_search;计费模型记为 grok-x-search
2026-08-13 07:49:19 +08:00
IanShaw027 363cc4994b fix: SuperGrokPro 用 4.5 窗口区分 Heavy,容量抖动只封单模型
JWT 的 SuperGrokPro 同时覆盖 SuperGrok 与 Heavy,账单月额度又滞后,
账号会误升/误降。multi-agent 容量抖动还会把整号提出调度。

- CanonicalGrokPlan:明确 JWT 优先;模糊档仅采新鲜的 grok-4.5 Responses 窗口
- 配额快照写入 plan_from_45_responses 时间戳,过期信号不用
- engine_overloaded 对 multi-agent 只封当前模型 0.5s–5min
2026-08-13 07:49:14 +08:00
IanShaw027 f3d9491071 feat: 分组支持逐模型定价,并可关闭长上下文阶梯
运营需要按分组覆盖渠道/内置价,且部分套餐不应自动吃 200k 倍率。
原先只能改渠道价卡,分组侧只剩 Voice 三列。

- groups 新增 model_pricing / long_context_pricing_enabled,解析链改为 Group → Channel → 内置
- 关闭长上下文时 token 模型只取最低档;video 按秒计费可写进同一价卡
- 回退价对齐官方卡:4.5 缓存 $0.30、4.3/imagine/audio/search 默认值一并校正
2026-08-13 07:49:08 +08:00
IanShaw027 69648476d4 fix: 账号徽章与用量格按实时档位展示,避免账单滞后误判
订阅失效刷新后凭证已是 free,但残留的 monthly_limit / usage_percent
仍会把用量格锁在付费 7d/30d,徽章也优先信滞后的 billing.plan。

- 凭证 subscription_tier 优先于账单/配额快照
- 用量格先看实时档位,再回退账单指标;Lite 仍走付费条
- 徽章补齐 SuperGrok Lite / Plus / X Basic
2026-08-13 01:23:54 +08:00
IanShaw027 a04ce49016 feat: 新增 grok-4.6 目录、官方定价与请求路径支持
官方目录价:<200k 为 $2 / 缓存读取 $0.50 / $6,≥200k 全部 2 倍。
原先没有独立价卡,/models 宣称可配 reasoning 但请求路径会剥掉 effort。

- 目录与别名接入 grok-4.6 / grok-4.6-latest,默认文本模型仍为 grok-4.5
- 独立回退价卡含缓存读取价与 200k 长上下文倍率
- 保留 reasoning_effort;Chat→Responses 的 cache/vision 桥接同样识别 4.6
2026-08-13 01:23:45 +08:00
IanShaw027 bb9e74285e feat: 从 JWT tier 识别 Grok 订阅档位,刷新后覆盖失效订阅
Grok Build access token 带有数字 tier claim(0=free、1=supergrok、5=heavy、6=lite),
原先只解 email/sub/team_id,refresh 还会把旧档位抄回去,订阅失效后一直显示 Heavy。

- 解码 JWT tier,并归一化 display / header 别名(含 free-tier、SuperGrok Lite)
- 仅从 access token 读取档位;新 JWT 覆盖已存凭证,AT 无 claim 时保留旧值
- 用量与 free-cache 判定以当前 AT 为准,账单月额度仅作降级
2026-08-13 01:23:37 +08:00
seng1eandClaude Opus 4.8 bba6a55e0f fix(backup): 定时备份加 leader 锁,避免多实例重复备份
仓库里所有周期任务都用 tryAcquireSingletonLeaderLock 选主、只让一个实例跑
(ops:*、dashboard:aggregation、subscription:expiry、payment:order:expiry 等),
唯独定时备份 runScheduledBackup 没接这套锁,每个实例都自己跑一遍。

多实例部署下同一时刻:
- 同一个库被 N 次 pg_dump;
- 内存峰值 ×N —— 归档整个读进内存再上传,内存吃紧的节点会直接 OOM;
- 上传的是同一个带时间戳的 key,N 份互相覆盖,白干一场连多余副本都留不下。

实测 3 节点:每天 backup_records 里三条一模一样的记录(同 key、同大小)。

改动:把 BackupService 接进和其它任务一样的 tryAcquireSingletonLeaderLock
(key=backup:scheduled:leader)。只锁定时备份,手动备份(CreateBackup/
StartBackup)不锁;无协调后端(cache/db 均 nil)时不加锁照常跑,单机/单测
行为不变;锁 TTL 35m > 备份自身 30m context 上限,防止大库 dump 中途锁过期。
wire_gen.go 由 wire 重新生成。

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-08-12 21:36:41 +08:00
shaw 5935e674a8 chore: update sponsors 2026-08-12 19:17:04 +08:00
github-actions[bot] ef4f99f292 chore: sync VERSION to 0.1.175 [skip ci] 2026-08-12 11:07:43 +00:00
Lucky 814ecfba7c fix(group): invalidate the channel cache when a group's platform changes
The channel cache holds a groupID -> platform map with a 10 minute TTL, and
only channel Create/Update/Delete call invalidateCache(). Changing a group's
platform through the admin API therefore leaves the cache pointing at the old
platform for up to 10 minutes.

Channel pricing, model mapping and the model whitelist are all matched per
platform, so during that window the lookups silently miss: pricing falls back
to the global LiteLLM price list, renames stop applying and the whitelist
stops restricting. Nothing is logged.

Inject a narrow ChannelCacheInvalidator into the admin service (same shape as
the existing APIKeyAuthCacheInvalidator) and call it from UpdateGroup only when
the platform actually changed. The dependency is optional -- when it is nil the
cache simply rebuilds on TTL expiry, as before.
2026-08-12 03:22:06 +00:00
Lucky bd404c16f2 fix(channel): align pricing conflict detection with the pricing cache key
`validateNoConflictingModels` used `toModelEntry`, which only lowercases,
while `expandPricingToCache` keys the cache with
`normalizeChannelPricingModelName` (lowercase + TrimSpace + `.` -> `-`
for `claude-*`). Two pricing entries that the validator considers distinct
therefore collapse onto the same cache key, and the one written later
silently overwrites the other.

Add `toPricingModelEntry`, which reuses `normalizeChannelPricingModelName`,
and use it for pricing conflict detection. `toModelEntry` is left untouched
for model *mapping* validation, because `expandMappingToCache` keys the
mapping cache with plain `strings.ToLower` -- normalizing mappings the same
way would reject configurations that the cache keeps separate.

This completes the fix for #4754, which added the normalization to the
lookup side only.
2026-08-12 02:17:09 +00:00
li fd9ce53281 fix(openai): 探测响应未跑完时不再落标为「上游不支持 Responses」
ProbeOpenAIAPIKeyResponsesSupport 的 2xx 分支靠「output 里有没有 function_call」
判定上游工具能力,但这只在响应真的跑完时成立。以下两种响应会被误判并落标:

- status=incomplete 且 incomplete_details.reason=max_output_tokens:探测请求自己
  只给了 512 的输出预算,推理型模型可能把预算全烧在 reasoning 上,还没轮到
  function_call 就被截断。
- status=failed:HTTP 200 携带的失败响应(上游瞬时故障)。

这个标记的代价很高:探测只在账号创建/更新时跑一次,写进
accounts.extra.openai_responses_supported 后不会自动重探,网关会长期把该账号的
/v1/responses 请求改走 /v1/chat/completions。对 Codex 客户端意味着 prompt 缓存
前缀被协议转换打散,issue #5371 实测缓存命中率从 ~85% 掉到 ~17%、成本涨 3.5 倍,
持续 32 小时无人察觉,最后是账号余额耗尽被停调度才自己弹回去。

同一个函数对非 2xx 早已写明「上游偶发故障不应误判为不支持」并保守返回 true,
2xx 分支缺的正是同一份保守。本次新增 responsesProbeVerdictIsConclusive:判据不
成立时保持 unknown 不写标记,与网络层失败、响应体读取失败的处理一致;网关按
「现状即证据」继续走 Responses。

status=completed 却只回 reasoning 的上游(火山方舟 coding/v3 × kimi-k2.6,即探测
最初要抓的目标)不受影响,仍落标为不支持。

另:落标为不支持时补一条 slog.Warn,让这次会长期改变成本结构的降级可被运维看到。

Fixes #5371
2026-08-11 10:46:50 +08:00
237 changed files with 14989 additions and 909 deletions
+2 -2
View File
@@ -32,7 +32,7 @@ jobs:
cache-dependency-path: backend/go.sum
- name: Verify Go version
run: |
go version | grep -q 'go1.26.5'
go version | grep -q 'go1.26.6'
- name: Unit tests
working-directory: backend
run: make test-unit
@@ -72,7 +72,7 @@ jobs:
cache-dependency-path: backend/go.sum
- name: Verify Go version
run: |
go version | grep -q 'go1.26.5'
go version | grep -q 'go1.26.6'
- name: golangci-lint
uses: golangci/golangci-lint-action@v9
with:
+1 -1
View File
@@ -115,7 +115,7 @@ jobs:
- name: Verify Go version
run: |
go version | grep -q 'go1.26.5'
go version | grep -q 'go1.26.6'
# Docker setup for GoReleaser
- name: Set up QEMU
+1 -1
View File
@@ -23,7 +23,7 @@ jobs:
cache-dependency-path: backend/go.sum
- name: Verify Go version
run: |
go version | grep -q 'go1.26.5'
go version | grep -q 'go1.26.6'
- name: Run govulncheck
working-directory: backend
run: |
+1 -1
View File
@@ -53,7 +53,7 @@ npm install -g pnpm
### CI 要求
- Go 版本必须是 **1.26.5**:三个 workflow 都用 `go-version-file: backend/go.mod` 取版本,随后硬断言 `go version | grep -q 'go1.26.5'`。升级 Go 时要同时改 `backend/go.mod` 和 `backend-ci.yml`(两处)、`release.yml`、`security-scan.yml` 里的这句断言,否则 CI 会在版本校验步骤直接失败。
- Go 版本必须是 **1.26.6**:三个 workflow 都用 `go-version-file: backend/go.mod` 取版本,随后硬断言 `go version | grep -q 'go1.26.6'`。升级 Go 时要同时改 `backend/go.mod`、`backend-ci.yml`(两处)、`release.yml`、`security-scan.yml` 里的这句断言,**以及三个 Dockerfile 里的 Go 构建镜像**(`Dockerfile` / `deploy/Dockerfile` 的 `ARG GOLANG_IMAGE`、`backend/Dockerfile` 的 `FROM golang:`)。前者漏了 CI 会在版本校验步骤直接失败;**后者漏了 CI 不会报,而是等到有人用这些 Dockerfile 构建时才失败**(`go.mod requires go >= X (running Y; GOTOOLCHAIN=local)`)。
- 前端使用 `pnpm install --frozen-lockfile`,必须提交 `pnpm-lock.yaml`
### 本地测试命令
+1 -1
View File
@@ -8,7 +8,7 @@
# =============================================================================
ARG NODE_IMAGE=node:24-alpine
ARG GOLANG_IMAGE=golang:1.26.5-alpine
ARG GOLANG_IMAGE=golang:1.26.6-alpine
ARG ALPINE_IMAGE=alpine:3.21
ARG POSTGRES_IMAGE=postgres:18-alpine
ARG GOPROXY=https://goproxy.cn,direct
+10
View File
@@ -164,6 +164,16 @@ Please read the following carefully before using this project:
<td><a href="https://hao.ai">hao.ai</a> is a high-speed, stable unified large-model API gateway for developers and teams. With a single API Key and a unified interface, you can access mainstream models such as GPT, Claude, and xAI Grok, with compatibility for common protocols and SDKs including OpenAI and Anthropic. The platform provides model routing, failover, team management, and complete request logs, with model prices as low as 15% of official reference pricing, helping users build AI applications more simply, more reliably, and at lower cost.</td>
</tr>
<tr>
<td width="180"><a href="https://www.swiftproxy.net/?ref=sub2api"><img src="assets/partners/logos/swiftprox.png" alt="Swiftproxy" width="150"></a></td>
<td>Swiftproxy is a high-performance proxy solution built for developers, providing stable and reliable residential and static residential proxy services. With 90M+ clean residential IPs, global coverage, flexible rotation, and precise geo-targeting, it helps projects such as web scraping, AI automation, browser automation, SEO monitoring, and multi-account management overcome access restrictions and improve workflow efficiency. It supports HTTP(S) and SOCKS5 protocols, integrates with popular automation tools like Playwright, Selenium, and Puppeteer, with dynamic proxy traffic that never expires until used and free testing available — <a href="https://www.swiftproxy.net/?ref=sub2api">start your free test now</a>!</td>
</tr>
<tr>
<td width="180"><a href="https://www.duckip.cn/?keyword=cu7oog6y"><img src="assets/partners/logos/duckip.png" alt="DuckIP" width="150"></a></td>
<td><a href="https://www.duckip.cn/?keyword=cu7oog6y">DuckIP</a> - 90M+ global residential network resources across 195+ countries and regions, with rotation and sticky sessions for public data collection, RAG updates, model evaluation, and multi-region data workloads. 🟢Residential Proxy - 20% Off; 🟢Static Residential Proxy - Starting at ¥50.00/IP; 🟢Unlimited Residential Proxy - Starting at ¥19.8/Hour. ✅Get 500M Free Trial.</td>
</tr>
</table>
## Overview
+10
View File
@@ -167,6 +167,16 @@
<td><a href="https://hao.ai">hao.ai</a> 是面向开发者与团队的高速、稳定大模型统一 API 网关。通过一个 API Key 和统一接口,即可接入 GPT、Claude、xAI Grok 等主流模型,兼容 OpenAI、Anthropic 等常用协议与 SDK。平台提供模型路由、故障回退、团队管理及完整调用日志,模型价格低至官方参考价的 1.5 折,帮助用户更简单、更稳定、更低成本地构建 AI 应用。</td>
</tr>
<tr>
<td width="180"><a href="https://www.swiftproxy.net/?ref=sub2api"><img src="assets/partners/logos/swiftprox.png" alt="Swiftproxy" width="150"></a></td>
<td>Swiftproxy 是面向开发者的高性能代理解决方案,提供稳定可靠的住宅代理和静态住宅代理服务。拥有 9000 万+ 纯净住宅 IP,覆盖全球,支持灵活轮换和精准地理定位,帮助网页抓取、AI 自动化、浏览器自动化、SEO 监控和多账号管理等项目突破访问限制,提升工作流效率。支持 HTTP(S) 和 SOCKS5 协议,兼容 Playwright、Selenium、Puppeteer 等主流自动化工具,动态代理流量用完为止永不过期,支持免费测试 — <a href="https://www.swiftproxy.net/?ref=sub2api">立即开始免费测试</a>!</td>
</tr>
<tr>
<td width="180"><a href="https://www.duckip.cn/?keyword=cu7oog6y"><img src="assets/partners/logos/duckip.png" alt="DuckIP" width="150"></a></td>
<td><a href="https://www.duckip.cn/?keyword=cu7oog6y">DuckIP</a> - 9000 万+ 全球住宅网络资源,覆盖 195+ 国家和地区,支持轮换和粘性会话,适用于公共数据采集、RAG 更新、模型评估和多区域数据工作负载。🟢住宅代理 - 8 折优惠;🟢静态住宅代理 - ¥50.00/IP 起;🟢无限住宅代理 - ¥19.8/小时 起。✅免费领取 500M 试用流量。</td>
</tr>
</table>
## 项目概述
+10
View File
@@ -166,6 +166,16 @@
<td><a href="https://hao.ai">hao.ai</a> は、開発者とチーム向けの高速で安定した大規模モデル統合 API ゲートウェイです。1 つの API Key と統一されたインターフェースで、GPT、Claude、xAI Grok などの主要モデルに接続でき、OpenAI や Anthropic などの一般的なプロトコルと SDK に対応しています。プラットフォームはモデルルーティング、フェイルオーバー、チーム管理、完全な呼び出しログを提供し、モデル価格は公式参考価格の 1.5 割から。よりシンプルに、より安定して、より低コストに AI アプリケーションを構築できます。</td>
</tr>
<tr>
<td width="180"><a href="https://www.swiftproxy.net/?ref=sub2api"><img src="assets/partners/logos/swiftprox.png" alt="Swiftproxy" width="150"></a></td>
<td>Swiftproxy は開発者向けの高性能プロキシソリューションで、安定して信頼できるレジデンシャルおよび静的レジデンシャルプロキシサービスを提供します。9,000 万以上のクリーンな住宅 IP を保有し、グローバルカバレッジ、柔軟なローテーション、精密なジオターゲティングにより、Web スクレイピング、AI オートメーション、ブラウザ自動化、SEO モニタリング、マルチアカウント管理などのプロジェクトがアクセス制限を克服し、ワークフロー効率を向上させます。HTTP(S) および SOCKS5 プロトコルに対応し、Playwright、Selenium、Puppeteer などの主要な自動化ツールと統合可能。動的プロキシトラフィックは使い切るまで期限切れなし、無料テストも利用可能 — <a href="https://www.swiftproxy.net/?ref=sub2api">今すぐ無料テストを開始</a>!</td>
</tr>
<tr>
<td width="180"><a href="https://www.duckip.cn/?keyword=cu7oog6y"><img src="assets/partners/logos/duckip.png" alt="DuckIP" width="150"></a></td>
<td><a href="https://www.duckip.cn/?keyword=cu7oog6y">DuckIP</a> - 195 以上の国と地域にわたる 9,000 万以上のグローバルレジデンシャルネットワークリソース。ローテーションとスティッキーセッションに対応し、パブリックデータ収集、RAG 更新、モデル評価、マルチリージョンデータワークロードに最適。🟢レジデンシャルプロキシ - 20% オフ;🟢スタティックレジデンシャルプロキシ - ¥50.00/IP から;🟢無制限レジデンシャルプロキシ - ¥19.8/時間 から。✅500M 無料トライアルを取得。</td>
</tr>
</table>
## 概要
Binary file not shown.

After

Width:  |  Height:  |  Size: 12 KiB

-32
View File
@@ -1,32 +0,0 @@
<svg
width="340"
height="112"
viewBox="0 0 340 112"
fill="none"
xmlns="http://www.w3.org/2000/svg"
role="img"
aria-label="Hao.ai"
>
<!-- Fill the compact landscape canvas while preserving the homepage lockup proportions. -->
<g transform="translate(170 56) scale(1.45) translate(-170 -56)">
<!-- Homepage HaoAI mark, rendered in the brand green. -->
<g transform="translate(66.55 12) scale(0.55)">
<path
d="M49 30C42.3726 30 37 35.3726 37 42V118C37 124.627 42.3726 130 49 130C55.6274 130 61 124.627 61 118V42C61 35.3726 55.6274 30 49 30ZM111 30C104.373 30 99 35.3726 99 42V118C99 124.627 104.373 130 111 130C117.627 130 123 124.627 123 118V42C123 35.3726 117.627 30 111 30ZM80 92.5C86.9036 92.5 92.5 86.9036 92.5 80C92.5 73.0964 86.9036 67.5 80 67.5C73.0964 67.5 67.5 73.0964 67.5 80C67.5 86.9036 73.0964 92.5 80 92.5Z"
fill="#1DB954"
fill-rule="evenodd"
clip-rule="evenodd"
/>
</g>
<!-- Homepage wordmark converted from Inter SemiBold to vector outlines. -->
<g transform="translate(150.35 69) scale(0.0166015625 -0.0166015625)" fill="#111111">
<path d="M150 0L150 1490L417 1490L417 874L1110 874L1110 1490L1377 1490L1377 0L1110 0L1110 649L417 649L417 0Z" />
<path transform="translate(1527 0)" d="M450 -23Q344 -23 259 15.5Q174 54 125 129Q76 204 76 314Q76 409 111.5 471Q147 533 207.5 570Q268 607 344 626Q420 645 502 653Q599 664 660 672Q721 680 749.5 698Q778 716 778 754L778 759Q778 814 756 852.5Q734 891 690.5 911.5Q647 932 583 932Q518 932 470.5 912Q423 892 393 860Q363 828 349 791L108 840Q141 939 210 1003.5Q279 1068 374.5 1100Q470 1132 582 1132Q661 1132 741.5 1113.5Q822 1095 889 1051Q956 1007 997 933Q1038 859 1038 747L1038 0L789 0L789 154L779 154Q754 107 710.5 66.5Q667 26 602.5 1.5Q538 -23 450 -23ZM517 170Q598 170 656.5 202Q715 234 747.5 286.5Q780 339 780 401L780 533Q767 523 736.5 514Q706 505 668.5 498Q631 491 594.5 486Q558 481 533 478Q475 470 428 452Q381 434 354.5 401.5Q328 369 328 317Q328 269 352.5 236.5Q377 204 419.5 187Q462 170 517 170Z" />
<path transform="translate(2703 0)" d="M623 -23Q460 -23 340 49.5Q220 122 154.5 251.5Q89 381 89 553Q89 727 154.5 857.5Q220 988 340 1060Q460 1132 623 1132Q787 1132 907 1060Q1027 988 1092.5 857.5Q1158 727 1158 553Q1158 381 1092.5 251.5Q1027 122 907 49.5Q787 -23 623 -23ZM623 188Q714 188 774.5 237Q835 286 864.5 369Q894 452 894 554Q894 657 864.5 740Q835 823 774.5 872Q714 921 623 921Q532 921 472.5 872Q413 823 383.5 740.5Q354 658 354 554Q354 452 383.5 369Q413 286 472.5 237Q532 188 623 188Z" />
<path transform="translate(3950 0)" fill="#1DB954" d="M326 -17Q259 -17 212.5 29Q166 75 166 143Q166 210 212.5 256Q259 302 326 302Q394 302 440.5 256Q487 210 487 143Q487 75 440.5 29Q394 -17 326 -17Z" />
<path transform="translate(4603 0)" d="M450 -23Q344 -23 259 15.5Q174 54 125 129Q76 204 76 314Q76 409 111.5 471Q147 533 207.5 570Q268 607 344 626Q420 645 502 653Q599 664 660 672Q721 680 749.5 698Q778 716 778 754L778 759Q778 814 756 852.5Q734 891 690.5 911.5Q647 932 583 932Q518 932 470.5 912Q423 892 393 860Q363 828 349 791L108 840Q141 939 210 1003.5Q279 1068 374.5 1100Q470 1132 582 1132Q661 1132 741.5 1113.5Q822 1095 889 1051Q956 1007 997 933Q1038 859 1038 747L1038 0L789 0L789 154L779 154Q754 107 710.5 66.5Q667 26 602.5 1.5Q538 -23 450 -23ZM517 170Q598 170 656.5 202Q715 234 747.5 286.5Q780 339 780 401L780 533Q767 523 736.5 514Q706 505 668.5 498Q631 491 594.5 486Q558 481 533 478Q475 470 428 452Q381 434 354.5 401.5Q328 369 328 317Q328 269 352.5 236.5Q377 204 419.5 187Q462 170 517 170Z" />
<path transform="translate(5779 0)" d="M138 0L138 1118L398 1118L398 0ZM268 1276Q206 1276 161.5 1318Q117 1360 117 1418Q117 1478 161.5 1519.5Q206 1561 268 1561Q331 1561 375.5 1519.5Q420 1478 420 1419Q420 1360 375.5 1318Q331 1276 268 1276Z" />
</g>
</g>
</svg>

Before

Width:  |  Height:  |  Size: 3.7 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 10 KiB

+1 -1
View File
@@ -1,4 +1,4 @@
FROM golang:1.26.5-alpine
FROM golang:1.26.6-alpine
WORKDIR /app
+1 -1
View File
@@ -1 +1 @@
0.1.173
0.1.177
+7
View File
@@ -88,6 +88,7 @@ func provideCleanup(
schedulerSnapshot *service.SchedulerSnapshotService,
tokenRefresh *service.TokenRefreshService,
accountExpiry *service.AccountExpiryService,
cnProviderBalanceCheck *service.CNProviderBalanceCheckService,
codexVersionSync *service.OpenAICodexVersionSyncService,
proxyExpiry *service.ProxyExpiryService,
subscriptionExpiry *service.SubscriptionExpiryService,
@@ -238,6 +239,12 @@ func provideCleanup(
accountExpiry.Stop()
return nil
}},
{"CNProviderBalanceCheckService", func() error {
if cnProviderBalanceCheck != nil {
cnProviderBalanceCheck.Stop()
}
return nil
}},
{"OpenAICodexVersionSyncService", func() error {
codexVersionSync.Stop()
return nil
+15 -4
View File
@@ -189,7 +189,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
adminAccountRepository := repository.NewAdminAccountRepository(client, db, schedulerCache)
proxyExitInfoProber := repository.NewProxyExitInfoProber(configConfig)
proxyLatencyCache := repository.NewProxyLatencyCache(redisClient)
adminService := service.NewAdminService(userRepository, adminGroupRepository, adminAccountRepository, proxyRepository, apiKeyRepository, redeemCodeRepository, userGroupRateRepository, userRPMCache, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator, client, settingService, subscriptionService, userSubscriptionRepository, privacyClientFactory, openAIGatewayService, affiliateService, compositeModelRouteRepository, compositeRouteResolver)
adminService := service.NewAdminService(userRepository, adminGroupRepository, adminAccountRepository, proxyRepository, apiKeyRepository, redeemCodeRepository, userGroupRateRepository, userRPMCache, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator, client, settingService, subscriptionService, userSubscriptionRepository, privacyClientFactory, openAIGatewayService, affiliateService, compositeModelRouteRepository, compositeRouteResolver, channelService)
adminUserHandler := admin.NewUserHandler(adminService, concurrencyService, serviceUserPlatformQuotaRepository, billingCache, totpService, userService, settingService)
groupCapacityService := service.NewGroupCapacityService(accountRepository, groupRepository, concurrencyService, sessionLimitCache, rpmCache)
groupHandler := admin.NewGroupHandler(adminService, dashboardService, groupCapacityService)
@@ -208,7 +208,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
dataManagementHandler := admin.NewDataManagementHandler(dataManagementService)
backupObjectStoreFactory := repository.NewS3BackupStoreFactory()
dbDumper := repository.NewPgDumper(configConfig)
backupService := service.ProvideBackupService(settingRepository, configConfig, secretEncryptor, backupObjectStoreFactory, dbDumper)
backupService := service.ProvideBackupService(settingRepository, configConfig, secretEncryptor, backupObjectStoreFactory, dbDumper, leaderLockCache, db)
imageStorageFactory := repository.ProvideImageStorageFactory()
imageStorageSettingService := service.ProvideImageStorageSettingService(settingRepository, secretEncryptor, backupService, imageStorageFactory, configConfig)
backupHandler := admin.NewBackupHandler(backupService, userService, imageStorageSettingService)
@@ -218,6 +218,9 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
antigravityOAuthHandler := admin.NewAntigravityOAuthHandler(antigravityOAuthService)
tokenRefreshService := service.ProvideTokenRefreshService(accountRepository, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, compositeTokenCacheInvalidator, schedulerCache, configConfig, tempUnschedCache, privacyClientFactory, proxyRepository, oAuthRefreshAPI, openAIGatewayService)
grokOAuthHandler := admin.NewGrokOAuthHandler(grokOAuthService, adminService, grokQuotaService, tokenRefreshService)
cnProviderQuotaService := service.ProvideCNProviderQuotaService(accountRepository, proxyRepository, httpUpstream, configConfig)
cnProviderBalanceService := service.ProvideCNProviderBalanceService(accountRepository, proxyRepository, httpUpstream, configConfig)
cnProviderHandler := admin.NewCNProviderHandler(cnProviderQuotaService, cnProviderBalanceService)
proxyHandler := admin.NewProxyHandler(adminService)
adminRedeemHandler := admin.NewRedeemHandler(adminService, redeemService)
promoHandler := admin.NewPromoHandler(promoService)
@@ -277,7 +280,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
auditLogHandler := admin.NewAuditLogHandler(auditLogService, totpService)
upstreamBillingProbeService := service.ProvideUpstreamBillingProbeService(accountRepository, accountTestService, settingService, leaderLockCache, db)
ollamaCloudUsageService := service.ProvideOllamaCloudUsageService(accountRepository, httpUpstream, settingService, secretEncryptor, configConfig, 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, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService, ollamaCloudUsageService)
adminHandlers := handler.ProvideAdminHandlers(dashboardHandler, adminUserHandler, groupHandler, accountHandler, adminAnnouncementHandler, dataManagementHandler, backupHandler, oAuthHandler, openAIOAuthHandler, geminiOAuthHandler, antigravityOAuthHandler, grokOAuthHandler, cnProviderHandler, proxyHandler, adminRedeemHandler, promoHandler, settingHandler, opsHandler, systemHandler, adminSubscriptionHandler, adminUsageHandler, userAttributeHandler, errorPassthroughHandler, tlsFingerprintProfileHandler, adminAPIKeyHandler, scheduledTestHandler, channelHandler, channelMonitorHandler, channelMonitorRequestTemplateHandler, contentModerationHandler, promptAdminHandler, paymentHandler, affiliateHandler, complianceHandler, auditLogHandler, upstreamBillingProbeService, ollamaCloudUsageService)
usageRecordWorkerPool := service.NewUsageRecordWorkerPool(configConfig)
userMsgQueueCache := repository.NewUserMsgQueueCache(redisClient)
userMessageQueueService := service.ProvideUserMessageQueueService(userMsgQueueCache, rpmCache, configConfig)
@@ -327,6 +330,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
opsScheduledReportService := service.ProvideOpsScheduledReportService(opsService, userService, emailService, redisClient, configConfig)
opsIngressRejectAggregator := service.ProvideOpsIngressRejectAggregator(opsRepository, opsService)
accountExpiryService := service.ProvideAccountExpiryService(accountRepository)
cnProviderBalanceCheckService := service.ProvideCNProviderBalanceCheckService(accountRepository, cnProviderBalanceService, cnProviderQuotaService, configConfig)
openAICodexVersionSyncService := service.ProvideOpenAICodexVersionSyncService(settingRepository, settingService, gitHubReleaseClient)
proxyExpiryService := service.ProvideProxyExpiryService(proxyRepository)
subscriptionExpiryService := service.ProvideSubscriptionExpiryService(userSubscriptionRepository, settingRepository, notificationEmailService, leaderLockCache, db)
@@ -336,7 +340,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
channelMonitorRunner := service.ProvideChannelMonitorRunner(channelMonitorService, settingService)
channelMonitorV2Aggregator := service.ProvideChannelMonitorV2Aggregator(channelMonitorV2Repository, db, settingService)
userPlatformQuotaUsageFlusher := service.ProvideUserPlatformQuotaUsageFlusher(configConfig, billingCache, serviceUserPlatformQuotaRepository, timingWheelService)
v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, openAICodexVersionSyncService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, channelMonitorV2Aggregator, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, ollamaCloudUsageService, auditLogService, promptService)
v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, opsService, opsIngressRejectAggregator, apiKeyService, authCacheInvalidationWorker, schedulerSnapshotService, tokenRefreshService, accountExpiryService, cnProviderBalanceCheckService, openAICodexVersionSyncService, proxyExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, batchImageCleanupService, batchImageWorkerRuntime, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, paymentOrderExpiryService, channelMonitorRunner, channelMonitorV2Aggregator, userPlatformQuotaUsageFlusher, upstreamBillingProbeService, ollamaCloudUsageService, auditLogService, promptService)
application := &Application{
Server: httpServer,
PromptAudit: promptService,
@@ -380,6 +384,7 @@ func provideCleanup(
schedulerSnapshot *service.SchedulerSnapshotService,
tokenRefresh *service.TokenRefreshService,
accountExpiry *service.AccountExpiryService,
cnProviderBalanceCheck *service.CNProviderBalanceCheckService,
codexVersionSync *service.OpenAICodexVersionSyncService,
proxyExpiry *service.ProxyExpiryService,
subscriptionExpiry *service.SubscriptionExpiryService,
@@ -529,6 +534,12 @@ func provideCleanup(
accountExpiry.Stop()
return nil
}},
{"CNProviderBalanceCheckService", func() error {
if cnProviderBalanceCheck != nil {
cnProviderBalanceCheck.Stop()
}
return nil
}},
{"OpenAICodexVersionSyncService", func() error {
codexVersionSync.Stop()
return nil
+1
View File
@@ -66,6 +66,7 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) {
schedulerSnapshotSvc,
tokenRefreshSvc,
accountExpirySvc,
nil, // cnProviderBalanceCheck
codexVersionSyncSvc,
proxyExpirySvc,
subscriptionExpirySvc,
+26 -2
View File
@@ -97,6 +97,10 @@ type Group struct {
AudioTtsPricePerMillionChars *float64 `json:"audio_tts_price_per_million_chars,omitempty"`
// STT 每小时价格(USD)
AudioSttPricePerHour *float64 `json:"audio_stt_price_per_hour,omitempty"`
// 是否按上下文长度应用模型阶梯价格
LongContextPricingEnabled bool `json:"long_context_pricing_enabled,omitempty"`
// 分组逐模型定价;优先级高于渠道和内置定价
ModelPricing json.RawMessage `json:"model_pricing,omitempty"`
// 是否仅允许 Claude Code 客户端
ClaudeCodeOnly bool `json:"claude_code_only,omitempty"`
// 非 Claude Code 请求降级使用的分组 ID
@@ -245,9 +249,9 @@ func (*Group) scanValues(columns []string) ([]any, error) {
values := make([]any, len(columns))
for i := range columns {
switch columns[i] {
case group.FieldVideoModelPrices, group.FieldModelRouting, group.FieldSupportedModelScopes, group.FieldMessagesDispatchModelConfig, group.FieldModelsListConfig, group.FieldReasoningEffortMappings:
case group.FieldVideoModelPrices, group.FieldModelPricing, group.FieldModelRouting, group.FieldSupportedModelScopes, group.FieldMessagesDispatchModelConfig, group.FieldModelsListConfig, group.FieldReasoningEffortMappings:
values[i] = new([]byte)
case group.FieldPeakRateEnabled, group.FieldIsExclusive, group.FieldAllowImageGeneration, group.FieldAllowBatchImageGeneration, group.FieldImageRateIndependent, group.FieldVideoRateIndependent, group.FieldClaudeCodeOnly, group.FieldModelRoutingEnabled, group.FieldMcpXMLInject, group.FieldAllowMessagesDispatch, group.FieldAllowLive, group.FieldRequireOauthOnly, group.FieldRequirePrivacySet, group.FieldProfitControlEnabled:
case group.FieldPeakRateEnabled, group.FieldIsExclusive, group.FieldAllowImageGeneration, group.FieldAllowBatchImageGeneration, group.FieldImageRateIndependent, group.FieldVideoRateIndependent, group.FieldLongContextPricingEnabled, group.FieldClaudeCodeOnly, group.FieldModelRoutingEnabled, group.FieldMcpXMLInject, group.FieldAllowMessagesDispatch, group.FieldAllowLive, group.FieldRequireOauthOnly, group.FieldRequirePrivacySet, group.FieldProfitControlEnabled:
values[i] = new(sql.NullBool)
case group.FieldRateMultiplier, group.FieldPeakRateMultiplier, group.FieldDailyLimitUsd, group.FieldWeeklyLimitUsd, group.FieldMonthlyLimitUsd, group.FieldImageRateMultiplier, group.FieldImagePrice1k, group.FieldImagePrice2k, group.FieldImagePrice4k, group.FieldBatchImageDiscountMultiplier, group.FieldBatchImageHoldMultiplier, group.FieldVideoRateMultiplier, group.FieldVideoPrice480p, group.FieldVideoPrice720p, group.FieldVideoPrice1080p, group.FieldWebSearchPricePerCall, group.FieldSearchPricePer1k, group.FieldAudioRealtimePricePerMin, group.FieldAudioTtsPricePerMillionChars, group.FieldAudioSttPricePerHour, group.FieldProfitMinMargin, group.FieldProfitSafetyBuffer:
values[i] = new(sql.NullFloat64)
@@ -531,6 +535,20 @@ func (_m *Group) assignValues(columns []string, values []any) error {
_m.AudioSttPricePerHour = new(float64)
*_m.AudioSttPricePerHour = value.Float64
}
case group.FieldLongContextPricingEnabled:
if value, ok := values[i].(*sql.NullBool); !ok {
return fmt.Errorf("unexpected type %T for field long_context_pricing_enabled", values[i])
} else if value.Valid {
_m.LongContextPricingEnabled = value.Bool
}
case group.FieldModelPricing:
if value, ok := values[i].(*[]byte); !ok {
return fmt.Errorf("unexpected type %T for field model_pricing", values[i])
} else if value != nil && len(*value) > 0 {
if err := json.Unmarshal(*value, &_m.ModelPricing); err != nil {
return fmt.Errorf("unmarshal field model_pricing: %w", err)
}
}
case group.FieldClaudeCodeOnly:
if value, ok := values[i].(*sql.NullBool); !ok {
return fmt.Errorf("unexpected type %T for field claude_code_only", values[i])
@@ -896,6 +914,12 @@ func (_m *Group) String() string {
builder.WriteString(fmt.Sprintf("%v", *v))
}
builder.WriteString(", ")
builder.WriteString("long_context_pricing_enabled=")
builder.WriteString(fmt.Sprintf("%v", _m.LongContextPricingEnabled))
builder.WriteString(", ")
builder.WriteString("model_pricing=")
builder.WriteString(fmt.Sprintf("%v", _m.ModelPricing))
builder.WriteString(", ")
builder.WriteString("claude_code_only=")
builder.WriteString(fmt.Sprintf("%v", _m.ClaudeCodeOnly))
builder.WriteString(", ")
+13
View File
@@ -94,6 +94,10 @@ const (
FieldAudioTtsPricePerMillionChars = "audio_tts_price_per_million_chars"
// FieldAudioSttPricePerHour holds the string denoting the audio_stt_price_per_hour field in the database.
FieldAudioSttPricePerHour = "audio_stt_price_per_hour"
// FieldLongContextPricingEnabled holds the string denoting the long_context_pricing_enabled field in the database.
FieldLongContextPricingEnabled = "long_context_pricing_enabled"
// FieldModelPricing holds the string denoting the model_pricing field in the database.
FieldModelPricing = "model_pricing"
// FieldClaudeCodeOnly holds the string denoting the claude_code_only field in the database.
FieldClaudeCodeOnly = "claude_code_only"
// FieldFallbackGroupID holds the string denoting the fallback_group_id field in the database.
@@ -250,6 +254,8 @@ var Columns = []string{
FieldAudioRealtimePricePerMin,
FieldAudioTtsPricePerMillionChars,
FieldAudioSttPricePerHour,
FieldLongContextPricingEnabled,
FieldModelPricing,
FieldClaudeCodeOnly,
FieldFallbackGroupID,
FieldFallbackGroupIDOnInvalidRequest,
@@ -364,6 +370,8 @@ var (
AudioTtsPricePerMillionCharsValidator func(float64) error
// AudioSttPricePerHourValidator is a validator for the "audio_stt_price_per_hour" field. It is called by the builders before save.
AudioSttPricePerHourValidator func(float64) error
// DefaultLongContextPricingEnabled holds the default value on creation for the "long_context_pricing_enabled" field.
DefaultLongContextPricingEnabled bool
// DefaultClaudeCodeOnly holds the default value on creation for the "claude_code_only" field.
DefaultClaudeCodeOnly bool
// DefaultModelRoutingEnabled holds the default value on creation for the "model_routing_enabled" field.
@@ -604,6 +612,11 @@ func ByAudioSttPricePerHour(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldAudioSttPricePerHour, opts...).ToFunc()
}
// ByLongContextPricingEnabled orders the results by the long_context_pricing_enabled field.
func ByLongContextPricingEnabled(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldLongContextPricingEnabled, opts...).ToFunc()
}
// ByClaudeCodeOnly orders the results by the claude_code_only field.
func ByClaudeCodeOnly(opts ...sql.OrderTermOption) OrderOption {
return sql.OrderByField(FieldClaudeCodeOnly, opts...).ToFunc()
+25
View File
@@ -245,6 +245,11 @@ func AudioSttPricePerHour(v float64) predicate.Group {
return predicate.Group(sql.FieldEQ(FieldAudioSttPricePerHour, v))
}
// LongContextPricingEnabled applies equality check predicate on the "long_context_pricing_enabled" field. It's identical to LongContextPricingEnabledEQ.
func LongContextPricingEnabled(v bool) predicate.Group {
return predicate.Group(sql.FieldEQ(FieldLongContextPricingEnabled, v))
}
// ClaudeCodeOnly applies equality check predicate on the "claude_code_only" field. It's identical to ClaudeCodeOnlyEQ.
func ClaudeCodeOnly(v bool) predicate.Group {
return predicate.Group(sql.FieldEQ(FieldClaudeCodeOnly, v))
@@ -2045,6 +2050,26 @@ func AudioSttPricePerHourNotNil() predicate.Group {
return predicate.Group(sql.FieldNotNull(FieldAudioSttPricePerHour))
}
// LongContextPricingEnabledEQ applies the EQ predicate on the "long_context_pricing_enabled" field.
func LongContextPricingEnabledEQ(v bool) predicate.Group {
return predicate.Group(sql.FieldEQ(FieldLongContextPricingEnabled, v))
}
// LongContextPricingEnabledNEQ applies the NEQ predicate on the "long_context_pricing_enabled" field.
func LongContextPricingEnabledNEQ(v bool) predicate.Group {
return predicate.Group(sql.FieldNEQ(FieldLongContextPricingEnabled, v))
}
// ModelPricingIsNil applies the IsNil predicate on the "model_pricing" field.
func ModelPricingIsNil() predicate.Group {
return predicate.Group(sql.FieldIsNull(FieldModelPricing))
}
// ModelPricingNotNil applies the NotNil predicate on the "model_pricing" field.
func ModelPricingNotNil() predicate.Group {
return predicate.Group(sql.FieldNotNull(FieldModelPricing))
}
// ClaudeCodeOnlyEQ applies the EQ predicate on the "claude_code_only" field.
func ClaudeCodeOnlyEQ(v bool) predicate.Group {
return predicate.Group(sql.FieldEQ(FieldClaudeCodeOnly, v))
+136
View File
@@ -4,6 +4,7 @@ package ent
import (
"context"
"encoding/json"
"errors"
"fmt"
"time"
@@ -559,6 +560,26 @@ func (_c *GroupCreate) SetNillableAudioSttPricePerHour(v *float64) *GroupCreate
return _c
}
// SetLongContextPricingEnabled sets the "long_context_pricing_enabled" field.
func (_c *GroupCreate) SetLongContextPricingEnabled(v bool) *GroupCreate {
_c.mutation.SetLongContextPricingEnabled(v)
return _c
}
// SetNillableLongContextPricingEnabled sets the "long_context_pricing_enabled" field if the given value is not nil.
func (_c *GroupCreate) SetNillableLongContextPricingEnabled(v *bool) *GroupCreate {
if v != nil {
_c.SetLongContextPricingEnabled(*v)
}
return _c
}
// SetModelPricing sets the "model_pricing" field.
func (_c *GroupCreate) SetModelPricing(v json.RawMessage) *GroupCreate {
_c.mutation.SetModelPricing(v)
return _c
}
// SetClaudeCodeOnly sets the "claude_code_only" field.
func (_c *GroupCreate) SetClaudeCodeOnly(v bool) *GroupCreate {
_c.mutation.SetClaudeCodeOnly(v)
@@ -1042,6 +1063,10 @@ func (_c *GroupCreate) defaults() error {
v := group.DefaultVideoRateMultiplier
_c.mutation.SetVideoRateMultiplier(v)
}
if _, ok := _c.mutation.LongContextPricingEnabled(); !ok {
v := group.DefaultLongContextPricingEnabled
_c.mutation.SetLongContextPricingEnabled(v)
}
if _, ok := _c.mutation.ClaudeCodeOnly(); !ok {
v := group.DefaultClaudeCodeOnly
_c.mutation.SetClaudeCodeOnly(v)
@@ -1237,6 +1262,9 @@ func (_c *GroupCreate) check() error {
return &ValidationError{Name: "audio_stt_price_per_hour", err: fmt.Errorf(`ent: validator failed for field "Group.audio_stt_price_per_hour": %w`, err)}
}
}
if _, ok := _c.mutation.LongContextPricingEnabled(); !ok {
return &ValidationError{Name: "long_context_pricing_enabled", err: errors.New(`ent: missing required field "Group.long_context_pricing_enabled"`)}
}
if _, ok := _c.mutation.ClaudeCodeOnly(); !ok {
return &ValidationError{Name: "claude_code_only", err: errors.New(`ent: missing required field "Group.claude_code_only"`)}
}
@@ -1484,6 +1512,14 @@ func (_c *GroupCreate) createSpec() (*Group, *sqlgraph.CreateSpec) {
_spec.SetField(group.FieldAudioSttPricePerHour, field.TypeFloat64, value)
_node.AudioSttPricePerHour = &value
}
if value, ok := _c.mutation.LongContextPricingEnabled(); ok {
_spec.SetField(group.FieldLongContextPricingEnabled, field.TypeBool, value)
_node.LongContextPricingEnabled = value
}
if value, ok := _c.mutation.ModelPricing(); ok {
_spec.SetField(group.FieldModelPricing, field.TypeJSON, value)
_node.ModelPricing = value
}
if value, ok := _c.mutation.ClaudeCodeOnly(); ok {
_spec.SetField(group.FieldClaudeCodeOnly, field.TypeBool, value)
_node.ClaudeCodeOnly = value
@@ -2396,6 +2432,36 @@ func (u *GroupUpsert) ClearAudioSttPricePerHour() *GroupUpsert {
return u
}
// SetLongContextPricingEnabled sets the "long_context_pricing_enabled" field.
func (u *GroupUpsert) SetLongContextPricingEnabled(v bool) *GroupUpsert {
u.Set(group.FieldLongContextPricingEnabled, v)
return u
}
// UpdateLongContextPricingEnabled sets the "long_context_pricing_enabled" field to the value that was provided on create.
func (u *GroupUpsert) UpdateLongContextPricingEnabled() *GroupUpsert {
u.SetExcluded(group.FieldLongContextPricingEnabled)
return u
}
// SetModelPricing sets the "model_pricing" field.
func (u *GroupUpsert) SetModelPricing(v json.RawMessage) *GroupUpsert {
u.Set(group.FieldModelPricing, v)
return u
}
// UpdateModelPricing sets the "model_pricing" field to the value that was provided on create.
func (u *GroupUpsert) UpdateModelPricing() *GroupUpsert {
u.SetExcluded(group.FieldModelPricing)
return u
}
// ClearModelPricing clears the value of the "model_pricing" field.
func (u *GroupUpsert) ClearModelPricing() *GroupUpsert {
u.SetNull(group.FieldModelPricing)
return u
}
// SetClaudeCodeOnly sets the "claude_code_only" field.
func (u *GroupUpsert) SetClaudeCodeOnly(v bool) *GroupUpsert {
u.Set(group.FieldClaudeCodeOnly, v)
@@ -3534,6 +3600,41 @@ func (u *GroupUpsertOne) ClearAudioSttPricePerHour() *GroupUpsertOne {
})
}
// SetLongContextPricingEnabled sets the "long_context_pricing_enabled" field.
func (u *GroupUpsertOne) SetLongContextPricingEnabled(v bool) *GroupUpsertOne {
return u.Update(func(s *GroupUpsert) {
s.SetLongContextPricingEnabled(v)
})
}
// UpdateLongContextPricingEnabled sets the "long_context_pricing_enabled" field to the value that was provided on create.
func (u *GroupUpsertOne) UpdateLongContextPricingEnabled() *GroupUpsertOne {
return u.Update(func(s *GroupUpsert) {
s.UpdateLongContextPricingEnabled()
})
}
// SetModelPricing sets the "model_pricing" field.
func (u *GroupUpsertOne) SetModelPricing(v json.RawMessage) *GroupUpsertOne {
return u.Update(func(s *GroupUpsert) {
s.SetModelPricing(v)
})
}
// UpdateModelPricing sets the "model_pricing" field to the value that was provided on create.
func (u *GroupUpsertOne) UpdateModelPricing() *GroupUpsertOne {
return u.Update(func(s *GroupUpsert) {
s.UpdateModelPricing()
})
}
// ClearModelPricing clears the value of the "model_pricing" field.
func (u *GroupUpsertOne) ClearModelPricing() *GroupUpsertOne {
return u.Update(func(s *GroupUpsert) {
s.ClearModelPricing()
})
}
// SetClaudeCodeOnly sets the "claude_code_only" field.
func (u *GroupUpsertOne) SetClaudeCodeOnly(v bool) *GroupUpsertOne {
return u.Update(func(s *GroupUpsert) {
@@ -4889,6 +4990,41 @@ func (u *GroupUpsertBulk) ClearAudioSttPricePerHour() *GroupUpsertBulk {
})
}
// SetLongContextPricingEnabled sets the "long_context_pricing_enabled" field.
func (u *GroupUpsertBulk) SetLongContextPricingEnabled(v bool) *GroupUpsertBulk {
return u.Update(func(s *GroupUpsert) {
s.SetLongContextPricingEnabled(v)
})
}
// UpdateLongContextPricingEnabled sets the "long_context_pricing_enabled" field to the value that was provided on create.
func (u *GroupUpsertBulk) UpdateLongContextPricingEnabled() *GroupUpsertBulk {
return u.Update(func(s *GroupUpsert) {
s.UpdateLongContextPricingEnabled()
})
}
// SetModelPricing sets the "model_pricing" field.
func (u *GroupUpsertBulk) SetModelPricing(v json.RawMessage) *GroupUpsertBulk {
return u.Update(func(s *GroupUpsert) {
s.SetModelPricing(v)
})
}
// UpdateModelPricing sets the "model_pricing" field to the value that was provided on create.
func (u *GroupUpsertBulk) UpdateModelPricing() *GroupUpsertBulk {
return u.Update(func(s *GroupUpsert) {
s.UpdateModelPricing()
})
}
// ClearModelPricing clears the value of the "model_pricing" field.
func (u *GroupUpsertBulk) ClearModelPricing() *GroupUpsertBulk {
return u.Update(func(s *GroupUpsert) {
s.ClearModelPricing()
})
}
// SetClaudeCodeOnly sets the "claude_code_only" field.
func (u *GroupUpsertBulk) SetClaudeCodeOnly(v bool) *GroupUpsertBulk {
return u.Update(func(s *GroupUpsert) {
+93
View File
@@ -4,6 +4,7 @@ package ent
import (
"context"
"encoding/json"
"errors"
"fmt"
"time"
@@ -787,6 +788,38 @@ func (_u *GroupUpdate) ClearAudioSttPricePerHour() *GroupUpdate {
return _u
}
// SetLongContextPricingEnabled sets the "long_context_pricing_enabled" field.
func (_u *GroupUpdate) SetLongContextPricingEnabled(v bool) *GroupUpdate {
_u.mutation.SetLongContextPricingEnabled(v)
return _u
}
// SetNillableLongContextPricingEnabled sets the "long_context_pricing_enabled" field if the given value is not nil.
func (_u *GroupUpdate) SetNillableLongContextPricingEnabled(v *bool) *GroupUpdate {
if v != nil {
_u.SetLongContextPricingEnabled(*v)
}
return _u
}
// SetModelPricing sets the "model_pricing" field.
func (_u *GroupUpdate) SetModelPricing(v json.RawMessage) *GroupUpdate {
_u.mutation.SetModelPricing(v)
return _u
}
// AppendModelPricing appends value to the "model_pricing" field.
func (_u *GroupUpdate) AppendModelPricing(v json.RawMessage) *GroupUpdate {
_u.mutation.AppendModelPricing(v)
return _u
}
// ClearModelPricing clears the value of the "model_pricing" field.
func (_u *GroupUpdate) ClearModelPricing() *GroupUpdate {
_u.mutation.ClearModelPricing()
return _u
}
// SetClaudeCodeOnly sets the "claude_code_only" field.
func (_u *GroupUpdate) SetClaudeCodeOnly(v bool) *GroupUpdate {
_u.mutation.SetClaudeCodeOnly(v)
@@ -1697,6 +1730,20 @@ func (_u *GroupUpdate) sqlSave(ctx context.Context) (_node int, err error) {
if _u.mutation.AudioSttPricePerHourCleared() {
_spec.ClearField(group.FieldAudioSttPricePerHour, field.TypeFloat64)
}
if value, ok := _u.mutation.LongContextPricingEnabled(); ok {
_spec.SetField(group.FieldLongContextPricingEnabled, field.TypeBool, value)
}
if value, ok := _u.mutation.ModelPricing(); ok {
_spec.SetField(group.FieldModelPricing, field.TypeJSON, value)
}
if value, ok := _u.mutation.AppendedModelPricing(); ok {
_spec.AddModifier(func(u *sql.UpdateBuilder) {
sqljson.Append(u, group.FieldModelPricing, value)
})
}
if _u.mutation.ModelPricingCleared() {
_spec.ClearField(group.FieldModelPricing, field.TypeJSON)
}
if value, ok := _u.mutation.ClaudeCodeOnly(); ok {
_spec.SetField(group.FieldClaudeCodeOnly, field.TypeBool, value)
}
@@ -2862,6 +2909,38 @@ func (_u *GroupUpdateOne) ClearAudioSttPricePerHour() *GroupUpdateOne {
return _u
}
// SetLongContextPricingEnabled sets the "long_context_pricing_enabled" field.
func (_u *GroupUpdateOne) SetLongContextPricingEnabled(v bool) *GroupUpdateOne {
_u.mutation.SetLongContextPricingEnabled(v)
return _u
}
// SetNillableLongContextPricingEnabled sets the "long_context_pricing_enabled" field if the given value is not nil.
func (_u *GroupUpdateOne) SetNillableLongContextPricingEnabled(v *bool) *GroupUpdateOne {
if v != nil {
_u.SetLongContextPricingEnabled(*v)
}
return _u
}
// SetModelPricing sets the "model_pricing" field.
func (_u *GroupUpdateOne) SetModelPricing(v json.RawMessage) *GroupUpdateOne {
_u.mutation.SetModelPricing(v)
return _u
}
// AppendModelPricing appends value to the "model_pricing" field.
func (_u *GroupUpdateOne) AppendModelPricing(v json.RawMessage) *GroupUpdateOne {
_u.mutation.AppendModelPricing(v)
return _u
}
// ClearModelPricing clears the value of the "model_pricing" field.
func (_u *GroupUpdateOne) ClearModelPricing() *GroupUpdateOne {
_u.mutation.ClearModelPricing()
return _u
}
// SetClaudeCodeOnly sets the "claude_code_only" field.
func (_u *GroupUpdateOne) SetClaudeCodeOnly(v bool) *GroupUpdateOne {
_u.mutation.SetClaudeCodeOnly(v)
@@ -3802,6 +3881,20 @@ func (_u *GroupUpdateOne) sqlSave(ctx context.Context) (_node *Group, err error)
if _u.mutation.AudioSttPricePerHourCleared() {
_spec.ClearField(group.FieldAudioSttPricePerHour, field.TypeFloat64)
}
if value, ok := _u.mutation.LongContextPricingEnabled(); ok {
_spec.SetField(group.FieldLongContextPricingEnabled, field.TypeBool, value)
}
if value, ok := _u.mutation.ModelPricing(); ok {
_spec.SetField(group.FieldModelPricing, field.TypeJSON, value)
}
if value, ok := _u.mutation.AppendedModelPricing(); ok {
_spec.AddModifier(func(u *sql.UpdateBuilder) {
sqljson.Append(u, group.FieldModelPricing, value)
})
}
if _u.mutation.ModelPricingCleared() {
_spec.ClearField(group.FieldModelPricing, field.TypeJSON)
}
if value, ok := _u.mutation.ClaudeCodeOnly(); ok {
_spec.SetField(group.FieldClaudeCodeOnly, field.TypeBool, value)
}
+3 -1
View File
@@ -934,6 +934,8 @@ var (
{Name: "audio_realtime_price_per_min", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}},
{Name: "audio_tts_price_per_million_chars", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}},
{Name: "audio_stt_price_per_hour", Type: field.TypeFloat64, Nullable: true, SchemaType: map[string]string{"postgres": "decimal(20,8)"}},
{Name: "long_context_pricing_enabled", Type: field.TypeBool, Default: true},
{Name: "model_pricing", Type: field.TypeJSON, Nullable: true, SchemaType: map[string]string{"postgres": "jsonb"}},
{Name: "claude_code_only", Type: field.TypeBool, Default: false},
{Name: "fallback_group_id", Type: field.TypeInt64, Nullable: true},
{Name: "fallback_group_id_on_invalid_request", Type: field.TypeInt64, Nullable: true},
@@ -990,7 +992,7 @@ var (
{
Name: "group_sort_order",
Unique: false,
Columns: []*schema.Column{GroupsColumns[47]},
Columns: []*schema.Column{GroupsColumns[49]},
},
{
Name: "idx_groups_duplicate_operation_id_active",
+145 -1
View File
@@ -21907,6 +21907,9 @@ type GroupMutation struct {
addaudio_tts_price_per_million_chars *float64
audio_stt_price_per_hour *float64
addaudio_stt_price_per_hour *float64
long_context_pricing_enabled *bool
model_pricing *json.RawMessage
appendmodel_pricing json.RawMessage
claude_code_only *bool
fallback_group_id *int64
addfallback_group_id *int64
@@ -24130,6 +24133,107 @@ func (m *GroupMutation) ResetAudioSttPricePerHour() {
delete(m.clearedFields, group.FieldAudioSttPricePerHour)
}
// SetLongContextPricingEnabled sets the "long_context_pricing_enabled" field.
func (m *GroupMutation) SetLongContextPricingEnabled(b bool) {
m.long_context_pricing_enabled = &b
}
// LongContextPricingEnabled returns the value of the "long_context_pricing_enabled" field in the mutation.
func (m *GroupMutation) LongContextPricingEnabled() (r bool, exists bool) {
v := m.long_context_pricing_enabled
if v == nil {
return
}
return *v, true
}
// OldLongContextPricingEnabled returns the old "long_context_pricing_enabled" 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) OldLongContextPricingEnabled(ctx context.Context) (v bool, err error) {
if !m.op.Is(OpUpdateOne) {
return v, errors.New("OldLongContextPricingEnabled is only allowed on UpdateOne operations")
}
if m.id == nil || m.oldValue == nil {
return v, errors.New("OldLongContextPricingEnabled requires an ID field in the mutation")
}
oldValue, err := m.oldValue(ctx)
if err != nil {
return v, fmt.Errorf("querying old value for OldLongContextPricingEnabled: %w", err)
}
return oldValue.LongContextPricingEnabled, nil
}
// ResetLongContextPricingEnabled resets all changes to the "long_context_pricing_enabled" field.
func (m *GroupMutation) ResetLongContextPricingEnabled() {
m.long_context_pricing_enabled = nil
}
// SetModelPricing sets the "model_pricing" field.
func (m *GroupMutation) SetModelPricing(jm json.RawMessage) {
m.model_pricing = &jm
m.appendmodel_pricing = nil
}
// ModelPricing returns the value of the "model_pricing" field in the mutation.
func (m *GroupMutation) ModelPricing() (r json.RawMessage, exists bool) {
v := m.model_pricing
if v == nil {
return
}
return *v, true
}
// OldModelPricing returns the old "model_pricing" 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) OldModelPricing(ctx context.Context) (v json.RawMessage, err error) {
if !m.op.Is(OpUpdateOne) {
return v, errors.New("OldModelPricing is only allowed on UpdateOne operations")
}
if m.id == nil || m.oldValue == nil {
return v, errors.New("OldModelPricing requires an ID field in the mutation")
}
oldValue, err := m.oldValue(ctx)
if err != nil {
return v, fmt.Errorf("querying old value for OldModelPricing: %w", err)
}
return oldValue.ModelPricing, nil
}
// AppendModelPricing adds jm to the "model_pricing" field.
func (m *GroupMutation) AppendModelPricing(jm json.RawMessage) {
m.appendmodel_pricing = append(m.appendmodel_pricing, jm...)
}
// AppendedModelPricing returns the list of values that were appended to the "model_pricing" field in this mutation.
func (m *GroupMutation) AppendedModelPricing() (json.RawMessage, bool) {
if len(m.appendmodel_pricing) == 0 {
return nil, false
}
return m.appendmodel_pricing, true
}
// ClearModelPricing clears the value of the "model_pricing" field.
func (m *GroupMutation) ClearModelPricing() {
m.model_pricing = nil
m.appendmodel_pricing = nil
m.clearedFields[group.FieldModelPricing] = struct{}{}
}
// ModelPricingCleared returns if the "model_pricing" field was cleared in this mutation.
func (m *GroupMutation) ModelPricingCleared() bool {
_, ok := m.clearedFields[group.FieldModelPricing]
return ok
}
// ResetModelPricing resets all changes to the "model_pricing" field.
func (m *GroupMutation) ResetModelPricing() {
m.model_pricing = nil
m.appendmodel_pricing = nil
delete(m.clearedFields, group.FieldModelPricing)
}
// SetClaudeCodeOnly sets the "claude_code_only" field.
func (m *GroupMutation) SetClaudeCodeOnly(b bool) {
m.claude_code_only = &b
@@ -25435,7 +25539,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, 60)
fields := make([]string, 0, 62)
if m.created_at != nil {
fields = append(fields, group.FieldCreatedAt)
}
@@ -25553,6 +25657,12 @@ func (m *GroupMutation) Fields() []string {
if m.audio_stt_price_per_hour != nil {
fields = append(fields, group.FieldAudioSttPricePerHour)
}
if m.long_context_pricing_enabled != nil {
fields = append(fields, group.FieldLongContextPricingEnabled)
}
if m.model_pricing != nil {
fields = append(fields, group.FieldModelPricing)
}
if m.claude_code_only != nil {
fields = append(fields, group.FieldClaudeCodeOnly)
}
@@ -25702,6 +25812,10 @@ func (m *GroupMutation) Field(name string) (ent.Value, bool) {
return m.AudioTtsPricePerMillionChars()
case group.FieldAudioSttPricePerHour:
return m.AudioSttPricePerHour()
case group.FieldLongContextPricingEnabled:
return m.LongContextPricingEnabled()
case group.FieldModelPricing:
return m.ModelPricing()
case group.FieldClaudeCodeOnly:
return m.ClaudeCodeOnly()
case group.FieldFallbackGroupID:
@@ -25831,6 +25945,10 @@ func (m *GroupMutation) OldField(ctx context.Context, name string) (ent.Value, e
return m.OldAudioTtsPricePerMillionChars(ctx)
case group.FieldAudioSttPricePerHour:
return m.OldAudioSttPricePerHour(ctx)
case group.FieldLongContextPricingEnabled:
return m.OldLongContextPricingEnabled(ctx)
case group.FieldModelPricing:
return m.OldModelPricing(ctx)
case group.FieldClaudeCodeOnly:
return m.OldClaudeCodeOnly(ctx)
case group.FieldFallbackGroupID:
@@ -26155,6 +26273,20 @@ func (m *GroupMutation) SetField(name string, value ent.Value) error {
}
m.SetAudioSttPricePerHour(v)
return nil
case group.FieldLongContextPricingEnabled:
v, ok := value.(bool)
if !ok {
return fmt.Errorf("unexpected type %T for field %s", value, name)
}
m.SetLongContextPricingEnabled(v)
return nil
case group.FieldModelPricing:
v, ok := value.(json.RawMessage)
if !ok {
return fmt.Errorf("unexpected type %T for field %s", value, name)
}
m.SetModelPricing(v)
return nil
case group.FieldClaudeCodeOnly:
v, ok := value.(bool)
if !ok {
@@ -26713,6 +26845,9 @@ func (m *GroupMutation) ClearedFields() []string {
if m.FieldCleared(group.FieldAudioSttPricePerHour) {
fields = append(fields, group.FieldAudioSttPricePerHour)
}
if m.FieldCleared(group.FieldModelPricing) {
fields = append(fields, group.FieldModelPricing)
}
if m.FieldCleared(group.FieldFallbackGroupID) {
fields = append(fields, group.FieldFallbackGroupID)
}
@@ -26790,6 +26925,9 @@ func (m *GroupMutation) ClearField(name string) error {
case group.FieldAudioSttPricePerHour:
m.ClearAudioSttPricePerHour()
return nil
case group.FieldModelPricing:
m.ClearModelPricing()
return nil
case group.FieldFallbackGroupID:
m.ClearFallbackGroupID()
return nil
@@ -26924,6 +27062,12 @@ func (m *GroupMutation) ResetField(name string) error {
case group.FieldAudioSttPricePerHour:
m.ResetAudioSttPricePerHour()
return nil
case group.FieldLongContextPricingEnabled:
m.ResetLongContextPricingEnabled()
return nil
case group.FieldModelPricing:
m.ResetModelPricing()
return nil
case group.FieldClaudeCodeOnly:
m.ResetClaudeCodeOnly()
return nil
+22 -18
View File
@@ -1133,80 +1133,84 @@ func init() {
groupDescAudioSttPricePerHour := groupFields[35].Descriptor()
// group.AudioSttPricePerHourValidator is a validator for the "audio_stt_price_per_hour" field. It is called by the builders before save.
group.AudioSttPricePerHourValidator = groupDescAudioSttPricePerHour.Validators[0].(func(float64) error)
// groupDescLongContextPricingEnabled is the schema descriptor for long_context_pricing_enabled field.
groupDescLongContextPricingEnabled := groupFields[36].Descriptor()
// group.DefaultLongContextPricingEnabled holds the default value on creation for the long_context_pricing_enabled field.
group.DefaultLongContextPricingEnabled = groupDescLongContextPricingEnabled.Default.(bool)
// groupDescClaudeCodeOnly is the schema descriptor for claude_code_only field.
groupDescClaudeCodeOnly := groupFields[36].Descriptor()
groupDescClaudeCodeOnly := groupFields[38].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[40].Descriptor()
groupDescModelRoutingEnabled := groupFields[42].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[41].Descriptor()
groupDescMcpXMLInject := groupFields[43].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[42].Descriptor()
groupDescSupportedModelScopes := groupFields[44].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[43].Descriptor()
groupDescSortOrder := groupFields[45].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[44].Descriptor()
groupDescAllowMessagesDispatch := groupFields[46].Descriptor()
// group.DefaultAllowMessagesDispatch holds the default value on creation for the allow_messages_dispatch field.
group.DefaultAllowMessagesDispatch = groupDescAllowMessagesDispatch.Default.(bool)
// groupDescAllowLive is the schema descriptor for allow_live field.
groupDescAllowLive := groupFields[45].Descriptor()
groupDescAllowLive := groupFields[47].Descriptor()
// group.DefaultAllowLive holds the default value on creation for the allow_live field.
group.DefaultAllowLive = groupDescAllowLive.Default.(bool)
// groupDescRequireOauthOnly is the schema descriptor for require_oauth_only field.
groupDescRequireOauthOnly := groupFields[46].Descriptor()
groupDescRequireOauthOnly := groupFields[48].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[47].Descriptor()
groupDescRequirePrivacySet := groupFields[49].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[48].Descriptor()
groupDescDefaultMappedModel := groupFields[50].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[49].Descriptor()
groupDescMessagesDispatchModelConfig := groupFields[51].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[50].Descriptor()
groupDescModelsListConfig := groupFields[52].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[51].Descriptor()
groupDescRpmLimit := groupFields[53].Descriptor()
// group.DefaultRpmLimit holds the default value on creation for the rpm_limit field.
group.DefaultRpmLimit = groupDescRpmLimit.Default.(int)
// groupDescMaxReasoningEffort is the schema descriptor for max_reasoning_effort field.
groupDescMaxReasoningEffort := groupFields[52].Descriptor()
groupDescMaxReasoningEffort := groupFields[54].Descriptor()
// group.DefaultMaxReasoningEffort holds the default value on creation for the max_reasoning_effort field.
group.DefaultMaxReasoningEffort = groupDescMaxReasoningEffort.Default.(string)
// group.MaxReasoningEffortValidator is a validator for the "max_reasoning_effort" field. It is called by the builders before save.
group.MaxReasoningEffortValidator = groupDescMaxReasoningEffort.Validators[0].(func(string) error)
// groupDescReasoningEffortMappings is the schema descriptor for reasoning_effort_mappings field.
groupDescReasoningEffortMappings := groupFields[53].Descriptor()
groupDescReasoningEffortMappings := groupFields[55].Descriptor()
// group.DefaultReasoningEffortMappings holds the default value on creation for the reasoning_effort_mappings field.
group.DefaultReasoningEffortMappings = groupDescReasoningEffortMappings.Default.([]domain.ReasoningEffortMapping)
// groupDescProfitControlEnabled is the schema descriptor for profit_control_enabled field.
groupDescProfitControlEnabled := groupFields[54].Descriptor()
groupDescProfitControlEnabled := groupFields[56].Descriptor()
// group.DefaultProfitControlEnabled holds the default value on creation for the profit_control_enabled field.
group.DefaultProfitControlEnabled = groupDescProfitControlEnabled.Default.(bool)
// groupDescProfitMinMargin is the schema descriptor for profit_min_margin field.
groupDescProfitMinMargin := groupFields[55].Descriptor()
groupDescProfitMinMargin := groupFields[57].Descriptor()
// group.DefaultProfitMinMargin holds the default value on creation for the profit_min_margin field.
group.DefaultProfitMinMargin = groupDescProfitMinMargin.Default.(float64)
// groupDescProfitSafetyBuffer is the schema descriptor for profit_safety_buffer field.
groupDescProfitSafetyBuffer := groupFields[56].Descriptor()
groupDescProfitSafetyBuffer := groupFields[58].Descriptor()
// group.DefaultProfitSafetyBuffer holds the default value on creation for the profit_safety_buffer field.
group.DefaultProfitSafetyBuffer = groupDescProfitSafetyBuffer.Default.(float64)
idempotencyrecordMixin := schema.IdempotencyRecord{}.Mixin()
+9
View File
@@ -1,6 +1,8 @@
package schema
import (
"encoding/json"
"github.com/Wei-Shaw/sub2api/ent/schema/mixins"
"github.com/Wei-Shaw/sub2api/internal/domain"
@@ -185,6 +187,13 @@ func (Group) Fields() []ent.Field {
Min(0).
SchemaType(map[string]string{dialect.Postgres: "decimal(20,8)"}).
Comment("STT 每小时价格(USD)"),
field.Bool("long_context_pricing_enabled").
Default(true).
Comment("是否按上下文长度应用模型阶梯价格;默认开启以保持官方/渠道长上下文价"),
field.JSON("model_pricing", json.RawMessage{}).
Optional().
SchemaType(map[string]string{dialect.Postgres: "jsonb"}).
Comment("分组逐模型定价;优先级高于渠道和内置定价"),
// Claude Code 客户端限制 (added by migration 029)
field.Bool("claude_code_only").
+2 -1
View File
@@ -41,7 +41,8 @@ func (UserPlatformQuota) Fields() []ent.Field {
// 注意:平台列表的单一权威源为 service.AllowedQuotaPlatforms;
// 此处为 ent 构建期约束,需与 service.AllowedQuotaPlatforms 保持同步。
switch s {
case "anthropic", "openai", "gemini", "antigravity", "grok":
case "anthropic", "openai", "gemini", "antigravity", "grok",
"kimi", "zhipu", "deepseek":
return nil
default:
return fmt.Errorf("platform %q is not allowed", s)
+1 -1
View File
@@ -1,6 +1,6 @@
module github.com/Wei-Shaw/sub2api
go 1.26.5
go 1.26.6
require (
entgo.io/ent v0.14.5
+24
View File
@@ -1025,6 +1025,10 @@ type GatewayConfig struct {
// Grok: Grok/xAI gateway scheduling and free-tier soft-gate settings.
Grok GatewayGrokConfig `mapstructure:"grok"`
// CNProviders: 国产 OpenAI 兼容供应商(kimi/zhipu/deepseek)的余额检测配置。
// 仅作用于 payg(按量付费)账号:周期探测余额,低于阈值则临时停调。
CNProviders GatewayCNProvidersConfig `mapstructure:"cn_providers"`
}
// GatewayGrokConfig holds Grok-specific gateway scheduling knobs.
@@ -1057,6 +1061,18 @@ type GatewayGrokConfig struct {
FreeQuotaStatsCacheSeconds int `mapstructure:"free_quota_stats_cache_seconds"`
}
// GatewayCNProvidersConfig 国产 OpenAI 兼容供应商(kimi/zhipu/deepseek)的余额检测配置。
//
// 仅作用于 payg(按量付费)账号(kimi/deepseek 有公开余额端点;zhipu 无,仅靠响应式 429/402)。
// - balance_check_enabled: 是否启用周期余额检测(默认 true)
// - balance_threshold: 余额低于此值(账户货币单位,默认 0.5)触发临时停调
// - balance_check_interval_minutes: 余额检测周期(分钟,默认 10)
type GatewayCNProvidersConfig struct {
BalanceCheckEnabled bool `mapstructure:"balance_check_enabled"`
BalanceThreshold float64 `mapstructure:"balance_threshold"`
BalanceCheckIntervalMinutes int `mapstructure:"balance_check_interval_minutes"`
}
type GatewayLiveConfig struct {
// MaxSessionDurationSeconds 是 Live 会话的硬上限。
MaxSessionDurationSeconds int `mapstructure:"max_session_duration_seconds"`
@@ -1703,6 +1719,10 @@ func load(allowMissingJWTSecret bool) (*Config, error) {
// 环境变量支持
viper.AutomaticEnv()
viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
if tz, ok := os.LookupEnv("TZ"); ok && strings.TrimSpace(tz) != "" {
// AutomaticEnv 会先把 timezone 映射到 TIMEZONE;显式 Set 保证标准 TZ 变量优先。
viper.Set("timezone", strings.TrimSpace(tz))
}
if err := viper.BindEnv("server.enable_server_timing", "ENABLE_SERVER_TIMING"); err != nil {
return nil, fmt.Errorf("bind ENABLE_SERVER_TIMING: %w", err)
}
@@ -2351,6 +2371,10 @@ func setDefaults() {
viper.SetDefault("gateway.grok.free_quota_soft_gate_percent", 95)
viper.SetDefault("gateway.grok.free_quota_window_hours", 24)
viper.SetDefault("gateway.grok.free_quota_stats_cache_seconds", 60)
// 国产供应商余额检测(kimi/deepseek payg;zhipu 无余额端点,仅靠响应式 429/402)。
viper.SetDefault("gateway.cn_providers.balance_check_enabled", true)
viper.SetDefault("gateway.cn_providers.balance_threshold", 0.5)
viper.SetDefault("gateway.cn_providers.balance_check_interval_minutes", 10)
viper.SetDefault("gateway.image_concurrency.enabled", false)
viper.SetDefault("gateway.image_concurrency.max_concurrent_requests", 0)
viper.SetDefault("gateway.image_concurrency.overflow_mode", ImageConcurrencyOverflowModeReject)
+32
View File
@@ -23,6 +23,38 @@ func resetViperWithJWTSecret(t *testing.T) {
t.Setenv("JWT_SECRET", strings.Repeat("x", 32))
}
func TestLoadTimezonePrecedence(t *testing.T) {
tests := []struct {
name string
fileTimezone string
timezoneEnv string
tzEnv string
want string
}{
{name: "default", want: "Asia/Shanghai"},
{name: "config_file", fileTimezone: "Europe/London", want: "Europe/London"},
{name: "timezone_env", fileTimezone: "Europe/London", timezoneEnv: "UTC", want: "UTC"},
{name: "tz_env", fileTimezone: "Europe/London", timezoneEnv: "UTC", tzEnv: "America/New_York", want: "America/New_York"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
resetViperWithJWTSecret(t)
t.Setenv("TIMEZONE", tt.timezoneEnv)
t.Setenv("TZ", tt.tzEnv)
if tt.fileTimezone != "" {
configFile := filepath.Join(t.TempDir(), "config.yaml")
require.NoError(t, os.WriteFile(configFile, []byte("timezone: "+tt.fileTimezone+"\n"), 0o600))
t.Setenv("CONFIG_FILE", configFile)
}
cfg, err := Load()
require.NoError(t, err)
require.Equal(t, tt.want, cfg.Timezone)
})
}
}
func TestLoadServerTimingConfig(t *testing.T) {
t.Run("disabled by default", func(t *testing.T) {
resetViperWithJWTSecret(t)
+21 -1
View File
@@ -23,7 +23,27 @@ const (
PlatformGemini = "gemini"
PlatformAntigravity = "antigravity"
PlatformGrok = "grok"
PlatformComposite = "composite"
// 国产 OpenAI 兼容供应商(经 OpenAI 网关转发,按 Chat Completions 协议)。
PlatformKimi = "kimi" // Kimi (月之暗面 / Moonshot)
PlatformZhipu = "zhipu" // 智谱 GLM (bigmodel)
PlatformDeepseek = "deepseek" // DeepSeek
PlatformComposite = "composite"
)
// Account mode constants 区分国产供应商的「按量付费(余额)」与「Coding Plan」两种接入方式。
// 存储于 credentials["account_mode"],决定 base_url 预设与额度监控方式。
const (
AccountModePayG = "payg" // 按量付费:消耗余额,做余额检测冷却
AccountModeCoding = "coding" // Coding Plan:滚动用量窗口冷却(5h / weekly)
)
// API protocol constants 国产供应商的上游 API 协议维度。存储于
// credentials["api_protocol"],与 account_mode 正交:协议决定转发端点与格式,
// 模式决定额度监控方式。同协议请求零转换直通;跨协议组合才走转换链。
const (
APIProtocolChatCompletions = "chat_completions" // OpenAI Chat Completions(默认)
APIProtocolAnthropic = "anthropic" // 原生 Anthropic /v1/messages(适配 Claude Code)
APIProtocolResponses = "responses" // OpenAI Responses(仅 deepseek,适配 Codex)
)
// Account type constants
@@ -1025,7 +1025,8 @@ func (h *AccountHandler) Update(c *gin.Context) {
// 当前请求。探测错误仅记录日志,不向上下文传播:探测失败时标记保持缺失,
// 网关会按"现状即证据"默认走 Responses。
func (h *AccountHandler) scheduleOpenAIResponsesProbe(account *service.Account) {
if account == nil || account.Platform != service.PlatformOpenAI || account.Type != service.AccountTypeAPIKey {
if account == nil || account.Type != service.AccountTypeAPIKey ||
(account.Platform != service.PlatformOpenAI && !service.IsCNProvider(account.Platform)) {
return
}
if h.accountTestService == nil {
@@ -0,0 +1,69 @@
package admin
import (
"strconv"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
)
// CNProviderHandler 暴露国产供应商(kimi/zhipu/deepseek)的额度与余额查询端点。
//
// - GET /admin/cn-providers/accounts/:id/quota Coding Plan 滚动窗口用量(kimi/zhipu)
// - GET /admin/cn-providers/accounts/:id/balance payg 账号余额(kimi/deepseek)
//
// 智谱(zhipu)无余额端点,故同一账号仅 quota 或 balance 其一可用:服务端按账号
// platform + account_mode 校验并返回明确错误(见 CNProvider*Service 的 load*Account)。
type CNProviderHandler struct {
quotaService *service.CNProviderQuotaService
balanceService *service.CNProviderBalanceService
}
func NewCNProviderHandler(
quotaService *service.CNProviderQuotaService,
balanceService *service.CNProviderBalanceService,
) *CNProviderHandler {
return &CNProviderHandler{
quotaService: quotaService,
balanceService: balanceService,
}
}
// QueryQuota 查询 Coding Plan 滚动窗口用量(5h + weekly)。
func (h *CNProviderHandler) QueryQuota(c *gin.Context) {
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil {
response.BadRequest(c, "Invalid account ID")
return
}
if h == nil || h.quotaService == nil {
response.BadRequest(c, "cn provider quota service is not enabled")
return
}
result, err := h.quotaService.QueryUsage(c.Request.Context(), accountID)
if err != nil {
response.ErrorFrom(c, err)
return
}
response.Success(c, result)
}
// QueryBalance 查询 payg 账号余额。
func (h *CNProviderHandler) QueryBalance(c *gin.Context) {
accountID, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil {
response.BadRequest(c, "Invalid account ID")
return
}
if h == nil || h.balanceService == nil {
response.BadRequest(c, "cn provider balance service is not enabled")
return
}
result, err := h.balanceService.QueryBalance(c.Request.Context(), accountID)
if err != nil {
response.ErrorFrom(c, err)
return
}
response.Success(c, result)
}
+31 -25
View File
@@ -8,11 +8,11 @@ import (
"log/slog"
"strconv"
"strings"
"time"
"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/platform/liveattestation"
"github.com/Wei-Shaw/sub2api/internal/service"
@@ -96,15 +96,17 @@ func NewGroupHandler(adminService service.AdminService, dashboardService *servic
// CreateGroupRequest represents create group request
type CreateGroupRequest struct {
Name string `json:"name" binding:"required"`
Description string `json:"description"`
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok composite"`
RateMultiplier float64 `json:"rate_multiplier"`
IsExclusive bool `json:"is_exclusive"`
SubscriptionType string `json:"subscription_type" binding:"omitempty,oneof=standard subscription"`
DailyLimitUSD optionalLimitField `json:"daily_limit_usd"`
WeeklyLimitUSD optionalLimitField `json:"weekly_limit_usd"`
MonthlyLimitUSD optionalLimitField `json:"monthly_limit_usd"`
Name string `json:"name" binding:"required"`
Description string `json:"description"`
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok kimi zhipu deepseek composite"`
RateMultiplier float64 `json:"rate_multiplier"`
IsExclusive bool `json:"is_exclusive"`
SubscriptionType string `json:"subscription_type" binding:"omitempty,oneof=standard subscription"`
DailyLimitUSD optionalLimitField `json:"daily_limit_usd"`
WeeklyLimitUSD optionalLimitField `json:"weekly_limit_usd"`
MonthlyLimitUSD optionalLimitField `json:"monthly_limit_usd"`
LongContextPricingEnabled bool `json:"long_context_pricing_enabled"`
ModelPricing []service.ChannelModelPricing `json:"model_pricing"`
// 图片生成计费配置(antigravity 和 gemini 平台使用,负数表示清除配置)
AllowImageGeneration bool `json:"allow_image_generation"`
AllowBatchImageGeneration bool `json:"allow_batch_image_generation"`
@@ -162,16 +164,18 @@ type CreateGroupRequest struct {
// UpdateGroupRequest represents update group request
type UpdateGroupRequest struct {
Name string `json:"name"`
Description *string `json:"description"`
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok composite"`
RateMultiplier *float64 `json:"rate_multiplier"`
IsExclusive *bool `json:"is_exclusive"`
Status string `json:"status" binding:"omitempty,oneof=active inactive"`
SubscriptionType string `json:"subscription_type" binding:"omitempty,oneof=standard subscription"`
DailyLimitUSD optionalLimitField `json:"daily_limit_usd"`
WeeklyLimitUSD optionalLimitField `json:"weekly_limit_usd"`
MonthlyLimitUSD optionalLimitField `json:"monthly_limit_usd"`
Name string `json:"name"`
Description *string `json:"description"`
Platform string `json:"platform" binding:"omitempty,oneof=anthropic openai gemini antigravity grok kimi zhipu deepseek composite"`
RateMultiplier *float64 `json:"rate_multiplier"`
IsExclusive *bool `json:"is_exclusive"`
Status string `json:"status" binding:"omitempty,oneof=active inactive"`
SubscriptionType string `json:"subscription_type" binding:"omitempty,oneof=standard subscription"`
DailyLimitUSD optionalLimitField `json:"daily_limit_usd"`
WeeklyLimitUSD optionalLimitField `json:"weekly_limit_usd"`
MonthlyLimitUSD optionalLimitField `json:"monthly_limit_usd"`
LongContextPricingEnabled *bool `json:"long_context_pricing_enabled"`
ModelPricing *[]service.ChannelModelPricing `json:"model_pricing"`
// 图片生成计费配置(antigravity 和 gemini 平台使用,负数表示清除配置)
AllowImageGeneration *bool `json:"allow_image_generation"`
AllowBatchImageGeneration *bool `json:"allow_batch_image_generation"`
@@ -508,6 +512,8 @@ func (h *GroupHandler) Create(c *gin.Context) {
DailyLimitUSD: req.DailyLimitUSD.ToServiceInput(),
WeeklyLimitUSD: req.WeeklyLimitUSD.ToServiceInput(),
MonthlyLimitUSD: req.MonthlyLimitUSD.ToServiceInput(),
LongContextPricingEnabled: req.LongContextPricingEnabled,
ModelPricing: req.ModelPricing,
AllowImageGeneration: req.AllowImageGeneration,
AllowBatchImageGeneration: req.AllowBatchImageGeneration,
ImageRateIndependent: req.ImageRateIndependent,
@@ -635,6 +641,8 @@ func (h *GroupHandler) Update(c *gin.Context) {
DailyLimitUSD: req.DailyLimitUSD.ToServiceInput(),
WeeklyLimitUSD: req.WeeklyLimitUSD.ToServiceInput(),
MonthlyLimitUSD: req.MonthlyLimitUSD.ToServiceInput(),
LongContextPricingEnabled: req.LongContextPricingEnabled,
ModelPricing: req.ModelPricing,
AllowImageGeneration: req.AllowImageGeneration,
AllowBatchImageGeneration: req.AllowBatchImageGeneration,
ImageRateIndependent: req.ImageRateIndependent,
@@ -726,12 +734,10 @@ func (h *GroupHandler) GetStats(c *gin.Context) {
_ = groupID // TODO: implement actual stats
}
// GetUsageSummary returns today's and cumulative cost for all groups.
// GET /api/v1/admin/groups/usage-summary?timezone=Asia/Shanghai
// GetUsageSummary returns today's, yesterday's, and cumulative cost for all groups.
// GET /api/v1/admin/groups/usage-summary
func (h *GroupHandler) GetUsageSummary(c *gin.Context) {
userTZ := c.Query("timezone")
now := timezone.NowInUserLocation(userTZ)
todayStart := timezone.StartOfDayInUserLocation(now, userTZ)
todayStart := service.GroupUsageTodayStart(time.Now())
results, err := h.dashboardService.GetGroupUsageSummary(c.Request.Context(), todayStart)
if err != nil {
@@ -0,0 +1,83 @@
//go:build unit
package admin
import (
"bytes"
"fmt"
"net/http/httptest"
"testing"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
// 回归分组平台枚举:kimi/zhipu/deepseek 必须能通过 Create/Update 的 binding 校验
// (历史 bug:调度/路由链路已支持 CN 平台分组,但 oneof 白名单漏加三平台,导致
// 平台分组无法创建、CN 账号"无可用分组");非法值仍须被拒。
func bindGroupPlatformJSON(t *testing.T, target any, body string) error {
t.Helper()
gin.SetMode(gin.TestMode)
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest("POST", "/", bytes.NewBufferString(body))
c.Request.Header.Set("Content-Type", "application/json")
return c.ShouldBindJSON(target)
}
func TestGroupPlatformBinding_AllowedPlatforms(t *testing.T) {
allowed := []string{
"anthropic", "openai", "gemini", "antigravity", "grok",
"kimi", "zhipu", "deepseek", "composite",
}
for _, platform := range allowed {
t.Run("create_"+platform, func(t *testing.T) {
var req CreateGroupRequest
body := fmt.Sprintf(`{"name":"g","platform":%q}`, platform)
require.NoError(t, bindGroupPlatformJSON(t, &req, body),
"platform %q 应通过 CreateGroupRequest 校验", platform)
require.Equal(t, platform, req.Platform)
})
t.Run("update_"+platform, func(t *testing.T) {
var req UpdateGroupRequest
body := fmt.Sprintf(`{"platform":%q}`, platform)
require.NoError(t, bindGroupPlatformJSON(t, &req, body),
"platform %q 应通过 UpdateGroupRequest 校验", platform)
require.Equal(t, platform, req.Platform)
})
}
}
func TestGroupPlatformBinding_RejectsInvalidPlatforms(t *testing.T) {
invalid := []string{
"moonshot", // 厂商别名,不是平台标识
"Kimi", // 大小写敏感
"openai ", // 尾随空格
"glm",
"bogus",
}
for _, platform := range invalid {
t.Run("create_"+platform, func(t *testing.T) {
var req CreateGroupRequest
body := fmt.Sprintf(`{"name":"g","platform":%q}`, platform)
require.Error(t, bindGroupPlatformJSON(t, &req, body),
"platform %q 应被 CreateGroupRequest 拒绝", platform)
})
t.Run("update_"+platform, func(t *testing.T) {
var req UpdateGroupRequest
body := fmt.Sprintf(`{"platform":%q}`, platform)
require.Error(t, bindGroupPlatformJSON(t, &req, body),
"platform %q 应被 UpdateGroupRequest 拒绝", platform)
})
}
}
// 守住 composite 路由目标不放行 CN:CN 平台不可作为 composite 路由目标
// (DetectModelPlatform/isConcreteRequestPlatform 均无 CN 分支,放行即打开半实现路径)。
func TestCompositeRouteTargetPlatform_StillExcludesCNProviders(t *testing.T) {
for _, platform := range []string{"kimi", "zhipu", "deepseek"} {
var req CompositeRouteRequest
body := fmt.Sprintf(`{"public_model":"m","target_platform":%q}`, platform)
require.Error(t, bindGroupPlatformJSON(t, &req, body),
"composite target_platform %q 应保持被拒", platform)
}
}
@@ -112,12 +112,13 @@ func TestUpdateUserPlatformQuotas_Success(t *testing.T) {
if len(repo.upsertCalls) != 1 {
t.Fatalf("UpsertForUser should be called once, got %d", len(repo.upsertCalls))
}
if repo.upsertCalls[0].userID != 42 || len(repo.upsertCalls[0].records) != len(service.AllowedQuotaPlatforms) {
// upsert 记录数 = 请求体中给出的平台数(未给出的平台不落库)。
if repo.upsertCalls[0].userID != 42 || len(repo.upsertCalls[0].records) != 5 {
t.Errorf("unexpected upsert call: %+v", repo.upsertCalls[0])
}
// 缓存失效:按全部允许平台统一失效。
if len(cache.deleteCalls) != 5 {
t.Errorf("expected 5 cache delete calls, got %d: %+v", len(cache.deleteCalls), cache.deleteCalls)
// 缓存失效:按全部允许平台统一失效(含 kimi/zhipu/deepseek)。
if len(cache.deleteCalls) != len(service.AllowedQuotaPlatforms) {
t.Errorf("expected %d cache delete calls, got %d: %+v", len(service.AllowedQuotaPlatforms), len(cache.deleteCalls), cache.deleteCalls)
}
}
+2
View File
@@ -149,6 +149,7 @@ func GroupFromServiceAdmin(g *service.Group) *AdminGroup {
ProfitControlEnabled: g.ProfitControlEnabled,
ProfitMinMargin: g.ProfitMinMargin,
ProfitSafetyBuffer: g.ProfitSafetyBuffer,
ModelPricing: g.ModelPricing,
ModelRouting: g.ModelRouting,
ModelRoutingEnabled: g.ModelRoutingEnabled,
MCPXMLInject: g.MCPXMLInject,
@@ -184,6 +185,7 @@ func groupFromServiceBase(g *service.Group) Group {
DailyLimitUSD: g.DailyLimitUSD,
WeeklyLimitUSD: g.WeeklyLimitUSD,
MonthlyLimitUSD: g.MonthlyLimitUSD,
LongContextPricingEnabled: g.LongContextPricingEnabled,
AllowImageGeneration: g.AllowImageGeneration,
AllowBatchImageGeneration: g.AllowBatchImageGeneration,
ImageRateIndependent: g.ImageRateIndependent,
+9 -7
View File
@@ -96,10 +96,11 @@ type Group struct {
IsExclusive bool `json:"is_exclusive"`
Status string `json:"status"`
SubscriptionType string `json:"subscription_type"`
DailyLimitUSD *float64 `json:"daily_limit_usd"`
WeeklyLimitUSD *float64 `json:"weekly_limit_usd"`
MonthlyLimitUSD *float64 `json:"monthly_limit_usd"`
SubscriptionType string `json:"subscription_type"`
DailyLimitUSD *float64 `json:"daily_limit_usd"`
WeeklyLimitUSD *float64 `json:"weekly_limit_usd"`
MonthlyLimitUSD *float64 `json:"monthly_limit_usd"`
LongContextPricingEnabled bool `json:"long_context_pricing_enabled"`
// 图片生成计费配置(仅 antigravity 平台使用)
AllowImageGeneration bool `json:"allow_image_generation"`
@@ -164,9 +165,10 @@ type AdminGroup struct {
// 分组利润控制(五个 token 平台分组可启用;margin/buffer 为小数存储)。
// 仅管理员可见:这三个字段与同响应中的 rate_multiplier 相乘即可反推出
// 运营方的上游成本上限,属于内部经营信息,不得下放到 dto.Group。
ProfitControlEnabled bool `json:"profit_control_enabled"`
ProfitMinMargin float64 `json:"profit_min_margin"`
ProfitSafetyBuffer float64 `json:"profit_safety_buffer"`
ProfitControlEnabled bool `json:"profit_control_enabled"`
ProfitMinMargin float64 `json:"profit_min_margin"`
ProfitSafetyBuffer float64 `json:"profit_safety_buffer"`
ModelPricing []service.ChannelModelPricing `json:"model_pricing"`
// 模型路由配置(仅 anthropic 平台使用)
ModelRouting map[string][]int64 `json:"model_routing"`
+1 -1
View File
@@ -1249,7 +1249,7 @@ func writeGrokModelsList(c *gin.Context, modelIDs []string) {
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":
case "grok-4.6", "grok-4.6-latest", "grok-4.5", "grok-4.5-latest", "grok", "grok-latest", "grok-build", "grok-build-latest", "grok-build-0.1":
return true
default:
return false
+56 -17
View File
@@ -28,12 +28,8 @@ const (
)
func (h *GatewayHandler) WebSearch(c *gin.Context) {
type webSearchReq struct {
Query string `json:"query" binding:"required"`
MaxResults int `json:"max_results"`
}
var req webSearchReq
isXSearch := c.GetBool("grok_x_search_endpoint")
var req grokStandaloneSearchRequest
if err := c.ShouldBindJSON(&req); err != nil {
c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{
"type": "invalid_request_error",
@@ -41,7 +37,28 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
}})
return
}
req.MaxResults = normalizeGrokWebSearchMaxResults(req.MaxResults)
query := strings.TrimSpace(req.Query)
if query == "" {
query = strings.TrimSpace(req.Input)
}
if query == "" {
c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{
"type": "invalid_request_error",
"message": "query is required",
}})
return
}
req.Query = query
maxResults := 0
if req.MaxResults != nil {
maxResults = *req.MaxResults
}
maxResults = normalizeGrokWebSearchMaxResults(maxResults)
searchModel := resolveGrokStandaloneSearchModel()
searchLabel := "web_search"
if isXSearch {
searchLabel = "x_search"
}
apiKey, ok := middleware2.GetAPIKeyFromContext(c)
if !ok || apiKey == nil {
@@ -55,7 +72,7 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
if apiKey.Group == nil || apiKey.Group.Platform != "grok" {
c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{
"type": "invalid_request_error",
"message": "web search is only supported for grok groups",
"message": searchLabel + " is only supported for grok groups",
}})
return
}
@@ -79,7 +96,7 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
"role": "user", "content": req.Query,
}},
})
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIChat, xai.DefaultTextModel, auditBody); decision != nil && !decision.AllowNextStage {
if decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIChat, searchModel, auditBody); decision != nil && !decision.AllowNextStage {
status := decision.HTTPStatus
if status == 0 {
status = http.StatusForbidden
@@ -123,7 +140,7 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
// First attempt + up to 3 failover accounts (max 4 total).
for attempt := 0; attempt < 4; attempt++ {
selected, selectErr := h.gatewayService.SelectAccountWithLoadAwareness(
c.Request.Context(), groupID, "", xai.DefaultTextModel, failedAccounts, "", 0,
c.Request.Context(), groupID, "", searchModel, failedAccounts, "", 0,
)
if selectErr != nil {
if attempt == 0 {
@@ -159,7 +176,11 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
account = selected.Account
accountReleaseFunc = release
nativeResp, providerName, err = h.doGrokNativeWebSearch(c.Request.Context(), c, account, req.Query, req.MaxResults)
if isXSearch {
nativeResp, providerName, err = h.doGrokNativeXSearch(c.Request.Context(), c, account, req, searchModel, maxResults)
} else {
nativeResp, providerName, err = h.doGrokNativeWebSearch(c.Request.Context(), c, account, req.Query, maxResults, searchModel)
}
if err == nil {
break
}
@@ -198,7 +219,7 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
// Request IDs are billing idempotency keys, so they must be unique per invocation.
// Query/IP/UA hashes would collapse repeated identical searches into one charge.
searchRequestID := "web_search:" + uuid.NewString()
searchRequestID := searchLabel + ":" + uuid.NewString()
if apiKey.Group != nil {
if p := apiKey.Group.GetSearchPricePer1k(); p != nil && *p == 0 {
logger.L().With(
@@ -211,7 +232,7 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.RecordUsageInput{
Result: &service.ForwardResult{
RequestID: searchRequestID,
Model: "grok-web-search",
Model: "grok-" + strings.ReplaceAll(searchLabel, "_", "-"),
SearchCount: 1,
Duration: 0,
},
@@ -240,7 +261,7 @@ func (h *GatewayHandler) WebSearch(c *gin.Context) {
"query": req.Query,
"results": nativeResp.Results,
"provider": providerName,
"max_results": req.MaxResults,
"max_results": maxResults,
})
}
@@ -299,13 +320,13 @@ func (h *GatewayHandler) acquireWebSearchAccountSlot(
// doGrokNativeWebSearch executes web search using the Grok account's native capability
// by calling the responses endpoint with web_search tool, then normalizes sources to unified format.
func (h *GatewayHandler) doGrokNativeWebSearch(ctx context.Context, c *gin.Context, account *service.Account, query string, maxResults int) (*websearch.SearchResponse, string, error) {
func (h *GatewayHandler) doGrokNativeWebSearch(ctx context.Context, c *gin.Context, account *service.Account, query string, maxResults int, model string) (*websearch.SearchResponse, string, error) {
maxResults = normalizeGrokWebSearchMaxResults(maxResults)
// Build a minimal responses request that triggers Grok web search tool.
// Ask for structured metadata because xAI action.sources commonly contains URLs only.
searchBody := map[string]any{
"model": xai.DefaultTextModel,
"model": xai.ResolveDefaultTextModel(model),
"input": buildGrokWebSearchPrompt(query, maxResults),
"tools": []map[string]any{{"type": "web_search"}},
"include": []string{"web_search_call.action.sources"},
@@ -329,6 +350,23 @@ func (h *GatewayHandler) doGrokNativeWebSearch(ctx context.Context, c *gin.Conte
}, "grok-native", nil
}
func (h *GatewayHandler) doGrokNativeXSearch(ctx context.Context, c *gin.Context, account *service.Account, req grokStandaloneSearchRequest, model string, maxResults int) (*websearch.SearchResponse, string, error) {
maxResults = normalizeGrokWebSearchMaxResults(maxResults)
bodyBytes, err := buildGrokXSearchResponsesBody(req, model)
if err != nil {
return nil, "", err
}
respBytes, err := h.gatewayService.DoGrokNativeResponsesJSON(ctx, account, bodyBytes)
if err != nil {
return nil, "", err
}
results := extractGrokWebSearchSources(respBytes, maxResults)
return &websearch.SearchResponse{
Results: results,
Query: req.Query,
}, "grok-native", nil
}
func normalizeGrokWebSearchMaxResults(maxResults int) int {
if maxResults <= 0 {
return defaultGrokWebSearchResults
@@ -377,7 +415,8 @@ func extractGrokWebSearchSources(body []byte, maxResults int) []websearch.Search
output := gjson.GetBytes(body, "output")
output.ForEach(func(_, item gjson.Result) bool {
if item.Get("type").String() == "web_search_call" {
callType := item.Get("type").String()
if callType == "web_search_call" || callType == "x_search_call" {
sources := item.Get("action.sources")
if sources.IsArray() {
sources.ForEach(func(_, src gjson.Result) bool {
+14 -11
View File
@@ -89,7 +89,7 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) {
model = "grok-voice-latest"
}
started := time.Now()
proxyErr := h.gatewayService.ProxyGrokRealtime(c.Request.Context(), c, conn, selection.Account, token, model)
audioObserved, proxyErr := h.gatewayService.ProxyGrokRealtime(c.Request.Context(), c, conn, selection.Account, token, model)
elapsed := time.Since(started)
if proxyErr != nil {
reqLog.Info("grok_realtime.proxy_failed", zap.Error(proxyErr))
@@ -98,20 +98,23 @@ func (h *OpenAIGatewayHandler) GrokRealtime(c *gin.Context) {
return
}
}
// A relay normally returns a close error when either side closes normally.
// Those sessions still consumed upstream audio time and must be billed.
if elapsed > 0 {
result := &service.OpenAIForwardResult{
// One durable id per WS session so retries cannot collapse or double under client ids.
RequestID: service.StableGrokRealtimeBillingRequestID(""),
Model: model,
Duration: elapsed,
AudioUsage: &service.AudioUsage{Mode: "realtime", DurationOrUnits: elapsed.Minutes()},
}
if result := grokRealtimeBillingResult(model, elapsed, audioObserved); result != nil {
h.recordGrokVoiceUsage(c, apiKey, selection.Account, subscription, "realtime", nil, result)
}
}
func grokRealtimeBillingResult(model string, elapsed time.Duration, audioObserved bool) *service.OpenAIForwardResult {
if !audioObserved || elapsed <= 0 {
return nil
}
return &service.OpenAIForwardResult{
RequestID: service.StableGrokRealtimeBillingRequestID(""),
Model: model,
Duration: elapsed,
AudioUsage: &service.AudioUsage{Mode: "realtime", DurationOrUnits: elapsed.Minutes()},
}
}
func isExpectedGrokRealtimeClose(err error) bool {
if err == nil {
return true
@@ -4,6 +4,7 @@ package handler
import (
"testing"
"time"
coderws "github.com/coder/websocket"
)
@@ -23,3 +24,29 @@ func TestIsExpectedGrokRealtimeClose(t *testing.T) {
t.Fatal("policy violations must not be treated as billable normal closes")
}
}
func TestGrokRealtimeBillingResultRequiresObservedAudio(t *testing.T) {
if grokRealtimeBillingResult("grok-voice-latest", time.Second, false) != nil {
t.Fatal("a session without observed audio must not be billed")
}
if grokRealtimeBillingResult("grok-voice-latest", 0, true) != nil {
t.Fatal("zero-duration sessions must not be billed")
}
}
func TestGrokRealtimeBillingResultUsesForcedUniqueID(t *testing.T) {
first := grokRealtimeBillingResult("grok-voice-latest", 90*time.Second, true)
second := grokRealtimeBillingResult("grok-voice-latest", 90*time.Second, true)
if first == nil || second == nil {
t.Fatal("observed audio sessions should be billable")
}
if first.RequestID == "" {
t.Fatalf("unexpected billing request ID %q", first.RequestID)
}
if first.RequestID == second.RequestID {
t.Fatal("independent realtime connections must not share a billing request ID")
}
if first.AudioUsage == nil || first.AudioUsage.Mode != "realtime" || first.AudioUsage.DurationOrUnits != 1.5 {
t.Fatalf("unexpected audio usage: %#v", first.AudioUsage)
}
}
+1
View File
@@ -19,6 +19,7 @@ type AdminHandlers struct {
GeminiOAuth *admin.GeminiOAuthHandler
AntigravityOAuth *admin.AntigravityOAuthHandler
GrokOAuth *admin.GrokOAuthHandler
CNProvider *admin.CNProviderHandler
Proxy *admin.ProxyHandler
Redeem *admin.RedeemHandler
Promo *admin.PromoHandler
@@ -255,6 +255,48 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
if err == nil && result != nil && result.FirstTokenMs != nil {
service.SetOpsLatencyMs(c, service.OpsTimeToFirstTokenMsKey, int64(*result.FirstTokenMs))
}
// #5148 对齐:错误返回携带的部分 result(流中断前上游已计量的 usage)照常
// 入账;failover 错误恒定 result=nil,不会重复计费。
submitChatUsage := func(res *service.OpenAIForwardResult) {
if res == nil {
return
}
userAgent := c.GetHeader("User-Agent")
clientIP := ip.GetClientIP(c)
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, res)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
h.submitOpenAIUsageRecordTask(c.Request.Context(), res, func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
Result: res,
APIKey: apiKey,
User: apiKey.User,
Account: account,
Subscription: subscription,
InboundEndpoint: inboundEndpoint,
UpstreamEndpoint: upstreamEndpoint,
UserAgent: userAgent,
IPAddress: clientIP,
APIKeyService: h.apiKeyService,
QuotaPlatform: quotaPlatform,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, res.UpstreamModel),
PricingAt: pricingAt,
CyberBlocked: cyberBlocked,
}); err != nil {
logger.L().With(
zap.String("component", "handler.openai_gateway.chat_completions"),
zap.Int64("user_id", subject.UserID),
zap.Int64("api_key_id", apiKey.ID),
zap.Any("group_id", apiKey.GroupID),
zap.String("model", reqModel),
zap.Int64("account_id", account.ID),
).Error("openai_chat_completions.record_usage_failed", zap.Error(err))
}
})
}
if err != nil {
if result != nil && result.ImageCount > 0 {
reqLog.Warn("openai_chat_completions.forward_partial_error_with_image_result",
@@ -339,6 +381,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
zap.Bool("upstream_error_response_already_written", upstreamErrorAlreadyCommunicated),
zap.Error(err),
)
submitChatUsage(result)
return
}
}
@@ -348,42 +391,7 @@ func (h *OpenAIGatewayHandler) ChatCompletions(c *gin.Context) {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), true, nil)
}
userAgent := c.GetHeader("User-Agent")
clientIP := ip.GetClientIP(c)
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
Result: result,
APIKey: apiKey,
User: apiKey.User,
Account: account,
Subscription: subscription,
InboundEndpoint: inboundEndpoint,
UpstreamEndpoint: upstreamEndpoint,
UserAgent: userAgent,
IPAddress: clientIP,
APIKeyService: h.apiKeyService,
QuotaPlatform: quotaPlatform,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
PricingAt: pricingAt,
CyberBlocked: cyberBlocked,
}); err != nil {
logger.L().With(
zap.String("component", "handler.openai_gateway.chat_completions"),
zap.Int64("user_id", subject.UserID),
zap.Int64("api_key_id", apiKey.ID),
zap.Any("group_id", apiKey.GroupID),
zap.String("model", reqModel),
zap.Int64("account_id", account.ID),
).Error("openai_chat_completions.record_usage_failed", zap.Error(err))
}
})
submitChatUsage(result)
reqLog.Debug("openai_chat_completions.request_completed",
zap.Int64("account_id", account.ID),
zap.Int("switch_count", switchCount),
@@ -0,0 +1,29 @@
package handler
// CN 分组 /v1/messages 调度闸门回归(修复:正常途径创建的 CN 分组曾恒 403):
// sanitizeGroupMessagesDispatchFields 对非 openai 平台强制 AllowMessagesDispatch
// =false,故 CN 分组必须与 grok 一样在闸门处豁免,否则原生 Anthropic 直通
//(Claude Code 主用例)永远不可达。
import (
"testing"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
func TestAllowOpenAICompatibleMessagesDispatch_CNProvidersExempt(t *testing.T) {
require.True(t, allowOpenAICompatibleMessagesDispatch(nil), "无 key 保持放行")
for _, platform := range []string{service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek, service.PlatformGrok} {
apiKey := &service.APIKey{Group: &service.Group{Platform: platform, AllowMessagesDispatch: false}}
require.True(t, allowOpenAICompatibleMessagesDispatch(apiKey),
"%s 分组必须豁免 allow_messages_dispatch 闸门", platform)
}
// 非回归:openai 分组仍受开关控制。
openaiOff := &service.APIKey{Group: &service.Group{Platform: service.PlatformOpenAI, AllowMessagesDispatch: false}}
require.False(t, allowOpenAICompatibleMessagesDispatch(openaiOff))
openaiOn := &service.APIKey{Group: &service.Group{Platform: service.PlatformOpenAI, AllowMessagesDispatch: true}}
require.True(t, allowOpenAICompatibleMessagesDispatch(openaiOn))
}
@@ -25,54 +25,211 @@ func newCompactBodySignalTestContext(t *testing.T, path string, body []byte) *gi
func TestNormalizeOpenAIResponsesCompactRequest_RemoteV2StaysOnResponses(t *testing.T) {
h := &OpenAIGatewayHandler{}
body := []byte(`{
"model":"gpt-5.6-sol",
"stream":true,
"store":true,
"prompt_cache_key":"pck-signal-1",
"reasoning":{"effort":"max","context":"all_turns"},
"input":[
{"type":"message","role":"user","content":"hello"},
{"type":"compaction_trigger"}
]
}`)
c := newCompactBodySignalTestContext(t, "/v1/responses", body)
c.Request.Header.Set("x-codex-beta-features", "responses_websockets_v2, remote_compaction_v2, another_feature")
tests := []struct {
name string
betaHeader string
userAgent string
}{
{name: "headerless"},
{name: "unrelated_header", betaHeader: "responses_websockets_v2"},
{name: "wrong_case_header", betaHeader: "REMOTE_COMPACTION_V2"},
{name: "declared_header", betaHeader: "remote_compaction_v2"},
{name: "codex_cli_user_agent", userAgent: "codex_cli_rs/0.144.1 (Ubuntu 22.4.0; x86_64) xterm-256color"},
{name: "codex_desktop_user_agent", userAgent: "Codex Desktop/0.139.0 (Mac OS X 14; arm64) unknown"},
}
normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body)
require.True(t, ok)
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
body := []byte(`{
"model":"gpt-5.6-sol",
"stream":true,
"store":true,
"prompt_cache_key":"pck-signal-1",
"reasoning":{"effort":"max","context":"all_turns"},
"input":[
{"type":"message","role":"user","content":"hello"},
{"type":"compaction_trigger"}
]
}`)
c := newCompactBodySignalTestContext(t, "/v1/responses", body)
if tt.betaHeader != "" {
c.Request.Header.Set("x-codex-beta-features", tt.betaHeader)
}
if tt.userAgent != "" {
c.Request.Header.Set("User-Agent", tt.userAgent)
}
require.Equal(t, "/v1/responses", c.Request.URL.Path)
require.False(t, isOpenAIRemoteCompactPath(c))
require.Equal(t, body, normalized)
require.True(t, gjson.GetBytes(normalized, "stream").Bool())
require.True(t, gjson.GetBytes(normalized, "store").Bool())
require.Equal(t, "pck-signal-1", gjson.GetBytes(normalized, "prompt_cache_key").String())
require.Equal(t, "max", gjson.GetBytes(normalized, "reasoning.effort").String())
require.Equal(t, "all_turns", gjson.GetBytes(normalized, "reasoning.context").String())
normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body)
require.True(t, ok)
reqStream, streamOK := parseOpenAICompatibleStream(normalized)
require.True(t, streamOK)
require.True(t, reqStream)
require.Equal(t, "/v1/responses", c.Request.URL.Path)
require.False(t, isOpenAILegacyCompactPath(c))
require.Equal(t, body, normalized)
require.True(t, gjson.GetBytes(normalized, "stream").Bool())
require.True(t, gjson.GetBytes(normalized, "store").Bool())
require.Equal(t, "pck-signal-1", gjson.GetBytes(normalized, "prompt_cache_key").String())
require.Equal(t, "max", gjson.GetBytes(normalized, "reasoning.effort").String())
require.Equal(t, "all_turns", gjson.GetBytes(normalized, "reasoning.context").String())
legacyCompact := service.IsOpenAIResponsesCompactPath(c)
nativeV2 := isBareOpenAIResponsesPath(c) && isOpenAIRemoteCompactionV2Request(normalized)
require.False(t, legacyCompact)
require.True(t, nativeV2)
require.Equal(t, service.OpenAIEndpointCapabilityResponses,
openAIResponsesRequiredCapabilityForRequest(false, nativeV2 || legacyCompact, service.PlatformOpenAI))
_, seedExists := c.Get(service.OpenAICompactSessionSeedKeyForTest())
require.False(t, seedExists)
_, streamMarkerExists := c.Get(service.OpenAICompactClientStreamKeyForTest())
require.False(t, streamMarkerExists)
reqStream, streamOK := parseOpenAICompatibleStream(normalized)
require.True(t, streamOK)
require.True(t, reqStream)
_, seedExists := c.Get(service.OpenAICompactSessionSeedKeyForTest())
require.False(t, seedExists)
_, streamMarkerExists := c.Get(service.OpenAICompactClientStreamKeyForTest())
require.False(t, streamMarkerExists)
})
}
}
func TestNormalizeOpenAIResponsesCompactRequest_RemoteV2PathAliasesStayOnResponses(t *testing.T) {
h := &OpenAIGatewayHandler{}
body := []byte(`{"model":"gpt-5.6-sol","stream":true,"input":[{"type":"compaction_trigger"}]}`)
for _, path := range []string{"/v1/responses/", "/backend-api/codex/responses"} {
for _, path := range []string{
"/v1/responses/",
"/openai/v1/responses",
"/responses",
"/backend-api/codex/responses",
} {
t.Run(path, func(t *testing.T) {
c := newCompactBodySignalTestContext(t, path, body)
c.Request.Header.Set("x-codex-beta-features", "remote_compaction_v2")
normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body)
require.True(t, ok)
require.Equal(t, path, c.Request.URL.Path)
require.Equal(t, body, normalized)
legacyCompact := service.IsOpenAIResponsesCompactPath(c)
nativeV2 := isBareOpenAIResponsesPath(c) && isOpenAIRemoteCompactionV2Request(normalized)
require.False(t, legacyCompact)
require.True(t, nativeV2)
require.Equal(t, service.OpenAIEndpointCapabilityResponses,
openAIResponsesRequiredCapabilityForRequest(false, nativeV2 || legacyCompact, service.PlatformOpenAI))
})
}
}
func TestOpenAIResponsesCompactionRoutingFlags(t *testing.T) {
h := &OpenAIGatewayHandler{}
tests := []struct {
name string
body []byte
path string
wantLegacyBefore bool
wantNativeBefore bool
wantLegacyAfter bool
wantNativeAfter bool
wantCapabilityAfter service.OpenAIEndpointCapability
wantPathAfter string
wantBodyUnchanged bool
}{
{
name: "native_v2_stream_trigger",
body: []byte(`{"model":"gpt-5.6-sol","stream":true,"input":[{"type":"compaction_trigger"}]}`),
path: "/v1/responses",
wantLegacyBefore: false,
wantNativeBefore: true,
wantLegacyAfter: false,
wantNativeAfter: true,
wantCapabilityAfter: service.OpenAIEndpointCapabilityResponses,
wantPathAfter: "/v1/responses",
wantBodyUnchanged: true,
},
{
name: "native_v2_without_trigger",
body: []byte(`{"model":"gpt-5.6-sol","stream":true,"input":[{"type":"message","role":"user","content":"hello"}]}`),
path: "/v1/responses",
wantLegacyBefore: false,
wantNativeBefore: false,
wantLegacyAfter: false,
wantNativeAfter: false,
wantCapabilityAfter: service.OpenAIEndpointCapabilityChatCompletions,
wantPathAfter: "/v1/responses",
wantBodyUnchanged: true,
},
{
name: "explicit_compact",
body: []byte(`{"model":"gpt-5.6-sol","stream":true,"input":[{"type":"compaction_trigger"}]}`),
path: "/v1/responses/compact",
wantLegacyBefore: true,
wantNativeBefore: false,
wantLegacyAfter: true,
wantNativeAfter: false,
wantCapabilityAfter: service.OpenAIEndpointCapabilityResponses,
wantPathAfter: "/v1/responses/compact",
},
{
name: "nested_compact",
body: []byte(`{"model":"gpt-5.6-sol","stream":true,"input":[{"type":"compaction_trigger"}]}`),
path: "/v1/responses/compact/detail",
wantLegacyBefore: true,
wantNativeBefore: false,
wantLegacyAfter: true,
wantNativeAfter: false,
wantCapabilityAfter: service.OpenAIEndpointCapabilityResponses,
wantPathAfter: "/v1/responses/compact/detail",
},
{
name: "responses_subpath_with_native_signal",
body: []byte(`{"model":"gpt-5.6-sol","stream":true,"input":[{"type":"compaction_trigger"}]}`),
path: "/v1/responses/resp_123/responses",
wantLegacyBefore: false,
wantNativeBefore: false,
wantLegacyAfter: false,
wantNativeAfter: false,
wantCapabilityAfter: service.OpenAIEndpointCapabilityChatCompletions,
wantPathAfter: "/v1/responses/resp_123/responses",
wantBodyUnchanged: true,
},
{
name: "stream_false_promotes",
body: []byte(`{"model":"gpt-5.6-sol","stream":false,"input":[{"type":"compaction_trigger"}]}`),
path: "/v1/responses",
wantLegacyBefore: false,
wantNativeBefore: false,
wantLegacyAfter: true,
wantNativeAfter: false,
wantCapabilityAfter: service.OpenAIEndpointCapabilityResponses,
wantPathAfter: "/v1/responses/compact",
},
{
name: "stream_absent_promotes",
body: []byte(`{"model":"gpt-5.6-sol","input":[{"type":"compaction_trigger"}]}`),
path: "/v1/responses",
wantLegacyBefore: false,
wantNativeBefore: false,
wantLegacyAfter: true,
wantNativeAfter: false,
wantCapabilityAfter: service.OpenAIEndpointCapabilityResponses,
wantPathAfter: "/v1/responses/compact",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
c := newCompactBodySignalTestContext(t, tt.path, tt.body)
legacyBefore := service.IsOpenAIResponsesCompactPath(c)
nativeBefore := isBareOpenAIResponsesPath(c) && isOpenAIRemoteCompactionV2Request(tt.body)
require.Equal(t, tt.wantLegacyBefore, legacyBefore)
require.Equal(t, tt.wantNativeBefore, nativeBefore)
normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), tt.body)
require.True(t, ok)
require.Equal(t, tt.wantPathAfter, c.Request.URL.Path)
legacyAfter := service.IsOpenAIResponsesCompactPath(c)
nativeAfter := isBareOpenAIResponsesPath(c) && isOpenAIRemoteCompactionV2Request(normalized)
require.Equal(t, tt.wantLegacyAfter, legacyAfter)
require.Equal(t, tt.wantNativeAfter, nativeAfter)
require.Equal(t, tt.wantCapabilityAfter,
openAIResponsesRequiredCapabilityForRequest(false, nativeAfter || legacyAfter, service.PlatformOpenAI))
if tt.wantBodyUnchanged {
require.Equal(t, tt.body, normalized)
}
})
}
}
@@ -106,31 +263,22 @@ func TestNormalizeOpenAIResponsesCompactRequest_NonRemoteV2BodySignalPromoted(t
wantMarked bool
}{
{
name: "no_header",
body: []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"compaction_trigger"}]}`),
wantMarked: true,
name: "stream_false_headerless",
body: []byte(`{"model":"gpt-5.5","stream":false,"input":[{"type":"compaction_trigger"}]}`),
},
{
name: "unrelated_header",
body: []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"compaction_trigger"}]}`),
betaHeader: "responses_websockets_v2",
wantMarked: true,
name: "stream_absent_headerless",
body: []byte(`{"model":"gpt-5.5","input":[{"type":"compaction_trigger"}]}`),
},
{
name: "wrong_case_header",
body: []byte(`{"model":"gpt-5.5","stream":true,"input":[{"type":"compaction_trigger"}]}`),
betaHeader: "REMOTE_COMPACTION_V2",
wantMarked: true,
},
{
name: "stream_false",
name: "stream_false_declared_header",
body: []byte(`{"model":"gpt-5.5","stream":false,"input":[{"type":"compaction_trigger"}]}`),
betaHeader: "remote_compaction_v2",
},
{
name: "stream_absent",
name: "stream_absent_wrong_case_header",
body: []byte(`{"model":"gpt-5.5","input":[{"type":"compaction_trigger"}]}`),
betaHeader: "remote_compaction_v2",
betaHeader: "REMOTE_COMPACTION_V2",
},
}
@@ -163,7 +311,7 @@ func TestNormalizeOpenAIResponsesCompactRequest_NoTriggerUntouched(t *testing.T)
normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body)
require.True(t, ok)
require.Equal(t, "/v1/responses", c.Request.URL.Path)
require.False(t, isOpenAIRemoteCompactPath(c))
require.False(t, isOpenAILegacyCompactPath(c))
require.Equal(t, body, normalized)
require.True(t, gjson.GetBytes(normalized, "stream").Bool())
}
@@ -172,7 +320,6 @@ func TestNormalizeOpenAIResponsesCompactRequest_PathBasedNoDoubleSuffix(t *testi
h := &OpenAIGatewayHandler{}
body := []byte(`{"model":"gpt-5.5","stream":true,"store":true,"input":[{"type":"message","role":"user","content":"hello"}]}`)
c := newCompactBodySignalTestContext(t, "/v1/responses/compact", body)
c.Request.Header.Set("x-codex-beta-features", "remote_compaction_v2")
normalized, ok := h.normalizeOpenAIResponsesCompactRequest(c, zap.NewNop(), body)
require.True(t, ok)
@@ -105,21 +105,29 @@ func captureHandlerStructuredLog(t *testing.T) (*handlerInMemoryLogSink, func())
}
}
func TestIsOpenAIRemoteCompactPath(t *testing.T) {
require.False(t, isOpenAIRemoteCompactPath(nil))
func TestIsOpenAILegacyCompactPath(t *testing.T) {
require.False(t, isOpenAILegacyCompactPath(nil))
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses/compact", nil)
require.True(t, isOpenAIRemoteCompactPath(c))
c.Request = httptest.NewRequest(http.MethodPost, "/responses/compact/", nil)
require.True(t, isOpenAIRemoteCompactPath(c))
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
require.False(t, isOpenAIRemoteCompactPath(c))
for _, test := range []struct {
path string
want bool
}{
{path: "/v1/responses/compact", want: true},
{path: "/v1/responses/compact/detail", want: true},
{path: "/responses/compact/", want: true},
{path: "/v1/responses", want: false},
{path: "/openai/v1/responses", want: false},
{path: "/responses", want: false},
{path: "/backend-api/codex/responses", want: false},
{path: "/v1/responses/resp_123/cancel", want: false},
} {
c.Request = httptest.NewRequest(http.MethodPost, test.path, nil)
require.Equal(t, test.want, isOpenAILegacyCompactPath(c), test.path)
}
}
func TestLogOpenAIRemoteCompactOutcome_Succeeded(t *testing.T) {
@@ -78,7 +78,7 @@ func (h *OpenAIGatewayHandler) CountTokens(c *gin.Context) {
zap.Any("group_id", apiKey.GroupID),
)
if apiKey.Group != nil && !apiKey.Group.AllowMessagesDispatch {
if !allowOpenAICompatibleMessagesDispatch(apiKey) {
h.anthropicErrorResponse(c, http.StatusForbidden, "permission_error",
"This group does not allow /v1/messages dispatch")
return
+155 -111
View File
@@ -164,13 +164,11 @@ func wrapUsageRecordTaskContext(parent context.Context, task service.UsageRecord
func openAICompatibleRequestPlatform(ctx context.Context, apiKey *service.APIKey) string {
if platform, ok := service.ResolvedTargetPlatformFromContext(ctx); ok {
if platform == service.PlatformGrok {
return service.PlatformGrok
}
return service.PlatformOpenAI
// 保留 grok 与国产供应商原值,其他归一为 openai(与调度器精确匹配语义一致)。
return service.NormalizeOpenAICompatiblePlatform(platform)
}
if apiKey != nil && apiKey.Group != nil && apiKey.Group.Platform == service.PlatformGrok {
return service.PlatformGrok
if apiKey != nil && apiKey.Group != nil {
return service.NormalizeOpenAICompatiblePlatform(apiKey.Group.Platform)
}
return service.PlatformOpenAI
}
@@ -182,6 +180,16 @@ func openAIResponsesRequiredCapability(imageIntent bool, platform string) servic
return service.OpenAIEndpointCapabilityChatCompletions
}
// openAIResponsesRequiredCapabilityForRequest returns the endpoint capability
// required by an image or Responses request. needsResponses includes both the
// legacy /responses/compact endpoint and native remote compaction v2.
func openAIResponsesRequiredCapabilityForRequest(imageIntent bool, needsResponses bool, platform string) service.OpenAIEndpointCapability {
if needsResponses && platform == service.PlatformOpenAI {
return service.OpenAIEndpointCapabilityResponses
}
return openAIResponsesRequiredCapability(imageIntent, platform)
}
func allowOpenAICompatibleMessagesDispatch(apiKey *service.APIKey) bool {
if apiKey == nil || apiKey.Group == nil {
return true
@@ -189,11 +197,20 @@ func allowOpenAICompatibleMessagesDispatch(apiKey *service.APIKey) bool {
if apiKey.Group.Platform == service.PlatformGrok {
return true
}
// 国产供应商分组与 grok 同语义:/v1/messages 就是其主要服务形态(anthropic
// 协议账号原生直通 Claude Code),无需 allow_messages_dispatch 开关授权——
// 该开关对非 openai 平台恒被 sanitizeGroupMessagesDispatchFields 置 false,
// 若不豁免,CN 分组将永远 403。
if service.IsCNProvider(apiKey.Group.Platform) {
return true
}
return apiKey.Group.AllowMessagesDispatch
}
func openAICompatibleTextTargetAllowed(c *gin.Context, apiKey *service.APIKey, model string) bool {
return compositeTargetPlatformAllowed(c, apiKey, model, service.PlatformOpenAI, service.PlatformGrok)
return compositeTargetPlatformAllowed(c, apiKey, model,
service.PlatformOpenAI, service.PlatformGrok,
service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek)
}
// NewOpenAIGatewayHandler creates a new OpenAIGatewayHandler
@@ -288,6 +305,13 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
if !ok {
return
}
legacyCompact := service.IsOpenAIResponsesCompactPath(c)
nativeV2 := isBareOpenAIResponsesPath(c) && isOpenAIRemoteCompactionV2Request(body)
if nativeV2 {
// 原生 v2 压缩出站前补注 x-codex-beta-features: remote_compaction_v2,
// 与真实 Codex 线型一致(网关链剥头后本级负责恢复,#5586)。
service.MarkOpenAINativeCompactionV2(c)
}
// body-signal compact:上游 unary 等待期间向下游发 SSE 注释行心跳,防止
// 反向代理空闲超时掐断长压缩连接(#3887)。首拍延迟一个心跳间隔,快速
// 失败仍走 JSON+状态码链路;未标记客户端流式或间隔为 0 时是 no-op。
@@ -377,6 +401,15 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
channelMapping, _ := h.gatewayService.ResolveChannelMappingAndRestrict(c.Request.Context(), apiKey.GroupID, reqModel)
forwardBody := openAIModelMappedBody(body, channelMapping.Mapped, channelMapping.MappedModel, h.gatewayService.ReplaceModelInBody)
seedOpenAIForwardImageIntentHint(c, channelMapping.Mapped, imageIntent)
forwardModel := reqModel
if channelMapping.Mapped {
forwardModel = channelMapping.MappedModel
}
c.Request = c.Request.WithContext(service.WithOpenAIForwardModel(
c.Request.Context(),
forwardModel,
legacyCompact,
))
// 提前校验 function_call_output 是否具备可关联上下文,避免上游 400。
if !h.validateFunctionCallOutputRequest(c, body, reqLog) {
@@ -420,7 +453,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
if h.rejectIfCyberSessionBlocked(c, apiKey, sessionHashBody, reqModel, cyberBlockFormatResponses) {
return
}
requireCompact := isOpenAIRemoteCompactPath(c)
requireCompact := legacyCompact
maxAccountSwitches := h.maxAccountSwitches
switchCount := 0
@@ -437,7 +470,8 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
// 仅对 OpenAI 平台生效:Grok 生图走独立的 forwardGrokResponses 路径,不应被过滤。
// 复用前置权限与并发阶段在未修改 body 上确认的显式生图意图,避免大 tools 请求重复扫描。
// 该判断已排除 Codex 被动 image_gen namespace,避免 CC-only 账号被误过滤(#4476)。
requiredCapability := openAIResponsesRequiredCapability(imageIntent, requestPlatform)
needsResponses := nativeV2 || legacyCompact
requiredCapability := openAIResponsesRequiredCapabilityForRequest(imageIntent, needsResponses, requestPlatform)
// 分组利润控制:请求级装配定价上下文——pricingAt 固定本请求的
// D 与计费高峰因子,选号、槽位终检与全部 failover 重入共用同一门与阈值。
@@ -479,7 +513,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
zap.Int("excluded_account_count", len(failedAccountIDs)),
)
if len(failedAccountIDs) == 0 {
if errors.Is(err, service.ErrNoAvailableCompactAccounts) {
if legacyCompact && errors.Is(err, service.ErrNoAvailableCompactAccounts) {
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "compact_not_supported", "No available accounts support /responses/compact", streamStarted)
return
@@ -570,6 +604,50 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
if err == nil && result != nil && result.FirstTokenMs != nil {
service.SetOpsLatencyMs(c, service.OpsTimeToFirstTokenMsKey, int64(*result.FirstTokenMs))
}
// #5148 对齐:错误返回携带的部分 result(流中断前上游已计量的 usage)照常
// 入账;failover 错误恒定 result=nil,不会重复计费。
submitResponsesUsage := func(res *service.OpenAIForwardResult) {
if res == nil {
return
}
userAgent := c.GetHeader("User-Agent")
clientIP := ip.GetClientIP(c)
requestPayloadHash := service.HashUsageRequestPayload(body)
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, res)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
h.submitOpenAIUsageRecordTask(c.Request.Context(), res, func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
Result: res,
APIKey: apiKey,
User: apiKey.User,
Account: account,
Subscription: subscription,
InboundEndpoint: inboundEndpoint,
UpstreamEndpoint: upstreamEndpoint,
UserAgent: userAgent,
IPAddress: clientIP,
RequestPayloadHash: requestPayloadHash,
APIKeyService: h.apiKeyService,
QuotaPlatform: quotaPlatform,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, res.UpstreamModel),
PricingAt: pricingAt,
CyberBlocked: cyberBlocked,
}); err != nil {
logger.L().With(
zap.String("component", "handler.openai_gateway.responses"),
zap.Int64("user_id", subject.UserID),
zap.Int64("api_key_id", apiKey.ID),
zap.Any("group_id", apiKey.GroupID),
zap.String("model", reqModel),
zap.Int64("account_id", account.ID),
).Error("openai.record_usage_failed", zap.Error(err))
}
})
}
if err != nil {
if result != nil && result.ImageCount > 0 {
reqLog.Warn("openai.forward_partial_error_with_image_result",
@@ -669,6 +747,7 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
zap.Bool("upstream_error_response_already_written", upstreamErrorAlreadyCommunicated),
zap.Error(err),
}
submitResponsesUsage(result)
if shouldLogOpenAIForwardFailureAsWarn(c, wroteFallback) {
reqLog.Warn("openai.forward_failed", fields...)
return
@@ -687,46 +766,8 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(reqModel), openAIForwardSucceededForScheduling(result), nil)
}
// 捕获请求信息(用于异步记录,避免在 goroutine 中访问 gin.Context)
userAgent := c.GetHeader("User-Agent")
clientIP := ip.GetClientIP(c)
requestPayloadHash := service.HashUsageRequestPayload(body)
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
// 使用量记录通过有界 worker 池提交,避免请求热路径创建无界 goroutine。
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
Result: result,
APIKey: apiKey,
User: apiKey.User,
Account: account,
Subscription: subscription,
InboundEndpoint: inboundEndpoint,
UpstreamEndpoint: upstreamEndpoint,
UserAgent: userAgent,
IPAddress: clientIP,
RequestPayloadHash: requestPayloadHash,
APIKeyService: h.apiKeyService,
QuotaPlatform: quotaPlatform,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMapping, reqModel, result.UpstreamModel),
PricingAt: pricingAt,
CyberBlocked: cyberBlocked,
}); err != nil {
logger.L().With(
zap.String("component", "handler.openai_gateway.responses"),
zap.Int64("user_id", subject.UserID),
zap.Int64("api_key_id", apiKey.ID),
zap.Any("group_id", apiKey.GroupID),
zap.String("model", reqModel),
zap.Int64("account_id", account.ID),
).Error("openai.record_usage_failed", zap.Error(err))
}
})
submitResponsesUsage(result)
reqLog.Debug("openai.request_completed",
zap.Int64("account_id", account.ID),
zap.Int("switch_count", switchCount),
@@ -735,12 +776,8 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
}
}
func isOpenAIRemoteCompactPath(c *gin.Context) bool {
if c == nil || c.Request == nil || c.Request.URL == nil {
return false
}
normalizedPath := strings.TrimRight(strings.TrimSpace(c.Request.URL.Path), "/")
return strings.HasSuffix(normalizedPath, "/responses/compact")
func isOpenAILegacyCompactPath(c *gin.Context) bool {
return service.IsOpenAIResponsesCompactPath(c)
}
// isBareOpenAIResponsesPath 仅匹配裸 /responses 端点(无 /compact 等子路径),
@@ -750,32 +787,27 @@ func isBareOpenAIResponsesPath(c *gin.Context) bool {
return false
}
normalizedPath := strings.TrimRight(strings.TrimSpace(c.Request.URL.Path), "/")
return strings.HasSuffix(normalizedPath, "/responses")
}
func isOpenAIRemoteCompactionV2Request(c *gin.Context, body []byte) bool {
stream, valid := parseOpenAICompatibleStream(body)
if !valid || !stream || c == nil || c.Request == nil {
switch normalizedPath {
case EndpointResponses, "/openai/v1/responses", "/responses", "/backend-api/codex/responses":
return true
default:
return false
}
for _, header := range c.Request.Header.Values("x-codex-beta-features") {
for _, feature := range strings.Split(header, ",") {
if strings.TrimSpace(feature) == "remote_compaction_v2" {
return true
}
}
}
return false
}
func isOpenAIRemoteCompactionV2Request(body []byte) bool {
stream, valid := parseOpenAICompatibleStream(body)
return valid && stream && service.HasCompactionTriggerInInput(body)
}
// normalizeOpenAIResponsesCompactRequest keeps Codex remote compaction v2 on
// its native streaming /responses wire and preserves the legacy body-signal
// promotion for clients that do not explicitly advertise that protocol.
// promotion for non-streaming requests.
// 返回归一化后的 body;ok=false 表示错误响应已写出,调用方应直接 return。
func (h *OpenAIGatewayHandler) normalizeOpenAIResponsesCompactRequest(c *gin.Context, reqLog *zap.Logger, body []byte) ([]byte, bool) {
isCompactRequest := service.IsOpenAIResponsesCompactPathForTest(c)
isCompactRequest := isOpenAILegacyCompactPath(c)
if !isCompactRequest && isBareOpenAIResponsesPath(c) && service.HasCompactionTriggerInInput(body) {
if isOpenAIRemoteCompactionV2Request(c, body) {
if isOpenAIRemoteCompactionV2Request(body) {
return body, true
}
c.Request.URL.Path = strings.TrimRight(c.Request.URL.Path, "/") + "/compact"
@@ -804,7 +836,7 @@ func (h *OpenAIGatewayHandler) normalizeOpenAIResponsesCompactRequest(c *gin.Con
}
func (h *OpenAIGatewayHandler) logOpenAIRemoteCompactOutcome(c *gin.Context, startedAt time.Time) {
if !isOpenAIRemoteCompactPath(c) {
if !isOpenAILegacyCompactPath(c) {
return
}
@@ -1127,6 +1159,51 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
if err == nil && result != nil && result.FirstTokenMs != nil {
service.SetOpsLatencyMs(c, service.OpsTimeToFirstTokenMsKey, int64(*result.FirstTokenMs))
}
// Forward 与错误一起返回的部分结果:流中断/客户端断开排水前上游已计量的
// usage 照常入账,避免上游已产生消耗的请求完全漏记(#5148,对齐 anthropic
// 网关同名修复)。failover 错误恒定 result=nil,不会重复计费。
submitMessagesUsage := func(res *service.OpenAIForwardResult) {
if res == nil {
return
}
userAgent := c.GetHeader("User-Agent")
clientIP := ip.GetClientIP(c)
requestPayloadHash := service.HashUsageRequestPayload(body)
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, res)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
h.submitOpenAIUsageRecordTask(c.Request.Context(), res, func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
Result: res,
APIKey: apiKey,
User: apiKey.User,
Account: account,
Subscription: subscription,
InboundEndpoint: inboundEndpoint,
UpstreamEndpoint: upstreamEndpoint,
UserAgent: userAgent,
IPAddress: clientIP,
RequestPayloadHash: requestPayloadHash,
APIKeyService: h.apiKeyService,
QuotaPlatform: quotaPlatform,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMappingMsg, reqModel, res.UpstreamModel),
PricingAt: pricingAt,
CyberBlocked: cyberBlocked,
}); err != nil {
logger.L().With(
zap.String("component", "handler.openai_gateway.messages"),
zap.Int64("user_id", subject.UserID),
zap.Int64("api_key_id", apiKey.ID),
zap.Any("group_id", apiKey.GroupID),
zap.String("model", reqModel),
zap.Int64("account_id", account.ID),
).Error("openai_messages.record_usage_failed", zap.Error(err))
}
})
}
if err != nil {
if result != nil && result.ImageCount > 0 {
reqLog.Warn("openai_messages.forward_partial_error_with_image_result",
@@ -1201,6 +1278,9 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
zap.Int64("account_id", account.ID),
zap.Error(err),
)
// 断开排水期间上游已计量的 usage 必须入账(此前直接 return 丢弃,
// payg 上游照常计费而平台漏记)。
submitMessagesUsage(result)
return
}
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), false, nil)
@@ -1210,6 +1290,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
zap.Bool("fallback_error_response_written", wroteFallback),
zap.Error(err),
)
submitMessagesUsage(result)
return
}
}
@@ -1219,44 +1300,7 @@ func (h *OpenAIGatewayHandler) Messages(c *gin.Context) {
h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, account.GetMappedModel(currentRoutingModel), true, nil)
}
userAgent := c.GetHeader("User-Agent")
clientIP := ip.GetClientIP(c)
requestPayloadHash := service.HashUsageRequestPayload(body)
inboundEndpoint := GetInboundEndpoint(c)
upstreamEndpoint := resolveOpenAIUpstreamEndpoint(c, account, result)
quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
sessionID := service.ExtractClientSessionID(c)
cyberBlocked := service.GetOpsCyberPolicy(c) != nil
h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) {
if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
Result: result,
APIKey: apiKey,
User: apiKey.User,
Account: account,
Subscription: subscription,
InboundEndpoint: inboundEndpoint,
UpstreamEndpoint: upstreamEndpoint,
UserAgent: userAgent,
IPAddress: clientIP,
RequestPayloadHash: requestPayloadHash,
APIKeyService: h.apiKeyService,
QuotaPlatform: quotaPlatform,
SessionID: sessionID,
ChannelUsageFields: clientRequestedUsageFields(c, channelMappingMsg, reqModel, result.UpstreamModel),
PricingAt: pricingAt,
CyberBlocked: cyberBlocked,
}); err != nil {
logger.L().With(
zap.String("component", "handler.openai_gateway.messages"),
zap.Int64("user_id", subject.UserID),
zap.Int64("api_key_id", apiKey.ID),
zap.Any("group_id", apiKey.GroupID),
zap.String("model", reqModel),
zap.Int64("account_id", account.ID),
).Error("openai_messages.record_usage_failed", zap.Error(err))
}
})
submitMessagesUsage(result)
reqLog.Debug("openai_messages.request_completed",
zap.Int64("account_id", account.ID),
zap.Int("switch_count", switchCount),
@@ -5,8 +5,7 @@ package handler
// 槽位终检与生图跳门回归(handler 半程):
// - 槽位获取成功后的利润终检:越线账号释放槽位并要求调用方排除重选,
// 不写响应、不绑定粘连;
// - openAIResponsesRequiredCapability 的生图意图映射钉死(scheduler 的
// 跳门条件依赖 CapabilityResponses ⇔ 显式生图意图这一耦合)。
// - openAIResponsesRequiredCapability 的请求能力映射覆盖生图与原生远程压缩。
import (
"context"
@@ -141,10 +140,10 @@ func TestAcquireResponsesAccountSlotProfitRecheck(t *testing.T) {
})
}
// scheduler 跳门条件依赖"CapabilityResponses 仅在显式生图意图时被要求"这一
// 映射;后续若扩展该 capability 的用途,本测试失败提示同步收窄跳门条件。
func TestOpenAIResponsesRequiredCapabilityPinsImageIntentMapping(t *testing.T) {
func TestOpenAIResponsesRequiredCapabilityForRequest(t *testing.T) {
require.Equal(t, service.OpenAIEndpointCapabilityResponses, openAIResponsesRequiredCapability(true, service.PlatformOpenAI))
require.Equal(t, service.OpenAIEndpointCapabilityChatCompletions, openAIResponsesRequiredCapability(false, service.PlatformOpenAI))
require.Equal(t, service.OpenAIEndpointCapabilityChatCompletions, openAIResponsesRequiredCapability(true, service.PlatformGrok))
require.Equal(t, service.OpenAIEndpointCapabilityResponses, openAIResponsesRequiredCapabilityForRequest(false, true, service.PlatformOpenAI))
require.Equal(t, service.OpenAIEndpointCapabilityChatCompletions, openAIResponsesRequiredCapabilityForRequest(false, true, service.PlatformGrok))
}
@@ -0,0 +1,79 @@
package handler
import (
"encoding/json"
"fmt"
"strings"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/gin-gonic/gin"
)
type grokStandaloneSearchRequest struct {
Query string `json:"query"`
Input string `json:"input"`
MaxResults *int `json:"max_results"`
AllowedXHandles []string `json:"allowed_x_handles"`
ExcludedXHandles []string `json:"excluded_x_handles"`
FromDate string `json:"from_date"`
ToDate string `json:"to_date"`
EnableImageUnderstanding *bool `json:"enable_image_understanding"`
EnableVideoUnderstanding *bool `json:"enable_video_understanding"`
}
// XSearch marks the standalone endpoint so WebSearch can use native x_search
// while retaining its dedicated per-call billing contract.
func (h *GatewayHandler) XSearch(c *gin.Context) {
c.Set("grok_x_search_endpoint", true)
h.WebSearch(c)
}
func resolveGrokStandaloneSearchModel() string {
return xai.ResolveDefaultTextModel(xai.RuntimeModelMappingOptions().DefaultText)
}
func buildGrokXSearchResponsesBody(req grokStandaloneSearchRequest, model string) ([]byte, error) {
input := strings.TrimSpace(req.Query)
if input == "" {
input = strings.TrimSpace(req.Input)
}
tool := map[string]any{"type": "x_search"}
if len(req.AllowedXHandles) > 0 {
tool["allowed_x_handles"] = req.AllowedXHandles
}
if len(req.ExcludedXHandles) > 0 {
tool["excluded_x_handles"] = req.ExcludedXHandles
}
if strings.TrimSpace(req.FromDate) != "" {
tool["from_date"] = strings.TrimSpace(req.FromDate)
}
if strings.TrimSpace(req.ToDate) != "" {
tool["to_date"] = strings.TrimSpace(req.ToDate)
}
if req.EnableImageUnderstanding != nil {
tool["enable_image_understanding"] = *req.EnableImageUnderstanding
}
if req.EnableVideoUnderstanding != nil {
tool["enable_video_understanding"] = *req.EnableVideoUnderstanding
}
maxResults := 0
if req.MaxResults != nil {
maxResults = *req.MaxResults
}
return json.Marshal(map[string]any{
"model": xai.ResolveDefaultTextModel(model),
"input": buildGrokXSearchPrompt(input, maxResults),
"tools": []map[string]any{tool},
"tool_choice": "required",
"include": []string{"x_search_call.action.sources"},
"store": false,
"stream": false,
})
}
func buildGrokXSearchPrompt(query string, maxResults int) string {
return fmt.Sprintf(`Search X for the user query below. Return ONLY valid JSON with this exact shape: {"results":[{"url":"https://...","title":"post or page title","snippet":"concise factual summary"}]}. Return at most %d unique results. Every URL must be an actual x_search source. Populate a non-empty title and snippet for every result. Do not wrap the JSON in markdown.
User query:
%s`, normalizeGrokWebSearchMaxResults(maxResults), query)
}
@@ -0,0 +1,58 @@
package handler
import (
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestBuildGrokXSearchResponsesBody(t *testing.T) {
t.Parallel()
understandImages := true
understandVideos := false
body, err := buildGrokXSearchResponsesBody(grokStandaloneSearchRequest{
Query: "latest posts from xAI",
AllowedXHandles: []string{"xai"},
ExcludedXHandles: []string{"spam"},
FromDate: "2026-08-01",
ToDate: "2026-08-10",
EnableImageUnderstanding: &understandImages,
EnableVideoUnderstanding: &understandVideos,
}, xai.DefaultTextModel)
require.NoError(t, err)
require.Equal(t, xai.DefaultTextModel, gjson.GetBytes(body, "model").String())
require.Contains(t, gjson.GetBytes(body, "input").String(), "latest posts from xAI")
require.Contains(t, gjson.GetBytes(body, "input").String(), "Return ONLY valid JSON")
require.Equal(t, "x_search_call.action.sources", gjson.GetBytes(body, "include.0").String())
require.Equal(t, "required", gjson.GetBytes(body, "tool_choice").String())
require.Equal(t, "x_search", gjson.GetBytes(body, "tools.0.type").String())
require.Equal(t, "xai", gjson.GetBytes(body, "tools.0.allowed_x_handles.0").String())
require.Equal(t, "spam", gjson.GetBytes(body, "tools.0.excluded_x_handles.0").String())
require.Equal(t, "2026-08-01", gjson.GetBytes(body, "tools.0.from_date").String())
require.Equal(t, "2026-08-10", gjson.GetBytes(body, "tools.0.to_date").String())
require.True(t, gjson.GetBytes(body, "tools.0.enable_image_understanding").Bool())
require.False(t, gjson.GetBytes(body, "tools.0.enable_video_understanding").Bool())
require.False(t, gjson.GetBytes(body, "store").Bool())
require.False(t, gjson.GetBytes(body, "stream").Bool())
}
func TestBuildGrokXSearchResponsesBodyAcceptsInputAlias(t *testing.T) {
t.Parallel()
body, err := buildGrokXSearchResponsesBody(grokStandaloneSearchRequest{Input: "latest posts from xAI"}, xai.DefaultTextModel)
require.NoError(t, err)
require.Contains(t, gjson.GetBytes(body, "input").String(), "latest posts from xAI")
}
func TestResolveGrokStandaloneSearchModelUsesRuntimeDefault(t *testing.T) {
original := xai.RuntimeModelMappingOptions()
t.Cleanup(func() { xai.SetRuntimeModelMappingOptions(original) })
xai.SetRuntimeModelMappingOptions(xai.ModelMappingOptions{DefaultText: "grok-4.6"})
model := resolveGrokStandaloneSearchModel()
body, err := buildGrokXSearchResponsesBody(grokStandaloneSearchRequest{Query: "latest posts from xAI"}, model)
require.NoError(t, err)
require.Equal(t, "grok-4.6", model)
require.Equal(t, model, gjson.GetBytes(body, "model").String())
}
+3
View File
@@ -23,6 +23,7 @@ func ProvideAdminHandlers(
geminiOAuthHandler *admin.GeminiOAuthHandler,
antigravityOAuthHandler *admin.AntigravityOAuthHandler,
grokOAuthHandler *admin.GrokOAuthHandler,
cnProviderHandler *admin.CNProviderHandler,
proxyHandler *admin.ProxyHandler,
redeemHandler *admin.RedeemHandler,
promoHandler *admin.PromoHandler,
@@ -63,6 +64,7 @@ func ProvideAdminHandlers(
GeminiOAuth: geminiOAuthHandler,
AntigravityOAuth: antigravityOAuthHandler,
GrokOAuth: grokOAuthHandler,
CNProvider: cnProviderHandler,
Proxy: proxyHandler,
Redeem: redeemHandler,
Promo: promoHandler,
@@ -253,6 +255,7 @@ var ProviderSet = wire.NewSet(
admin.NewGeminiOAuthHandler,
admin.NewAntigravityOAuthHandler,
admin.NewGrokOAuthHandler,
admin.NewCNProviderHandler,
admin.NewProxyHandler,
admin.NewRedeemHandler,
admin.NewPromoHandler,
@@ -63,6 +63,9 @@ func ResponsesToChatCompletionsRequest(req *ResponsesRequest) (*ChatCompletionsR
if tool.Function != nil {
declared[tool.Function.Name] = true
}
if strings.EqualFold(strings.TrimSpace(tool.Type), "x_search") {
declared["x_search"] = true
}
}
if tc := responsesToolChoiceToChatToolChoice(req.ToolChoice, declared); len(tc) > 0 {
out.ToolChoice = tc
@@ -847,6 +850,16 @@ func responsesToolsToChatTools(tools []ResponsesTool) ([]ChatTool, error) {
return nil, err
}
out = append(out, flattened...)
case "x_search":
out = append(out, ChatTool{
Type: "x_search",
AllowedXHandles: tool.AllowedXHandles,
ExcludedXHandles: tool.ExcludedXHandles,
FromDate: tool.FromDate,
ToDate: tool.ToDate,
EnableImageUnderstanding: tool.EnableImageUnderstanding,
EnableVideoUnderstanding: tool.EnableVideoUnderstanding,
})
}
// 其余类型(web_search、image_generation 等服务端工具)在 chat 上游没有
// 对应能力,维持丢弃。
@@ -948,6 +961,15 @@ func responsesToolChoiceToChatToolChoice(raw json.RawMessage, declared map[strin
}
var name string
switch rawString(choice["type"]) {
case "x_search":
if !declared["x_search"] {
return nil
}
out, err := json.Marshal(map[string]any{"type": "x_search"})
if err != nil {
return raw
}
return out
case "tool_search":
// tool_search 未被丢弃而是降级为同名 function 代理(见
// responsesToolsToChatTools),强制选择它同样降级为 function 选择,
@@ -661,6 +661,20 @@ func TestResponsesToChatCompletionsRequest_DropsToolChoiceForDroppedTool(t *test
require.Len(t, out.Tools, 1)
assert.Empty(t, out.ToolChoice, "指向被丢弃服务端工具的 tool_choice 必须丢弃")
out, err = ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "glm-5.2",
Input: json.RawMessage(`"hi"`),
Tools: []ResponsesTool{
{Type: "function", Name: "wait", Parameters: json.RawMessage(`{"type":"object","properties":{}}`)},
{Type: "web_search"},
{Type: "x_search"},
},
ToolChoice: json.RawMessage(`{"type":"function","name":"web_search"}`),
})
require.NoError(t, err)
require.Len(t, out.Tools, 2)
assert.Empty(t, out.ToolChoice, "surviving x_search must not keep a function tool_choice named web_search")
// 具名选择指向不存在的工具名。
out, err = ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "glm-5.2",
@@ -419,6 +419,18 @@ func convertChatToolsToResponses(tools []ChatTool, functions []ChatFunction) []R
var out []ResponsesTool
for _, t := range tools {
if strings.EqualFold(strings.TrimSpace(t.Type), "x_search") {
out = append(out, ResponsesTool{
Type: "x_search",
AllowedXHandles: t.AllowedXHandles,
ExcludedXHandles: t.ExcludedXHandles,
FromDate: t.FromDate,
ToDate: t.ToDate,
EnableImageUnderstanding: t.EnableImageUnderstanding,
EnableVideoUnderstanding: t.EnableVideoUnderstanding,
})
continue
}
if t.Type != "function" || t.Function == nil {
continue
}
@@ -0,0 +1,79 @@
package apicompat
import (
"encoding/json"
"testing"
"github.com/stretchr/testify/require"
)
func TestChatCompletionsToResponsesPreservesXSearchTool(t *testing.T) {
enabled := true
req := &ChatCompletionsRequest{
Model: "grok-4.5",
Messages: []ChatMessage{
{Role: "user", Content: json.RawMessage(`"latest xAI post"`)},
},
Tools: []ChatTool{{
Type: "x_search",
AllowedXHandles: []string{"xai"},
ExcludedXHandles: []string{"spam"},
FromDate: "2026-08-01",
ToDate: "2026-08-10",
EnableImageUnderstanding: &enabled,
EnableVideoUnderstanding: &enabled,
}},
ToolChoice: json.RawMessage(`{"type":"x_search"}`),
}
resp, err := ChatCompletionsToResponses(req)
require.NoError(t, err)
require.Len(t, resp.Tools, 1)
require.Equal(t, "x_search", resp.Tools[0].Type)
require.Equal(t, []string{"xai"}, resp.Tools[0].AllowedXHandles)
require.Equal(t, []string{"spam"}, resp.Tools[0].ExcludedXHandles)
require.Equal(t, "2026-08-01", resp.Tools[0].FromDate)
require.Equal(t, "2026-08-10", resp.Tools[0].ToDate)
require.NotNil(t, resp.Tools[0].EnableImageUnderstanding)
require.True(t, *resp.Tools[0].EnableImageUnderstanding)
require.NotNil(t, resp.Tools[0].EnableVideoUnderstanding)
require.True(t, *resp.Tools[0].EnableVideoUnderstanding)
require.JSONEq(t, `{"type":"x_search"}`, string(resp.ToolChoice))
}
func TestResponsesToChatCompletionsPreservesXSearchTool(t *testing.T) {
enabled := true
req := &ResponsesRequest{
Model: "grok-4.5",
Input: json.RawMessage(`"latest xAI post"`),
Tools: []ResponsesTool{{
Type: "x_search",
AllowedXHandles: []string{"xai"},
ExcludedXHandles: []string{"spam"},
FromDate: "2026-08-01",
ToDate: "2026-08-10",
EnableImageUnderstanding: &enabled,
EnableVideoUnderstanding: &enabled,
}},
ToolChoice: json.RawMessage(`{"type":"x_search"}`),
}
chat, err := ResponsesToChatCompletionsRequest(req)
require.NoError(t, err)
require.Len(t, chat.Tools, 1)
require.Equal(t, "x_search", chat.Tools[0].Type)
require.Equal(t, []string{"xai"}, chat.Tools[0].AllowedXHandles)
require.Equal(t, []string{"spam"}, chat.Tools[0].ExcludedXHandles)
require.JSONEq(t, `{"type":"x_search"}`, string(chat.ToolChoice))
}
func TestResponsesToChatCompletionsXSearchToolChoiceString(t *testing.T) {
chat, err := ResponsesToChatCompletionsRequest(&ResponsesRequest{
Model: "grok-4.5",
Input: json.RawMessage(`"latest xAI post"`),
Tools: []ResponsesTool{{Type: "x_search"}},
ToolChoice: json.RawMessage(`"x_search"`),
})
require.NoError(t, err)
require.JSONEq(t, `"x_search"`, string(chat.ToolChoice))
}
+18 -2
View File
@@ -302,7 +302,7 @@ type ResponsesContentPart struct {
// ResponsesTool describes a tool in the Responses API.
type ResponsesTool struct {
Type string `json:"type"` // "function" | "custom" | "web_search" | "local_shell" etc.
Type string `json:"type"` // "function" | "custom" | "web_search" | "x_search" | "local_shell" etc.
Name string `json:"name,omitempty"`
Description string `json:"description,omitempty"`
Parameters json.RawMessage `json:"parameters,omitempty"`
@@ -311,6 +311,14 @@ type ResponsesTool struct {
// type=namespace 的子工具列表(tools 与 children 二选一,语义相同)。
Tools []ResponsesTool `json:"tools,omitempty"`
Children []ResponsesTool `json:"children,omitempty"`
// type=x_search
AllowedXHandles []string `json:"allowed_x_handles,omitempty"`
ExcludedXHandles []string `json:"excluded_x_handles,omitempty"`
FromDate string `json:"from_date,omitempty"`
ToDate string `json:"to_date,omitempty"`
EnableImageUnderstanding *bool `json:"enable_image_understanding,omitempty"`
EnableVideoUnderstanding *bool `json:"enable_video_understanding,omitempty"`
}
// UnmarshalJSON 容忍字符串形式的工具声明:codex 会以 "name" 简写声明 custom 工具,
@@ -675,8 +683,16 @@ type ChatImageURL struct {
// ChatTool describes a tool available to the model.
type ChatTool struct {
Type string `json:"type"` // "function"
Type string `json:"type"` // "function" | "x_search"
Function *ChatFunction `json:"function,omitempty"`
// type=x_search
AllowedXHandles []string `json:"allowed_x_handles,omitempty"`
ExcludedXHandles []string `json:"excluded_x_handles,omitempty"`
FromDate string `json:"from_date,omitempty"`
ToDate string `json:"to_date,omitempty"`
EnableImageUnderstanding *bool `json:"enable_image_understanding,omitempty"`
EnableVideoUnderstanding *bool `json:"enable_video_understanding,omitempty"`
}
// ChatFunction describes a function tool definition.
@@ -115,11 +115,12 @@ type EndpointStat struct {
ActualCost float64 `json:"actual_cost"` // 实际扣除
}
// GroupUsageSummary represents today's and cumulative cost for a single group.
// GroupUsageSummary represents today's, yesterday's, and cumulative cost for a single group.
type GroupUsageSummary struct {
GroupID int64 `json:"group_id"`
TodayCost float64 `json:"today_cost"`
TotalCost float64 `json:"total_cost"`
GroupID int64 `json:"group_id"`
TodayCost float64 `json:"today_cost"`
YesterdayCost float64 `json:"yesterday_cost"`
TotalCost float64 `json:"total_cost"`
}
// GroupStat represents usage statistics for a single group
+3
View File
@@ -83,6 +83,7 @@ func (o ModelMappingOptions) defaultText() string {
var defaultModels = []Model{
// Text
{ID: "grok-4.6", Object: "model", Type: "model", OwnedBy: "xai", DisplayName: "Grok 4.6"},
{ID: "grok-4.5", Object: "model", Type: "model", OwnedBy: "xai", DisplayName: "Grok 4.5"},
{ID: "grok-4.3", Object: "model", Type: "model", OwnedBy: "xai", DisplayName: "Grok 4.3"},
{ID: "grok-3-mini", Object: "model", Type: "model", OwnedBy: "xai", DisplayName: "Grok 3 Mini"},
@@ -106,6 +107,8 @@ var defaultModels = []Model{
var grokTextResponsesModelAliases = map[string]string{
"grok": DefaultTextModel,
"grok-latest": DefaultTextModel,
"grok-4.6": "grok-4.6",
"grok-4.6-latest": "grok-4.6",
"grok-4.5": DefaultTextModel,
"grok-4.5-latest": DefaultTextModel,
"grok-4.3": "grok-4.3",
+9
View File
@@ -52,11 +52,20 @@ func TestCanonicalImagineVideoModel(t *testing.T) {
func TestIsGrokModelID(t *testing.T) {
t.Parallel()
require.True(t, IsGrokModelID("grok-4.5"))
require.True(t, IsGrokModelID("grok-4.6"))
require.True(t, IsGrokModelID("x-ai/grok-4.3"))
require.False(t, IsGrokModelID("gpt-5"))
require.False(t, IsGrokModelID("claude-sonnet-4"))
}
func TestDefaultModelsIncludesGrok46(t *testing.T) {
t.Parallel()
ids := DefaultModelIDs()
require.Contains(t, ids, "grok-4.6")
require.Equal(t, "grok-4.6", ResolveGrokTextResponsesModelID("grok-4.6"))
require.Equal(t, "grok-4.6", ResolveGrokTextResponsesModelID("grok-4.6-latest"))
}
func TestResolveGrokTextResponsesModelID(t *testing.T) {
t.Parallel()
require.Equal(t, "grok-4.5", ResolveGrokTextResponsesModelID(""))
+2
View File
@@ -352,6 +352,8 @@ func TestDefaultModelMappingIncludesGrokAliases(t *testing.T) {
mapping := DefaultModelMapping()
require.Equal(t, "grok-4.5", mapping["grok"])
require.Equal(t, "grok-4.5", mapping["grok-latest"])
require.Equal(t, "grok-4.6", mapping["grok-4.6"])
require.Equal(t, "grok-4.6", mapping["grok-4.6-latest"])
require.Equal(t, "grok-4.5", mapping["grok-4.5"])
require.Equal(t, "grok-4.5", mapping["grok-4.5-latest"])
require.Equal(t, "grok-build-0.1", mapping["grok-build"])
+6
View File
@@ -43,6 +43,12 @@ type QuotaSnapshot struct {
LastProbeAt string `json:"last_probe_at,omitempty"`
LastHeadersSeenAt string `json:"last_headers_seen_at,omitempty"`
UpdatedAt string `json:"updated_at"`
// Model is the upstream id that produced these rate-limit headers.
Model string `json:"model,omitempty"`
// PlanFrom45Responses is inferred from a grok-4.5 Responses window
// (8300/53M = Heavy). Carried across later non-4.5 overwrites.
PlanFrom45Responses string `json:"plan_from_45_responses,omitempty"`
PlanFrom45ResponsesAt string `json:"plan_from_45_responses_at,omitempty"`
}
func (s *QuotaSnapshot) HasObservedHeaders() bool {
@@ -0,0 +1,279 @@
package xai
import (
"encoding/json"
"strconv"
"strings"
"time"
)
// GrokQuotaSignalMaxAge bounds how long a grok-4.5 Responses window can
// influence SuperGrok vs Heavy inference.
const GrokQuotaSignalMaxAge = 24 * time.Hour
const (
grok45ResponsesModel = "grok-4.5"
grokHeavyQuotaRequestLimit int64 = 8_300
grokHeavyQuotaTokenLimit int64 = 53_000_000
)
// MapJWTSubscriptionTier maps prod_auth.SubscriptionTier numeric JWT claims
// to stable snake_case keys used by Grok Build / Mixpanel.
func MapJWTSubscriptionTier(tier uint64) string {
switch tier {
case 0:
return "free"
case 1:
return "supergrok"
case 2:
return "x_basic"
case 3:
return "x_premium"
case 4:
return "x_premium_plus"
case 5:
return "supergrok_heavy"
case 6:
return "supergrok_lite"
case 7:
return "supergrok_plus"
default:
return strconv.FormatUint(tier, 10)
}
}
// NormalizeSubscriptionTier canonicalizes display names, /user strings, and
// JWT-derived keys onto the same snake_case identifiers.
func NormalizeSubscriptionTier(raw string) string {
t := strings.ToLower(strings.TrimSpace(raw))
t = strings.ReplaceAll(t, "-", "_")
t = strings.Join(strings.Fields(t), "_")
switch t {
case "free", "grok_free", "grokfree", "free_tier", "freetier", "grok_basic", "grokbasic":
return "free"
case "supergrok", "grokpro":
return "supergrok"
case "supergrok_lite", "supergroklite":
return "supergrok_lite"
case "supergrok_heavy", "supergrokheavy":
return "supergrok_heavy"
case "supergrok_pro", "supergrokpro":
return "supergrok_pro"
case "supergrok_plus", "supergrokplus":
return "supergrok_plus"
case "x_basic", "xbasic", "basic":
return "x_basic"
case "x_premium", "xpremium":
return "x_premium"
case "x_premium_plus", "xpremiumplus", "x_premium+":
return "x_premium_plus"
default:
return t
}
}
// SubscriptionTierFromJWT decodes an access token payload (no signature check)
// and maps the numeric or string `tier` claim.
func SubscriptionTierFromJWT(jwt string) string {
claims := DecodeJWTClaims(jwt)
if claims == nil {
return ""
}
raw, ok := claims["tier"]
if !ok || raw == nil {
return ""
}
switch v := raw.(type) {
case float64:
if v < 0 {
return ""
}
return MapJWTSubscriptionTier(uint64(v))
case json.Number:
n, err := v.Int64()
if err != nil || n < 0 {
return NormalizeSubscriptionTier(v.String())
}
return MapJWTSubscriptionTier(uint64(n))
case string:
trimmed := strings.TrimSpace(v)
if trimmed == "" {
return ""
}
if n, err := strconv.ParseUint(trimmed, 10, 64); err == nil {
return MapJWTSubscriptionTier(n)
}
return NormalizeSubscriptionTier(trimmed)
default:
return ""
}
}
// CanonicalGrokPlan resolves SuperGrok vs Heavy when the provider label is
// ambiguous (SuperGrokPro). JWT numeric claims are applied by the caller first.
// Monthly $150/$1500 limits still win when present.
// Rate-limit windows are only used when they came from grok-4.5 Responses.
func CanonicalGrokPlan(monthlyLimitCents *float64, subscriptionTier string, quota *QuotaSnapshot) string {
if plan := resolvePlan(monthlyLimitCents); plan != "" {
return NormalizeSubscriptionTier(plan)
}
normalized := NormalizeSubscriptionTier(subscriptionTier)
switch normalized {
case "free", "x_basic":
return "free"
case "supergrok_heavy":
return "supergrok_heavy"
case "supergrok_lite":
return "supergrok_lite"
case "supergrok_plus":
return "supergrok_plus"
}
if isAmbiguousGrokPaidPlan(normalized) {
if hint := Grok45ResponsesPlanHint(quota, time.Time{}); hint != "" {
return hint
}
return "supergrok"
}
return ""
}
func isAmbiguousGrokPaidPlan(normalized string) bool {
switch normalized {
case "supergrok", "supergrok_pro", "paid", "pro":
return true
default:
return false
}
}
// IsGrok45ResponsesQuotaModel reports whether model is the grok-4.5 Responses
// id (or a dated grok-4.5-* variant). Empty and other families are false.
func IsGrok45ResponsesQuotaModel(model string) bool {
m := strings.ToLower(strings.TrimSpace(StripGrokProviderPrefix(model)))
return m == grok45ResponsesModel || strings.HasPrefix(m, grok45ResponsesModel+"-")
}
// Grok45ResponsesPlanHint returns SuperGrok / Heavy inferred from a grok-4.5
// Responses window. Other models' limits are ignored.
func Grok45ResponsesPlanHint(quota *QuotaSnapshot, now time.Time) string {
if quota == nil {
return ""
}
if plan := NormalizeSubscriptionTier(quota.PlanFrom45Responses); plan == "supergrok" || plan == "supergrok_heavy" {
if isQuotaTimestampFresh(quota.PlanFrom45ResponsesAt, now) {
return plan
}
}
if !IsGrok45ResponsesQuotaModel(quota.Model) || !IsQuotaSnapshotFresh(quota, now) {
return ""
}
if quotaLooksLikeGrokHeavy(quota) {
return "supergrok_heavy"
}
return ""
}
// ApplyGrok45ResponsesPlanSignal records a grok-4.5 Heavy/SuperGrok hint, or
// copies the previous 4.5 hint when this observation is a different model.
func (s *QuotaSnapshot) ApplyGrok45ResponsesPlanSignal(prev *QuotaSnapshot) {
if s == nil {
return
}
observedAt := firstNonEmptyQuotaTime(s.LastHeadersSeenAt, s.UpdatedAt)
if IsGrok45ResponsesQuotaModel(s.Model) && quotaHasLimitWindow(s) {
if quotaLooksLikeGrokHeavy(s) {
s.PlanFrom45Responses = "supergrok_heavy"
s.PlanFrom45ResponsesAt = observedAt
return
}
s.PlanFrom45Responses = "supergrok"
s.PlanFrom45ResponsesAt = observedAt
return
}
if prev != nil && strings.TrimSpace(prev.PlanFrom45Responses) != "" {
s.PlanFrom45Responses = prev.PlanFrom45Responses
s.PlanFrom45ResponsesAt = prev.PlanFrom45ResponsesAt
}
}
// QuotaSnapshotObservedAt prefers LastHeadersSeenAt over UpdatedAt so a later
// snapshot rewrite cannot refresh a stale Heavy window.
func QuotaSnapshotObservedAt(snapshot *QuotaSnapshot) (time.Time, bool) {
if snapshot == nil {
return time.Time{}, false
}
return parseQuotaTimestamp(firstNonEmptyQuotaTime(snapshot.LastHeadersSeenAt, snapshot.UpdatedAt))
}
// IsQuotaSnapshotFresh reports whether a quota signal is recent enough to
// distinguish SuperGrok from Heavy.
func IsQuotaSnapshotFresh(snapshot *QuotaSnapshot, now time.Time) bool {
observedAt, ok := QuotaSnapshotObservedAt(snapshot)
if !ok {
return false
}
return isTimeFresh(observedAt, now)
}
func isQuotaTimestampFresh(raw string, now time.Time) bool {
parsed, ok := parseQuotaTimestamp(raw)
if !ok {
return false
}
return isTimeFresh(parsed, now)
}
func parseQuotaTimestamp(raw string) (time.Time, bool) {
raw = strings.TrimSpace(raw)
if raw == "" {
return time.Time{}, false
}
parsed, err := time.Parse(time.RFC3339, raw)
if err != nil {
return time.Time{}, false
}
return parsed, true
}
func isTimeFresh(observedAt, now time.Time) bool {
if now.IsZero() {
now = time.Now()
}
age := now.Sub(observedAt)
return age <= GrokQuotaSignalMaxAge && age >= -5*time.Minute
}
func quotaHasLimitWindow(quota *QuotaSnapshot) bool {
if quota == nil {
return false
}
if quota.Requests != nil && quota.Requests.Limit != nil {
return true
}
return quota.Tokens != nil && quota.Tokens.Limit != nil
}
func quotaLooksLikeGrokHeavy(quota *QuotaSnapshot) bool {
if quota == nil {
return false
}
var requestLimit, tokenLimit int64
if quota.Requests != nil && quota.Requests.Limit != nil {
requestLimit = *quota.Requests.Limit
}
if quota.Tokens != nil && quota.Tokens.Limit != nil {
tokenLimit = *quota.Tokens.Limit
}
return requestLimit >= grokHeavyQuotaRequestLimit || tokenLimit >= grokHeavyQuotaTokenLimit
}
func firstNonEmptyQuotaTime(values ...string) string {
for _, value := range values {
if strings.TrimSpace(value) != "" {
return strings.TrimSpace(value)
}
}
return ""
}
@@ -0,0 +1,144 @@
//go:build unit
package xai
import (
"encoding/base64"
"encoding/json"
"testing"
"time"
"github.com/stretchr/testify/require"
)
func TestMapJWTSubscriptionTierNumber(t *testing.T) {
t.Parallel()
require.Equal(t, "free", MapJWTSubscriptionTier(0))
require.Equal(t, "supergrok", MapJWTSubscriptionTier(1))
require.Equal(t, "x_basic", MapJWTSubscriptionTier(2))
require.Equal(t, "x_premium", MapJWTSubscriptionTier(3))
require.Equal(t, "x_premium_plus", MapJWTSubscriptionTier(4))
require.Equal(t, "supergrok_heavy", MapJWTSubscriptionTier(5))
require.Equal(t, "supergrok_lite", MapJWTSubscriptionTier(6))
require.Equal(t, "supergrok_plus", MapJWTSubscriptionTier(7))
require.Equal(t, "9", MapJWTSubscriptionTier(9))
}
func TestNormalizeSubscriptionTierAliases(t *testing.T) {
t.Parallel()
require.Equal(t, "free", NormalizeSubscriptionTier("Free"))
require.Equal(t, "free", NormalizeSubscriptionTier(" FREE "))
require.Equal(t, "supergrok", NormalizeSubscriptionTier("SuperGrok"))
require.Equal(t, "supergrok_heavy", NormalizeSubscriptionTier("SuperGrok Heavy"))
require.Equal(t, "supergrok_pro", NormalizeSubscriptionTier("SuperGrokPro"))
require.Equal(t, "supergrok_lite", NormalizeSubscriptionTier("SuperGrok Lite"))
require.Equal(t, "supergrok_lite", NormalizeSubscriptionTier("SuperGrokLite"))
require.Equal(t, "x_basic", NormalizeSubscriptionTier("X Basic"))
require.Equal(t, "free", NormalizeSubscriptionTier("free-tier"))
require.Equal(t, "free", NormalizeSubscriptionTier("free_tier"))
require.Equal(t, "free", NormalizeSubscriptionTier("grok-basic"))
require.Equal(t, "free", NormalizeSubscriptionTier("grok_basic"))
require.Equal(t, "supergrok_lite", NormalizeSubscriptionTier("supergrok_lite"))
}
func TestSubscriptionTierFromJWTUsesNumericClaim(t *testing.T) {
t.Parallel()
require.Equal(t, "supergrok_heavy", SubscriptionTierFromJWT(jwtWithClaims(t, map[string]any{"tier": 5})))
require.Equal(t, "free", SubscriptionTierFromJWT(jwtWithClaims(t, map[string]any{"tier": 0})))
require.Equal(t, "supergrok_lite", SubscriptionTierFromJWT(jwtWithClaims(t, map[string]any{"tier": 6})))
require.Empty(t, SubscriptionTierFromJWT(jwtWithClaims(t, map[string]any{"sub": "user"})))
require.Empty(t, SubscriptionTierFromJWT("not-a-jwt"))
}
func TestCanonicalGrokPlanUsesOnlyGrok45ResponsesWindow(t *testing.T) {
t.Parallel()
zero := float64(0)
heavyReq, heavyTok := int64(8300), int64(53_000_000)
superReq, superTok := int64(900), int64(15_000_000)
fresh := time.Now().UTC().Format(time.RFC3339)
stale := time.Now().Add(-GrokQuotaSignalMaxAge - time.Hour).UTC().Format(time.RFC3339)
require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", nil))
require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrok", nil))
require.Equal(t, "supergrok_heavy", CanonicalGrokPlan(&zero, "SuperGrok Heavy", nil))
require.Empty(t, CanonicalGrokPlan(&zero, "", nil))
require.Equal(t, "free", CanonicalGrokPlan(&zero, "free", nil))
require.Equal(t, "supergrok_heavy", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{
Model: "grok-4.5",
Requests: &QuotaWindow{Limit: &heavyReq},
Tokens: &QuotaWindow{Limit: &heavyTok},
LastHeadersSeenAt: fresh,
}))
require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{
Model: "grok-4.6",
Requests: &QuotaWindow{Limit: &heavyReq},
Tokens: &QuotaWindow{Limit: &heavyTok},
LastHeadersSeenAt: fresh,
}))
require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{
Requests: &QuotaWindow{Limit: &heavyReq},
Tokens: &QuotaWindow{Limit: &heavyTok},
LastHeadersSeenAt: fresh,
}))
require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{
Model: "grok-4.5",
Requests: &QuotaWindow{Limit: &superReq},
Tokens: &QuotaWindow{Limit: &superTok},
LastHeadersSeenAt: fresh,
}))
require.Equal(t, "supergrok", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{
Model: "grok-4.5",
Requests: &QuotaWindow{Limit: &heavyReq},
LastHeadersSeenAt: stale,
}))
require.Equal(t, "supergrok_heavy", CanonicalGrokPlan(&zero, "SuperGrokPro", &QuotaSnapshot{
Model: "grok-4.6",
Requests: &QuotaWindow{Limit: &superReq},
PlanFrom45Responses: "supergrok_heavy",
PlanFrom45ResponsesAt: fresh,
}))
require.Equal(t, "free", CanonicalGrokPlan(&zero, "free", &QuotaSnapshot{
Model: "grok-4.5",
Requests: &QuotaWindow{Limit: &heavyReq},
LastHeadersSeenAt: fresh,
}))
heavyCents := float64(SuperGrokHeavyLimitCents)
require.Equal(t, "supergrok_heavy", CanonicalGrokPlan(&heavyCents, "SuperGrokPro", nil))
}
func TestApplyGrok45ResponsesPlanSignalCarriesHint(t *testing.T) {
t.Parallel()
heavyReq := int64(8300)
fresh := time.Now().UTC().Format(time.RFC3339)
prev := &QuotaSnapshot{
Model: "grok-4.5",
Requests: &QuotaWindow{Limit: &heavyReq},
LastHeadersSeenAt: fresh,
}
prev.ApplyGrok45ResponsesPlanSignal(nil)
require.Equal(t, "supergrok_heavy", prev.PlanFrom45Responses)
next := &QuotaSnapshot{
Model: "grok-4.6",
Requests: &QuotaWindow{Limit: int64Ptr(100)},
LastHeadersSeenAt: fresh,
}
next.ApplyGrok45ResponsesPlanSignal(prev)
require.Equal(t, "supergrok_heavy", next.PlanFrom45Responses)
require.Equal(t, prev.PlanFrom45ResponsesAt, next.PlanFrom45ResponsesAt)
}
func int64Ptr(v int64) *int64 { return &v }
func jwtWithClaims(t *testing.T, claims map[string]any) string {
t.Helper()
payload, err := json.Marshal(claims)
require.NoError(t, err)
return "header." + base64.RawURLEncoding.EncodeToString(payload) + ".sig"
}
@@ -3,8 +3,10 @@ package repository
import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"log/slog"
"strings"
"time"
@@ -196,6 +198,8 @@ func (r *apiKeyRepository) GetByKeyForAuth(ctx context.Context, key string) (*se
group.FieldAudioRealtimePricePerMin,
group.FieldAudioTtsPricePerMillionChars,
group.FieldAudioSttPricePerHour,
group.FieldLongContextPricingEnabled,
group.FieldModelPricing,
group.FieldClaudeCodeOnly,
group.FieldFallbackGroupID,
group.FieldFallbackGroupIDOnInvalidRequest,
@@ -946,6 +950,14 @@ func groupEntityToService(g *dbent.Group) *service.Group {
if g == nil {
return nil
}
var modelPricing []service.ChannelModelPricing
if len(g.ModelPricing) > 0 {
if err := json.Unmarshal(g.ModelPricing, &modelPricing); err != nil {
slog.Warn("group model_pricing unmarshal failed; falling back to channel/builtin pricing",
"group_id", g.ID, "error", err)
modelPricing = nil
}
}
return &service.Group{
ID: g.ID,
Name: g.Name,
@@ -980,6 +992,8 @@ func groupEntityToService(g *dbent.Group) *service.Group {
AudioRealtimePricePerMin: g.AudioRealtimePricePerMin,
AudioTTSPricePerMillionChars: g.AudioTtsPricePerMillionChars,
AudioSTTPricePerHour: g.AudioSttPricePerHour,
LongContextPricingEnabled: g.LongContextPricingEnabled,
ModelPricing: modelPricing,
DefaultValidityDays: g.DefaultValidityDays,
ClaudeCodeOnly: g.ClaudeCodeOnly,
FallbackGroupID: g.FallbackGroupID,
@@ -0,0 +1,259 @@
package repository
import (
"context"
"database/sql"
"fmt"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/usagestats"
"github.com/Wei-Shaw/sub2api/internal/service"
)
func (r *usageLogRepository) getAllGroupUsageSummaryFromRollups(ctx context.Context, todayStart time.Time) (results []usagestats.GroupUsageSummary, err error) {
todayStart = service.GroupUsageTodayStart(todayStart)
yesterdayStart := service.GroupUsageYesterdayStart(todayStart)
timezoneName := service.GroupUsageTimezoneName()
todayDate := service.GroupUsageDate(todayStart)
yesterdayDate := service.GroupUsageDate(yesterdayStart)
const query = `
WITH state_values AS (
SELECT
COUNT(*) = 1
AND MAX(timezone_name) = $3
AND MAX(closed_before) <= $4::date AS valid,
MAX(closed_before) AS closed_before,
MAX(retained_from) AS retained_from
FROM usage_group_rollup_state
WHERE id = 1
),
state AS (
SELECT
CASE WHEN valid THEN closed_before ELSE DATE '1970-01-01' END AS closed_before,
CASE WHEN valid THEN retained_from ELSE TIMESTAMPTZ '1970-01-01 00:00:00+00' END AS retained_from,
CASE
WHEN valid THEN closed_before::timestamp AT TIME ZONE $3::text
ELSE TIMESTAMPTZ '1970-01-01 00:00:00+00'
END AS tail_start,
valid
FROM state_values
),
historical AS (
SELECT
rollup.group_id,
COALESCE(SUM(rollup.actual_cost), 0) AS actual_cost,
COALESCE(SUM(rollup.actual_cost) FILTER (
WHERE rollup.bucket_date = $5::date
), 0) AS yesterday_cost
FROM usage_group_daily_rollups rollup
CROSS JOIN state
WHERE state.valid
AND rollup.bucket_date >= (state.retained_from AT TIME ZONE $3::text)::date
AND rollup.bucket_date < state.closed_before
GROUP BY rollup.group_id
),
tail AS (
SELECT
ul.group_id,
COALESCE(SUM(ul.actual_cost), 0) AS actual_cost,
COALESCE(SUM(ul.actual_cost) FILTER (WHERE ul.created_at >= $1), 0) AS today_cost,
COALESCE(SUM(ul.actual_cost) FILTER (
WHERE ul.created_at >= $2
AND ul.created_at < $1
), 0) AS yesterday_cost
FROM usage_logs ul
CROSS JOIN state
WHERE ul.created_at >= state.tail_start
GROUP BY ul.group_id
)
SELECT
g.id AS group_id,
COALESCE(historical.actual_cost, 0) + COALESCE(tail.actual_cost, 0) AS total_cost,
COALESCE(tail.today_cost, 0) AS today_cost,
COALESCE(historical.yesterday_cost, 0) + COALESCE(tail.yesterday_cost, 0) AS yesterday_cost
FROM groups g
LEFT JOIN historical ON historical.group_id = g.id
LEFT JOIN tail ON tail.group_id = g.id
ORDER BY g.id
`
rows, err := r.sql.QueryContext(
ctx,
query,
todayStart,
yesterdayStart,
timezoneName,
todayDate,
yesterdayDate,
)
if err != nil {
return nil, err
}
defer func() {
if closeErr := rows.Close(); closeErr != nil && err == nil {
err = closeErr
results = nil
}
}()
results = make([]usagestats.GroupUsageSummary, 0)
for rows.Next() {
var row usagestats.GroupUsageSummary
if err := rows.Scan(&row.GroupID, &row.TotalCost, &row.TodayCost, &row.YesterdayCost); err != nil {
return nil, err
}
results = append(results, row)
}
if err := rows.Err(); err != nil {
return nil, err
}
return results, nil
}
// SyncGroupUsageRollups 将服务端配置时区今日以前的用量发布为分组日桶。
func (r *dashboardAggregationRepository) SyncGroupUsageRollups(ctx context.Context, todayStart time.Time) error {
if r == nil || r.sql == nil {
return nil
}
todayStart = service.GroupUsageTodayStart(todayStart)
if db, ok := r.sql.(*sql.DB); ok {
tx, err := db.BeginTx(ctx, nil)
if err != nil {
return err
}
txRepo := newDashboardAggregationRepositoryWithSQL(tx)
if err := txRepo.syncGroupUsageRollupsInTx(ctx, todayStart); err != nil {
_ = tx.Rollback()
return err
}
return tx.Commit()
}
return r.syncGroupUsageRollupsInTx(ctx, todayStart)
}
func (r *dashboardAggregationRepository) syncGroupUsageRollupsInTx(ctx context.Context, todayStart time.Time) error {
var closedBefore string
var previousRetainedFrom time.Time
var stateTimezoneName string
if err := scanSingleRow(ctx, r.sql, `
SELECT closed_before::text, retained_from, timezone_name
FROM usage_group_rollup_state
WHERE id = 1
FOR UPDATE
`, nil, &closedBefore, &previousRetainedFrom, &stateTimezoneName); err != nil {
return fmt.Errorf("读取分组用量汇总水位: %w", err)
}
todayDate := service.GroupUsageDate(todayStart)
timezoneName := service.GroupUsageTimezoneName()
timezoneChanged := stateTimezoneName != timezoneName
var closedTime time.Time
if !timezoneChanged {
var err error
closedTime, err = service.ParseGroupUsageDate(closedBefore)
if err != nil {
return fmt.Errorf("解析分组用量汇总水位 %q: %w", closedBefore, err)
}
todayDateTime, err := service.ParseGroupUsageDate(todayDate)
if err != nil {
return err
}
if closedTime.After(todayDateTime) {
return fmt.Errorf("分组用量汇总水位位于未来: %s", closedBefore)
}
if closedBefore == todayDate {
return nil
}
}
var earliest sql.NullTime
if err := scanSingleRow(ctx, r.sql, "SELECT MIN(created_at) FROM usage_logs", nil, &earliest); err != nil {
return fmt.Errorf("读取最早用量记录: %w", err)
}
retainedFrom := todayStart
if earliest.Valid {
retainedFrom = earliest.Time.UTC()
}
retainedDate := service.GroupUsageDate(retainedFrom)
retainedDateTime, err := service.ParseGroupUsageDate(retainedDate)
if err != nil {
return err
}
rebuildStartDate := retainedDate
if !timezoneChanged && closedTime.After(retainedDateTime) {
rebuildStartDate = closedBefore
}
rebuildStart, err := service.ParseGroupUsageDate(rebuildStartDate)
if err != nil {
return err
}
if _, err := r.sql.ExecContext(ctx, `
DELETE FROM usage_group_daily_rollups
WHERE bucket_date < $1::date
OR (bucket_date >= $2::date AND bucket_date < $3::date)
OR bucket_date >= $3::date
`, retainedDate, rebuildStartDate, todayDate); err != nil {
return fmt.Errorf("清理分组用量日桶: %w", err)
}
if _, err := r.sql.ExecContext(ctx, `
INSERT INTO usage_group_daily_rollups (bucket_date, group_id, actual_cost, computed_at)
SELECT
(created_at AT TIME ZONE $3::text)::date AS bucket_date,
group_id,
COALESCE(SUM(actual_cost), 0) AS actual_cost,
NOW()
FROM usage_logs
WHERE group_id IS NOT NULL
AND created_at >= $1
AND created_at < $2
GROUP BY 1, 2
ON CONFLICT (bucket_date, group_id)
DO UPDATE SET
actual_cost = EXCLUDED.actual_cost,
computed_at = EXCLUDED.computed_at
`, rebuildStart.UTC(), todayStart.UTC(), timezoneName); err != nil {
return fmt.Errorf("重建分组用量日桶: %w", err)
}
if _, err := r.sql.ExecContext(ctx, `
UPDATE usage_group_rollup_state
SET closed_before = $1::date,
retained_from = $2,
timezone_name = $3,
updated_at = NOW()
WHERE id = 1
`, todayDate, retainedFrom, timezoneName); err != nil {
return fmt.Errorf("更新分组用量汇总水位: %w", err)
}
return nil
}
func lockGroupUsageRollupState(ctx context.Context, tx *sql.Tx) error {
var id int16
if err := tx.QueryRowContext(ctx, `
SELECT id
FROM usage_group_rollup_state
WHERE id = 1
FOR UPDATE
`).Scan(&id); err != nil {
return fmt.Errorf("锁定分组用量汇总水位: %w", err)
}
return nil
}
func invalidateGroupUsageRollupsAt(ctx context.Context, tx *sql.Tx, affectedAt time.Time) error {
timezoneName := service.GroupUsageTimezoneName()
_, err := tx.ExecContext(ctx, `
UPDATE usage_group_rollup_state
SET closed_before = LEAST(
closed_before,
($1::timestamptz AT TIME ZONE $2::text)::date
),
updated_at = NOW()
WHERE id = 1
`, affectedAt.UTC(), timezoneName)
return err
}
@@ -0,0 +1,16 @@
package repository
import (
"testing"
appTimezone "github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
"github.com/stretchr/testify/require"
)
func useGroupUsageRepositoryTestTimezone(t *testing.T, name string) {
t.Helper()
previousName := appTimezone.Name()
require.NoError(t, appTimezone.Init(name))
t.Cleanup(func() { require.NoError(t, appTimezone.Init(previousName)) })
}
@@ -0,0 +1,319 @@
//go:build unit
package repository
import (
"context"
"database/sql"
"errors"
"testing"
"time"
"github.com/DATA-DOG/go-sqlmock"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
func TestDashboardAggregationRepositorySyncGroupUsageRollupsNoopsAtCurrentDate(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
mock.ExpectBegin()
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
AddRow("2026-08-14", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
mock.ExpectCommit()
require.NoError(t, repo.SyncGroupUsageRollups(context.Background(), todayStart))
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositorySyncGroupUsageRollupsRebuildsWhenTimezoneChanges(t *testing.T) {
useGroupUsageRepositoryTestTimezone(t, "America/New_York")
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
retainedFrom := time.Date(2026, 3, 1, 5, 0, 0, 0, time.UTC)
mock.ExpectBegin()
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
AddRow("2026-03-09", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
mock.ExpectQuery(`SELECT MIN\(created_at\) FROM usage_logs`).
WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(retainedFrom))
mock.ExpectExec(`DELETE FROM usage_group_daily_rollups`).
WithArgs("2026-03-01", "2026-03-01", "2026-03-09").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`INSERT INTO usage_group_daily_rollups`).
WithArgs(retainedFrom, todayStart, "America/New_York").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs("2026-03-09", retainedFrom, "America/New_York").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
require.NoError(t, repo.SyncGroupUsageRollups(context.Background(), todayStart))
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositorySyncGroupUsageRollupsPublishesWatermarkLast(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
retainedFrom := time.Date(2026, 5, 1, 3, 0, 0, 0, time.UTC)
rebuildStart := time.Date(2026, 8, 12, 16, 0, 0, 0, time.UTC)
mock.ExpectBegin()
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
AddRow("2026-08-13", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
mock.ExpectQuery(`SELECT MIN\(created_at\) FROM usage_logs`).
WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(retainedFrom))
mock.ExpectExec(`DELETE FROM usage_group_daily_rollups`).
WithArgs("2026-05-01", "2026-08-13", "2026-08-14").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`INSERT INTO usage_group_daily_rollups`).
WithArgs(rebuildStart, todayStart, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 2))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs("2026-08-14", retainedFrom, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
require.NoError(t, repo.SyncGroupUsageRollups(context.Background(), todayStart))
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositorySyncGroupUsageRollupsRejectsFutureWatermark(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
mock.ExpectBegin()
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
AddRow("2026-08-15", time.Unix(0, 0).UTC(), "Asia/Shanghai"))
mock.ExpectRollback()
err := repo.SyncGroupUsageRollups(context.Background(), todayStart)
require.ErrorContains(t, err, "未来")
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositoryRecomputeRangeInvalidatesGroupRollupsBeforeDashboardRebuild(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
start := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
end := start.Add(time.Hour)
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(start, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`DELETE FROM usage_dashboard_hourly`).
WillReturnError(sql.ErrConnDone)
mock.ExpectRollback()
err := repo.RecomputeRange(context.Background(), start, end)
require.ErrorIs(t, err, sql.ErrConnDone)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositoryRecomputeRangeRebuildsGroupRollupsBeforeCommit(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
start := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
end := start.Add(time.Hour)
fixedNow := time.Date(2026, 8, 14, 8, 0, 0, 0, time.UTC)
repo.clock = func() time.Time { return fixedNow }
todayStart := service.GroupUsageTodayStart(fixedNow)
startDate := service.GroupUsageDate(start)
rebuildStart, err := service.ParseGroupUsageDate(startDate)
require.NoError(t, err)
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(start, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
for _, query := range []string{
`DELETE FROM usage_dashboard_hourly WHERE`,
`DELETE FROM usage_dashboard_hourly_users WHERE`,
`DELETE FROM usage_dashboard_daily WHERE`,
`DELETE FROM usage_dashboard_daily_users WHERE`,
`INSERT INTO usage_dashboard_hourly_users`,
`INSERT INTO usage_dashboard_daily_users`,
`INSERT INTO usage_dashboard_hourly`,
`INSERT INTO usage_dashboard_daily`,
} {
mock.ExpectExec(query).WillReturnResult(sqlmock.NewResult(0, 1))
}
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
AddRow(startDate, time.Unix(0, 0).UTC(), "Asia/Shanghai"))
mock.ExpectQuery(`SELECT MIN\(created_at\) FROM usage_logs`).
WillReturnRows(sqlmock.NewRows([]string{"min"}).AddRow(start))
mock.ExpectExec(`DELETE FROM usage_group_daily_rollups`).
WithArgs(startDate, startDate, service.GroupUsageDate(todayStart)).
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`INSERT INTO usage_group_daily_rollups`).
WithArgs(rebuildStart.UTC(), todayStart.UTC(), "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(service.GroupUsageDate(todayStart), start, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
require.NoError(t, repo.RecomputeRange(context.Background(), start, end))
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositoryCleanupUsageLogsNonPartitionedInvalidatesEachBatchAndSyncs(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
cutoff := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC)
earliestDeletedAt := time.Date(2026, 5, 3, 2, 0, 0, 0, time.UTC)
fixedNow := time.Date(2026, 8, 14, 8, 0, 0, 0, time.UTC)
repo.clock = func() time.Time { return fixedNow }
todayStart := service.GroupUsageTodayStart(fixedNow)
mock.ExpectQuery(`SELECT EXISTS`).
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false))
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectQuery(`(?s)DELETE FROM usage_logs.*RETURNING created_at`).
WithArgs(cutoff, usageLogsCleanupBatchSize).
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).
AddRow(earliestDeletedAt.Add(time.Hour)).
AddRow(earliestDeletedAt))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(earliestDeletedAt, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
mock.ExpectBegin()
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
AddRow(service.GroupUsageDate(todayStart), time.Unix(0, 0).UTC(), "Asia/Shanghai"))
mock.ExpectCommit()
require.NoError(t, repo.CleanupUsageLogs(context.Background(), cutoff))
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositoryCleanupUsageLogsPartitionedSortsAndInvalidatesEachDropBeforeSync(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
cutoff := time.Date(2026, 7, 18, 0, 0, 0, 0, time.UTC)
aprilStart := time.Date(2026, 4, 1, 0, 0, 0, 0, time.UTC)
juneStart := time.Date(2026, 6, 1, 0, 0, 0, 0, time.UTC)
fixedNow := time.Date(2026, 8, 14, 8, 0, 0, 0, time.UTC)
repo.clock = func() time.Time { return fixedNow }
todayStart := service.GroupUsageTodayStart(fixedNow)
mock.ExpectQuery(`SELECT EXISTS`).
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
mock.ExpectQuery(`SELECT c.relname`).
WillReturnRows(sqlmock.NewRows([]string{"relname"}).
AddRow("usage_logs_202606").
AddRow("usage_logs_invalid").
AddRow("usage_logs_202604").
AddRow("usage_logs_202607"))
for _, partition := range []struct {
name string
start time.Time
}{
{name: "usage_logs_202604", start: aprilStart},
{name: "usage_logs_202606", start: juneStart},
} {
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(partition.start, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`DROP TABLE IF EXISTS "` + partition.name + `"`).
WillReturnResult(sqlmock.NewResult(0, 0))
mock.ExpectCommit()
}
mock.ExpectBegin()
mock.ExpectQuery(`SELECT closed_before::text, retained_from.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"closed_before", "retained_from", "timezone_name"}).
AddRow(service.GroupUsageDate(todayStart), time.Unix(0, 0).UTC(), "Asia/Shanghai"))
mock.ExpectCommit()
require.NoError(t, repo.CleanupUsageLogs(context.Background(), cutoff))
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositoryCleanupUsageLogsNonPartitionedFailureRollsBackWithoutSync(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
cutoff := time.Date(2026, 7, 1, 0, 0, 0, 0, time.UTC)
deletedAt := time.Date(2026, 5, 3, 2, 0, 0, 0, time.UTC)
mock.ExpectQuery(`SELECT EXISTS`).
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false))
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectQuery(`(?s)SELECT ctid.*ORDER BY created_at ASC, id ASC.*DELETE FROM usage_logs.*RETURNING created_at`).
WithArgs(cutoff, usageLogsCleanupBatchSize).
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(deletedAt))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(deletedAt, "Asia/Shanghai").
WillReturnError(sql.ErrConnDone)
mock.ExpectRollback()
err := repo.CleanupUsageLogs(context.Background(), cutoff)
require.ErrorIs(t, err, sql.ErrConnDone)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestDashboardAggregationRepositoryCleanupUsageLogsPartitionFailureRollsBackAndStops(t *testing.T) {
setGroupUsageRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := newDashboardAggregationRepositoryWithSQL(db)
cutoff := time.Date(2026, 7, 18, 0, 0, 0, 0, time.UTC)
aprilStart := time.Date(2026, 4, 1, 0, 0, 0, 0, time.UTC)
dropErr := errors.New("drop partition failed")
mock.ExpectQuery(`SELECT EXISTS`).
WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(true))
mock.ExpectQuery(`SELECT c.relname`).
WillReturnRows(sqlmock.NewRows([]string{"relname"}).
AddRow("usage_logs_202606").
AddRow("usage_logs_202604"))
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(aprilStart, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectExec(`DROP TABLE IF EXISTS "usage_logs_202604"`).
WillReturnError(dropErr)
mock.ExpectRollback()
err := repo.CleanupUsageLogs(context.Background(), cutoff)
require.ErrorIs(t, err, dropErr)
require.NoError(t, mock.ExpectationsWereMet())
}
func setGroupUsageRollupTestTimezone(t *testing.T) {
t.Helper()
useGroupUsageRepositoryTestTimezone(t, "Asia/Shanghai")
}
@@ -5,6 +5,7 @@ import (
"database/sql"
"fmt"
"log"
"sort"
"strings"
"time"
@@ -14,7 +15,8 @@ import (
)
type dashboardAggregationRepository struct {
sql sqlExecutor
sql sqlExecutor
clock func() time.Time
}
const usageLogsCleanupBatchSize = 10000
@@ -33,7 +35,14 @@ func NewDashboardAggregationRepository(sqlDB *sql.DB) service.DashboardAggregati
}
func newDashboardAggregationRepositoryWithSQL(sqlq sqlExecutor) *dashboardAggregationRepository {
return &dashboardAggregationRepository{sql: sqlq}
return &dashboardAggregationRepository{sql: sqlq, clock: time.Now}
}
func (r *dashboardAggregationRepository) now() time.Time {
if r.clock != nil {
return r.clock()
}
return time.Now()
}
func isPostgresDriver(db *sql.DB) bool {
@@ -128,11 +137,23 @@ func (r *dashboardAggregationRepository) RecomputeRange(ctx context.Context, sta
if err != nil {
return err
}
if err := lockGroupUsageRollupState(ctx, tx); err != nil {
_ = tx.Rollback()
return err
}
if err := invalidateGroupUsageRollupsAt(ctx, tx, start); err != nil {
_ = tx.Rollback()
return err
}
txRepo := newDashboardAggregationRepositoryWithSQL(tx)
if err := txRepo.recomputeRangeInTx(ctx, hourStart, hourEnd, dayStart, dayEnd); err != nil {
_ = tx.Rollback()
return err
}
if err := txRepo.syncGroupUsageRollupsInTx(ctx, service.GroupUsageTodayStart(r.now())); err != nil {
_ = tx.Rollback()
return err
}
return tx.Commit()
}
return r.recomputeRangeInTx(ctx, hourStart, hourEnd, dayStart, dayEnd)
@@ -215,14 +236,35 @@ func (r *dashboardAggregationRepository) CleanupUsageLogs(ctx context.Context, c
return err
}
if isPartitioned {
return r.dropUsageLogsPartitions(ctx, cutoff)
if err := r.dropUsageLogsPartitions(ctx, cutoff); err != nil {
return err
}
} else if err := r.cleanupUsageLogsBatches(ctx, cutoff); err != nil {
return err
}
return r.SyncGroupUsageRollups(ctx, service.GroupUsageTodayStart(r.now()))
}
func (r *dashboardAggregationRepository) cleanupUsageLogsBatches(ctx context.Context, cutoff time.Time) error {
db, transactional := r.sql.(*sql.DB)
for {
if transactional {
affected, err := cleanupUsageLogsBatchWithRollupInvalidation(ctx, db, cutoff)
if err != nil {
return err
}
if affected < usageLogsCleanupBatchSize {
return nil
}
continue
}
res, err := r.sql.ExecContext(ctx, `
WITH victims AS (
SELECT ctid
FROM usage_logs
WHERE created_at < $1
ORDER BY created_at ASC, id ASC
LIMIT $2
)
DELETE FROM usage_logs
@@ -241,6 +283,66 @@ func (r *dashboardAggregationRepository) CleanupUsageLogs(ctx context.Context, c
}
}
func cleanupUsageLogsBatchWithRollupInvalidation(ctx context.Context, db *sql.DB, cutoff time.Time) (int64, error) {
tx, err := db.BeginTx(ctx, nil)
if err != nil {
return 0, err
}
rollback := func(err error) (int64, error) {
_ = tx.Rollback()
return 0, err
}
if err := lockGroupUsageRollupState(ctx, tx); err != nil {
return rollback(err)
}
rows, err := tx.QueryContext(ctx, `
WITH victims AS (
SELECT ctid
FROM usage_logs
WHERE created_at < $1
ORDER BY created_at ASC, id ASC
LIMIT $2
)
DELETE FROM usage_logs
WHERE ctid IN (SELECT ctid FROM victims)
RETURNING created_at
`, cutoff.UTC(), usageLogsCleanupBatchSize)
if err != nil {
return rollback(err)
}
var affected int64
var earliestDeletedAt time.Time
for rows.Next() {
var deletedAt time.Time
if err := rows.Scan(&deletedAt); err != nil {
_ = rows.Close()
return rollback(err)
}
affected++
if earliestDeletedAt.IsZero() || deletedAt.Before(earliestDeletedAt) {
earliestDeletedAt = deletedAt
}
}
if err := rows.Err(); err != nil {
_ = rows.Close()
return rollback(err)
}
if err := rows.Close(); err != nil {
return rollback(err)
}
if affected > 0 {
if err := invalidateGroupUsageRollupsAt(ctx, tx, earliestDeletedAt); err != nil {
return rollback(err)
}
}
if err := tx.Commit(); err != nil {
return 0, err
}
return affected, nil
}
func (r *dashboardAggregationRepository) CleanupUsageBillingDedup(ctx context.Context, cutoff time.Time) error {
for {
res, err := r.sql.ExecContext(ctx, `
@@ -489,14 +591,16 @@ func (r *dashboardAggregationRepository) dropUsageLogsPartitions(ctx context.Con
if err != nil {
return err
}
defer func() {
_ = rows.Close()
}()
cutoffMonth := truncateToMonthUTC(cutoff)
type usageLogsPartition struct {
name string
month time.Time
}
partitions := make([]usageLogsPartition, 0)
for rows.Next() {
var name string
if err := rows.Scan(&name); err != nil {
_ = rows.Close()
return err
}
if !strings.HasPrefix(name, "usage_logs_") {
@@ -509,12 +613,56 @@ func (r *dashboardAggregationRepository) dropUsageLogsPartitions(ctx context.Con
}
month = month.UTC()
if month.Before(cutoffMonth) {
if _, err := r.sql.ExecContext(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", pq.QuoteIdentifier(name))); err != nil {
partitions = append(partitions, usageLogsPartition{name: name, month: month})
}
}
if err := rows.Err(); err != nil {
_ = rows.Close()
return err
}
if err := rows.Close(); err != nil {
return err
}
sort.Slice(partitions, func(i, j int) bool {
return partitions[i].month.Before(partitions[j].month)
})
if db, ok := r.sql.(*sql.DB); ok {
for _, partition := range partitions {
if err := dropUsageLogsPartitionWithRollupInvalidation(ctx, db, partition.name, partition.month); err != nil {
return err
}
}
return nil
}
return rows.Err()
for _, partition := range partitions {
if _, err := r.sql.ExecContext(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", pq.QuoteIdentifier(partition.name))); err != nil {
return err
}
}
return nil
}
func dropUsageLogsPartitionWithRollupInvalidation(ctx context.Context, db *sql.DB, name string, monthStart time.Time) error {
tx, err := db.BeginTx(ctx, nil)
if err != nil {
return err
}
rollback := func(err error) error {
_ = tx.Rollback()
return err
}
if err := lockGroupUsageRollupState(ctx, tx); err != nil {
return rollback(err)
}
if err := invalidateGroupUsageRollupsAt(ctx, tx, monthStart); err != nil {
return rollback(err)
}
if _, err := tx.ExecContext(ctx, fmt.Sprintf("DROP TABLE IF EXISTS %s", pq.QuoteIdentifier(name))); err != nil {
return rollback(err)
}
return tx.Commit()
}
func (r *dashboardAggregationRepository) createUsageLogsPartition(ctx context.Context, month time.Time) error {
+13
View File
@@ -3,6 +3,7 @@ package repository
import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"sort"
@@ -56,6 +57,10 @@ func createGroupRecord(ctx context.Context, client *dbent.Client, groupIn *servi
if groupIn == nil {
return errors.New("group is nil")
}
modelPricing, err := json.Marshal(groupIn.ModelPricing)
if err != nil {
return fmt.Errorf("marshal group model pricing: %w", err)
}
builder := client.Group.Create().
SetName(groupIn.Name).
SetDescription(groupIn.Description).
@@ -88,6 +93,8 @@ func createGroupRecord(ctx context.Context, client *dbent.Client, groupIn *servi
SetNillableAudioRealtimePricePerMin(groupIn.AudioRealtimePricePerMin).
SetNillableAudioTtsPricePerMillionChars(groupIn.AudioTTSPricePerMillionChars).
SetNillableAudioSttPricePerHour(groupIn.AudioSTTPricePerHour).
SetLongContextPricingEnabled(groupIn.LongContextPricingEnabled).
SetModelPricing(modelPricing).
SetDefaultValidityDays(groupIn.DefaultValidityDays).
SetClaudeCodeOnly(groupIn.ClaudeCodeOnly).
SetNillableFallbackGroupID(groupIn.FallbackGroupID).
@@ -234,6 +241,10 @@ func (r *groupRepository) GetByIDLite(ctx context.Context, id int64) (*service.G
}
func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) error {
modelPricing, err := json.Marshal(groupIn.ModelPricing)
if err != nil {
return fmt.Errorf("marshal group model pricing: %w", err)
}
builder := r.client.Group.UpdateOneID(groupIn.ID).
SetName(groupIn.Name).
SetDescription(groupIn.Description).
@@ -260,6 +271,8 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er
SetNillableVideoPrice720p(groupIn.VideoPrice720P).
SetNillableVideoPrice1080p(groupIn.VideoPrice1080P).
SetVideoModelPrices(service.NormalizeVideoModelPrices(groupIn.VideoModelPrices)).
SetLongContextPricingEnabled(groupIn.LongContextPricingEnabled).
SetModelPricing(modelPricing).
SetDefaultValidityDays(groupIn.DefaultValidityDays).
SetClaudeCodeOnly(groupIn.ClaudeCodeOnly).
SetModelRoutingEnabled(groupIn.ModelRoutingEnabled).
@@ -0,0 +1,511 @@
//go:build integration
package repository
import (
"context"
"database/sql"
"fmt"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/migrations"
"github.com/lib/pq"
"github.com/stretchr/testify/require"
)
func TestGroupUsageRollupTriggerInvalidatesCascadedHistoricalDelete(t *testing.T) {
for _, partitioned := range []bool{false, true} {
name := "ordinary"
if partitioned {
name = "partitioned"
}
t.Run(name, func(t *testing.T) {
ctx := context.Background()
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, partitioned)
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
defer func() { _ = tx.Rollback() }()
_, err := tx.ExecContext(ctx, `
INSERT INTO groups (id) VALUES (10);
INSERT INTO users (id) VALUES (1);
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
VALUES (1, 1, 10, 1.25, TIMESTAMPTZ '2020-01-02 08:00:00+08');
UPDATE usage_group_rollup_state
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
WHERE id = 1;
DELETE FROM users WHERE id = 1;
`)
require.NoError(t, err)
var closedBefore string
err = tx.QueryRowContext(ctx, `
SELECT closed_before::text
FROM usage_group_rollup_state
WHERE id = 1
`).Scan(&closedBefore)
require.NoError(t, err)
require.Equal(t, "2020-01-02", closedBefore)
})
}
}
func TestGroupUsageRollupTriggerSerializesLateHistoricalInsertWithPublish(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
seedTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
_, err := seedTx.ExecContext(ctx, `
INSERT INTO groups (id) VALUES (10);
INSERT INTO users (id) VALUES (1);
UPDATE usage_group_rollup_state
SET closed_before = DATE '2020-01-02'
WHERE id = 1;
`)
require.NoError(t, err)
require.NoError(t, seedTx.Commit())
syncTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
defer func() { _ = syncTx.Rollback() }()
var stateID int16
require.NoError(t, syncTx.QueryRowContext(ctx, `
SELECT id
FROM usage_group_rollup_state
WHERE id = 1
FOR UPDATE
`).Scan(&stateID))
lateTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
defer func() { _ = lateTx.Rollback() }()
var lateBackendPID int
require.NoError(t, lateTx.QueryRowContext(ctx, "SELECT pg_backend_pid()").Scan(&lateBackendPID))
insertResult := make(chan error, 1)
go func() {
_, insertErr := lateTx.ExecContext(ctx, `
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
VALUES (1, 1, 10, 1.25, TIMESTAMPTZ '2020-01-02 09:00:00+08')
`)
insertResult <- insertErr
}()
blocked, err := waitForGroupUsageRollupStateLock(ctx, lateBackendPID, insertResult)
if err != nil || !blocked {
_ = syncTx.Rollback()
_ = lateTx.Rollback()
require.NoError(t, err)
require.True(t, blocked, "迟到写入必须等待正在发布水位的事务")
}
_, err = syncTx.ExecContext(ctx, `
UPDATE usage_group_rollup_state
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
WHERE id = 1
`)
require.NoError(t, err)
require.NoError(t, syncTx.Commit())
select {
case err = <-insertResult:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal("等待迟到写入完成超时")
}
require.NoError(t, lateTx.Commit())
var closedBefore string
err = integrationDB.QueryRowContext(ctx, fmt.Sprintf(
"SELECT closed_before::text FROM %s.usage_group_rollup_state WHERE id = 1",
pq.QuoteIdentifier(schema),
)).Scan(&closedBefore)
require.NoError(t, err)
require.Equal(t, "2020-01-02", closedBefore)
}
func TestGroupUsageRollupTriggerSerializesInsertTransactionAcrossMidnight(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
seedTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
_, err := seedTx.ExecContext(ctx, `
INSERT INTO groups (id) VALUES (10);
INSERT INTO users (id) VALUES (1);
UPDATE usage_group_rollup_state
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
WHERE id = 1;
`)
require.NoError(t, err)
require.NoError(t, seedTx.Commit())
syncTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
defer func() { _ = syncTx.Rollback() }()
var stateID int16
require.NoError(t, syncTx.QueryRowContext(ctx, `
SELECT id
FROM usage_group_rollup_state
WHERE id = 1
FOR UPDATE
`).Scan(&stateID))
insertTx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
defer func() { _ = insertTx.Rollback() }()
var insertBackendPID int
require.NoError(t, insertTx.QueryRowContext(ctx, "SELECT pg_backend_pid()").Scan(&insertBackendPID))
insertResult := make(chan error, 1)
go func() {
_, insertErr := insertTx.ExecContext(ctx, `
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
VALUES (1, 1, 10, 1.25, CURRENT_TIMESTAMP)
`)
insertResult <- insertErr
}()
blocked, err := waitForGroupUsageRollupStateLock(ctx, insertBackendPID, insertResult)
if err != nil || !blocked {
_ = syncTx.Rollback()
_ = insertTx.Rollback()
require.NoError(t, err)
require.True(t, blocked, "跨越零点的在途写入必须与水位发布串行化")
}
_, err = syncTx.ExecContext(ctx, `
UPDATE usage_group_rollup_state
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date + 1
WHERE id = 1
`)
require.NoError(t, err)
require.NoError(t, syncTx.Commit())
select {
case err = <-insertResult:
require.NoError(t, err)
case <-ctx.Done():
t.Fatal("等待跨零点写入完成超时")
}
require.NoError(t, insertTx.Commit())
var currentDate string
require.NoError(t, integrationDB.QueryRowContext(ctx, `
SELECT (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date::text
`).Scan(&currentDate))
var closedBefore string
err = integrationDB.QueryRowContext(ctx, fmt.Sprintf(
"SELECT closed_before::text FROM %s.usage_group_rollup_state WHERE id = 1",
pq.QuoteIdentifier(schema),
)).Scan(&closedBefore)
require.NoError(t, err)
require.Equal(t, currentDate, closedBefore)
}
func TestGroupUsageRollupTriggerKeepsWatermarkForTodayInsert(t *testing.T) {
ctx := context.Background()
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
defer func() { _ = tx.Rollback() }()
_, err := tx.ExecContext(ctx, `
INSERT INTO groups (id) VALUES (10);
INSERT INTO users (id) VALUES (1);
UPDATE usage_group_rollup_state
SET closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
WHERE id = 1;
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
VALUES (1, 1, 10, 1.25, CURRENT_TIMESTAMP);
`)
require.NoError(t, err)
var unchanged bool
err = tx.QueryRowContext(ctx, `
SELECT closed_before = (CURRENT_TIMESTAMP AT TIME ZONE 'Asia/Shanghai')::date
FROM usage_group_rollup_state
WHERE id = 1
`).Scan(&unchanged)
require.NoError(t, err)
require.True(t, unchanged)
}
func TestGroupUsageRollupTriggerUsesSessionTimezoneAcrossDST(t *testing.T) {
ctx := context.Background()
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
defer func() { _ = tx.Rollback() }()
_, err := tx.ExecContext(ctx, `
SET LOCAL TIME ZONE 'America/New_York';
INSERT INTO groups (id) VALUES (10);
INSERT INTO users (id) VALUES (1);
UPDATE usage_group_rollup_state
SET closed_before = DATE '2026-03-09',
timezone_name = 'America/New_York'
WHERE id = 1;
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at)
VALUES (1, 1, 10, 1.25, TIMESTAMPTZ '2026-03-08 04:30:00+00');
`)
require.NoError(t, err)
var closedBefore string
err = tx.QueryRowContext(ctx, `
SELECT closed_before::text
FROM usage_group_rollup_state
WHERE id = 1
`).Scan(&closedBefore)
require.NoError(t, err)
require.Equal(t, "2026-03-07", closedBefore)
}
func TestGroupUsageSummaryIncludesYesterdayAcrossWatermark(t *testing.T) {
ctx := context.Background()
useGroupUsageRepositoryTestTimezone(t, "Asia/Shanghai")
todayStart := time.Date(2026, 8, 13, 16, 0, 0, 0, time.UTC)
tests := []struct {
name string
closedBefore string
includeYesterday bool
}{
{name: "closed_rollup", closedBefore: "2026-08-14", includeYesterday: true},
{name: "raw_tail", closedBefore: "2026-08-13", includeYesterday: false},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
defer func() { _ = tx.Rollback() }()
_, err := tx.ExecContext(ctx, `
INSERT INTO groups (id) VALUES (10);
INSERT INTO users (id) VALUES (1);
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at) VALUES
(1, 1, 10, 2, TIMESTAMPTZ '2026-08-12 12:00:00+08'),
(2, 1, 10, 3, TIMESTAMPTZ '2026-08-13 12:00:00+08'),
(3, 1, 10, 4, TIMESTAMPTZ '2026-08-14 12:00:00+08');
INSERT INTO usage_group_daily_rollups (bucket_date, group_id, actual_cost, computed_at)
VALUES (DATE '2026-08-12', 10, 2, NOW());
`)
require.NoError(t, err)
if tt.includeYesterday {
_, err = tx.ExecContext(ctx, `
INSERT INTO usage_group_daily_rollups (bucket_date, group_id, actual_cost, computed_at)
VALUES (DATE '2026-08-13', 10, 3, NOW())
`)
require.NoError(t, err)
}
_, err = tx.ExecContext(ctx, `
UPDATE usage_group_rollup_state
SET closed_before = $1::date,
retained_from = TIMESTAMPTZ '2026-08-12 00:00:00+08'
WHERE id = 1
`, tt.closedBefore)
require.NoError(t, err)
repo := newUsageLogRepositoryWithSQL(nil, tx)
result, err := repo.GetAllGroupUsageSummary(ctx, todayStart)
require.NoError(t, err)
require.Len(t, result, 1)
require.InDelta(t, 9, result[0].TotalCost, 0.0000001)
require.InDelta(t, 4, result[0].TodayCost, 0.0000001)
require.InDelta(t, 3, result[0].YesterdayCost, 0.0000001)
})
}
}
func TestGroupUsageRollupSyncRebuildsAfterTimezoneChange(t *testing.T) {
ctx := context.Background()
useGroupUsageRepositoryTestTimezone(t, "America/New_York")
todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
defer func() { _ = tx.Rollback() }()
_, err := tx.ExecContext(ctx, `
SET LOCAL TIME ZONE 'UTC';
INSERT INTO groups (id) VALUES (10);
INSERT INTO users (id) VALUES (1);
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at) VALUES
(1, 1, 10, 3, TIMESTAMPTZ '2026-03-08 05:30:00+00'),
(2, 1, 10, 5, TIMESTAMPTZ '2026-03-09 04:30:00+00');
INSERT INTO usage_group_daily_rollups (bucket_date, group_id, actual_cost, computed_at)
VALUES (DATE '2026-03-08', 10, 99, NOW());
UPDATE usage_group_rollup_state
SET closed_before = DATE '2026-03-09',
retained_from = TIMESTAMPTZ '2026-03-08 05:30:00+00',
timezone_name = 'Asia/Shanghai'
WHERE id = 1;
`)
require.NoError(t, err)
repo := newDashboardAggregationRepositoryWithSQL(tx)
require.NoError(t, repo.SyncGroupUsageRollups(ctx, todayStart))
var stateTimezone string
var closedBefore string
require.NoError(t, tx.QueryRowContext(ctx, `
SELECT timezone_name, closed_before::text
FROM usage_group_rollup_state
WHERE id = 1
`).Scan(&stateTimezone, &closedBefore))
require.Equal(t, "America/New_York", stateTimezone)
require.Equal(t, "2026-03-09", closedBefore)
var rollupCost float64
require.NoError(t, tx.QueryRowContext(ctx, `
SELECT actual_cost
FROM usage_group_daily_rollups
WHERE bucket_date = DATE '2026-03-08' AND group_id = 10
`).Scan(&rollupCost))
require.InDelta(t, 3, rollupCost, 0.0000001)
usageRepo := newUsageLogRepositoryWithSQL(nil, tx)
result, err := usageRepo.GetAllGroupUsageSummary(ctx, todayStart)
require.NoError(t, err)
require.Len(t, result, 1)
require.InDelta(t, 8, result[0].TotalCost, 0.0000001)
require.InDelta(t, 5, result[0].TodayCost, 0.0000001)
require.InDelta(t, 3, result[0].YesterdayCost, 0.0000001)
}
func TestGroupUsageSummaryUsesConfiguredDSTBoundaries(t *testing.T) {
ctx := context.Background()
useGroupUsageRepositoryTestTimezone(t, "America/New_York")
todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
schema := createGroupUsageRollupTriggerTestSchema(t, ctx, false)
tx := beginGroupUsageRollupTriggerTestTx(t, ctx, schema)
defer func() { _ = tx.Rollback() }()
_, err := tx.ExecContext(ctx, `
SET LOCAL TIME ZONE 'UTC';
INSERT INTO groups (id) VALUES (10);
INSERT INTO users (id) VALUES (1);
INSERT INTO usage_logs (id, user_id, group_id, actual_cost, created_at) VALUES
(1, 1, 10, 100, TIMESTAMPTZ '2026-03-08 04:30:00+00'),
(2, 1, 10, 3, TIMESTAMPTZ '2026-03-08 05:30:00+00'),
(3, 1, 10, 4, TIMESTAMPTZ '2026-03-09 03:30:00+00'),
(4, 1, 10, 5, TIMESTAMPTZ '2026-03-09 04:30:00+00');
UPDATE usage_group_rollup_state
SET closed_before = DATE '1970-01-01',
retained_from = TIMESTAMPTZ '1970-01-01 00:00:00+00',
timezone_name = 'America/New_York'
WHERE id = 1;
`)
require.NoError(t, err)
repo := newUsageLogRepositoryWithSQL(nil, tx)
result, err := repo.GetAllGroupUsageSummary(ctx, todayStart)
require.NoError(t, err)
require.Len(t, result, 1)
require.InDelta(t, 112, result[0].TotalCost, 0.0000001)
require.InDelta(t, 5, result[0].TodayCost, 0.0000001)
require.InDelta(t, 7, result[0].YesterdayCost, 0.0000001)
}
func createGroupUsageRollupTriggerTestSchema(t *testing.T, ctx context.Context, partitioned bool) string {
t.Helper()
schema := fmt.Sprintf("group_usage_rollup_trigger_%d", time.Now().UnixNano())
quotedSchema := pq.QuoteIdentifier(schema)
_, err := integrationDB.ExecContext(ctx, "CREATE SCHEMA "+quotedSchema)
require.NoError(t, err)
t.Cleanup(func() {
_, _ = integrationDB.ExecContext(context.Background(), "DROP SCHEMA IF EXISTS "+quotedSchema+" CASCADE")
})
tx, err := integrationDB.BeginTx(ctx, nil)
require.NoError(t, err)
defer func() { _ = tx.Rollback() }()
require.NoError(t, setGroupUsageRollupTriggerSearchPath(ctx, tx, quotedSchema))
usageLogsDDL := `
CREATE TABLE usage_logs (
id BIGINT PRIMARY KEY,
user_id BIGINT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
group_id BIGINT REFERENCES groups(id) ON DELETE SET NULL,
actual_cost NUMERIC(20, 10) NOT NULL,
created_at TIMESTAMPTZ NOT NULL
);
`
if partitioned {
usageLogsDDL = `
CREATE TABLE usage_logs (
id BIGINT NOT NULL,
user_id BIGINT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
group_id BIGINT REFERENCES groups(id) ON DELETE SET NULL,
actual_cost NUMERIC(20, 10) NOT NULL,
created_at TIMESTAMPTZ NOT NULL
) PARTITION BY RANGE (created_at);
CREATE TABLE usage_logs_default PARTITION OF usage_logs DEFAULT;
`
}
_, err = tx.ExecContext(ctx, `
CREATE TABLE users (id BIGINT PRIMARY KEY);
CREATE TABLE groups (id BIGINT PRIMARY KEY);
`+usageLogsDDL)
require.NoError(t, err)
for _, migrationName := range []string{
"222_group_usage_daily_rollups.sql",
"223_group_usage_rollup_timezone.sql",
} {
migrationSQL, readErr := migrations.FS.ReadFile(migrationName)
require.NoError(t, readErr)
for range 2 {
_, err = tx.ExecContext(ctx, string(migrationSQL))
require.NoError(t, err)
}
}
require.NoError(t, tx.Commit())
return schema
}
func beginGroupUsageRollupTriggerTestTx(t *testing.T, ctx context.Context, schema string) *sql.Tx {
t.Helper()
tx, err := integrationDB.BeginTx(ctx, nil)
require.NoError(t, err)
require.NoError(t, setGroupUsageRollupTriggerSearchPath(ctx, tx, pq.QuoteIdentifier(schema)))
return tx
}
func setGroupUsageRollupTriggerSearchPath(ctx context.Context, tx *sql.Tx, quotedSchema string) error {
_, err := tx.ExecContext(ctx, "SET LOCAL search_path TO "+quotedSchema)
return err
}
func waitForGroupUsageRollupStateLock(
ctx context.Context,
backendPID int,
insertResult <-chan error,
) (bool, error) {
ticker := time.NewTicker(10 * time.Millisecond)
defer ticker.Stop()
for {
select {
case err := <-insertResult:
if err != nil {
return false, err
}
return false, nil
case <-ticker.C:
var waitEventType sql.NullString
err := integrationDB.QueryRowContext(ctx, `
SELECT wait_event_type
FROM pg_stat_activity
WHERE pid = $1
`, backendPID).Scan(&waitEventType)
if err != nil {
return false, err
}
if waitEventType.Valid && waitEventType.String == "Lock" {
return true, nil
}
case <-ctx.Done():
return false, ctx.Err()
}
}
}
@@ -291,6 +291,9 @@ func (r *usageCleanupRepository) DeleteUsageLogsBatch(ctx context.Context, filte
return 0, fmt.Errorf("cleanup filters missing time range")
}
args = append(args, limit)
if db, ok := r.sql.(*sql.DB); ok {
return r.deleteUsageLogsBatchWithRollupInvalidation(ctx, db, whereClause, args)
}
query := fmt.Sprintf(`
WITH target AS (
SELECT id
@@ -301,7 +304,7 @@ func (r *usageCleanupRepository) DeleteUsageLogsBatch(ctx context.Context, filte
)
DELETE FROM usage_logs
WHERE id IN (SELECT id FROM target)
RETURNING id
RETURNING created_at
`, whereClause, len(args))
rows, err := r.sql.QueryContext(ctx, query, args...)
@@ -320,6 +323,68 @@ func (r *usageCleanupRepository) DeleteUsageLogsBatch(ctx context.Context, filte
return deleted, nil
}
func (r *usageCleanupRepository) deleteUsageLogsBatchWithRollupInvalidation(ctx context.Context, db *sql.DB, whereClause string, args []any) (int64, error) {
tx, err := db.BeginTx(ctx, nil)
if err != nil {
return 0, err
}
rollback := func(err error) (int64, error) {
_ = tx.Rollback()
return 0, err
}
if err := lockGroupUsageRollupState(ctx, tx); err != nil {
return rollback(err)
}
query := fmt.Sprintf(`
WITH target AS (
SELECT id
FROM usage_logs
WHERE %s
ORDER BY created_at ASC, id ASC
LIMIT $%d
)
DELETE FROM usage_logs
WHERE id IN (SELECT id FROM target)
RETURNING created_at
`, whereClause, len(args))
rows, err := tx.QueryContext(ctx, query, args...)
if err != nil {
return rollback(err)
}
var deleted int64
var earliestDeletedAt time.Time
for rows.Next() {
var deletedAt time.Time
if err := rows.Scan(&deletedAt); err != nil {
_ = rows.Close()
return rollback(err)
}
deleted++
if earliestDeletedAt.IsZero() || deletedAt.Before(earliestDeletedAt) {
earliestDeletedAt = deletedAt
}
}
if err := rows.Err(); err != nil {
_ = rows.Close()
return rollback(err)
}
if err := rows.Close(); err != nil {
return rollback(err)
}
if deleted > 0 {
if err := invalidateGroupUsageRollupsAt(ctx, tx, earliestDeletedAt); err != nil {
return rollback(err)
}
}
if err := tx.Commit(); err != nil {
return 0, err
}
return deleted, nil
}
func buildUsageCleanupWhere(filters service.UsageCleanupFilters) (string, []any) {
conditions := make([]string, 0, 8)
args := make([]any, 0, 8)
@@ -21,6 +21,11 @@ func newSQLMock(t *testing.T) (*sql.DB, sqlmock.Sqlmock) {
return db, mock
}
func setUsageCleanupRollupTestTimezone(t *testing.T) {
t.Helper()
useGroupUsageRepositoryTestTimezone(t, "Asia/Shanghai")
}
func TestNewUsageCleanupRepository(t *testing.T) {
db, _ := newSQLMock(t)
repo := NewUsageCleanupRepository(nil, db)
@@ -398,6 +403,7 @@ func TestUsageCleanupRepositoryDeleteUsageLogsBatchMissingRange(t *testing.T) {
}
func TestUsageCleanupRepositoryDeleteUsageLogsBatch(t *testing.T) {
setUsageCleanupRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := &usageCleanupRepository{sql: db}
@@ -412,9 +418,16 @@ func TestUsageCleanupRepositoryDeleteUsageLogsBatch(t *testing.T) {
Model: &model,
}
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectQuery("DELETE FROM usage_logs").
WithArgs(start, end, userID, "gpt-4", 2).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(int64(1)).AddRow(int64(2)))
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(start.Add(time.Hour)).AddRow(start.Add(2 * time.Hour)))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(start.Add(time.Hour), "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
deleted, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 2)
require.NoError(t, err)
@@ -422,6 +435,62 @@ func TestUsageCleanupRepositoryDeleteUsageLogsBatch(t *testing.T) {
require.NoError(t, mock.ExpectationsWereMet())
}
func TestUsageCleanupRepositoryDeleteUsageLogsBatchAtomicallyInvalidatesGroupRollups(t *testing.T) {
setUsageCleanupRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := &usageCleanupRepository{sql: db}
start := time.Date(2026, 8, 1, 0, 0, 0, 0, time.UTC)
end := start.Add(24 * time.Hour)
firstDeletedAt := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
secondDeletedAt := firstDeletedAt.Add(time.Hour)
filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectQuery(`(?s)DELETE FROM usage_logs.*RETURNING created_at`).
WithArgs(start, end, 2).
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).
AddRow(firstDeletedAt).
AddRow(secondDeletedAt))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(firstDeletedAt, "Asia/Shanghai").
WillReturnResult(sqlmock.NewResult(0, 1))
mock.ExpectCommit()
deleted, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 2)
require.NoError(t, err)
require.Equal(t, int64(2), deleted)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestUsageCleanupRepositoryDeleteUsageLogsBatchRollsBackWhenInvalidationFails(t *testing.T) {
setUsageCleanupRollupTestTimezone(t)
db, mock := newSQLMock(t)
repo := &usageCleanupRepository{sql: db}
start := time.Date(2026, 8, 1, 0, 0, 0, 0, time.UTC)
end := start.Add(24 * time.Hour)
deletedAt := time.Date(2026, 8, 1, 3, 0, 0, 0, time.UTC)
filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectQuery(`(?s)DELETE FROM usage_logs.*RETURNING created_at`).
WithArgs(start, end, 1).
WillReturnRows(sqlmock.NewRows([]string{"created_at"}).AddRow(deletedAt))
mock.ExpectExec(`UPDATE usage_group_rollup_state`).
WithArgs(deletedAt, "Asia/Shanghai").
WillReturnError(sql.ErrConnDone)
mock.ExpectRollback()
_, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 1)
require.ErrorIs(t, err, sql.ErrConnDone)
require.NoError(t, mock.ExpectationsWereMet())
}
func TestUsageCleanupRepositoryDeleteUsageLogsBatchQueryError(t *testing.T) {
db, mock := newSQLMock(t)
repo := &usageCleanupRepository{sql: db}
@@ -430,9 +499,13 @@ func TestUsageCleanupRepositoryDeleteUsageLogsBatchQueryError(t *testing.T) {
end := start.Add(24 * time.Hour)
filters := service.UsageCleanupFilters{StartTime: start, EndTime: end}
mock.ExpectBegin()
mock.ExpectQuery(`SELECT id FROM usage_group_rollup_state.*FOR UPDATE`).
WillReturnRows(sqlmock.NewRows([]string{"id"}).AddRow(1))
mock.ExpectQuery("DELETE FROM usage_logs").
WithArgs(start, end, 5).
WillReturnError(sql.ErrConnDone)
mock.ExpectRollback()
_, err := repo.DeleteUsageLogsBatch(context.Background(), filters, 5)
require.Error(t, err)
@@ -0,0 +1,33 @@
//go:build unit
package repository
import (
"context"
"testing"
"time"
"github.com/DATA-DOG/go-sqlmock"
"github.com/stretchr/testify/require"
)
func TestUsageLogRepositoryGetAllGroupUsageSummaryUsesRollupTail(t *testing.T) {
db, mock := newSQLMock(t)
repo := newUsageLogRepositoryWithSQL(nil, db)
useGroupUsageRepositoryTestTimezone(t, "America/New_York")
todayStart := time.Date(2026, 3, 9, 4, 0, 0, 0, time.UTC)
yesterdayStart := time.Date(2026, 3, 8, 5, 0, 0, 0, time.UTC)
mock.ExpectQuery(`(?s)usage_group_rollup_state.*usage_group_daily_rollups.*created_at >= state\.tail_start`).
WithArgs(todayStart, yesterdayStart, "America/New_York", "2026-03-09", "2026-03-08").
WillReturnRows(sqlmock.NewRows([]string{"group_id", "total_cost", "today_cost", "yesterday_cost"}).
AddRow(int64(7), 12.5, 1.25, 2.5))
result, err := repo.GetAllGroupUsageSummary(context.Background(), todayStart)
require.NoError(t, err)
require.Equal(t, int64(7), result[0].GroupID)
require.InDelta(t, 12.5, result[0].TotalCost, 0.0000001)
require.InDelta(t, 1.25, result[0].TodayCost, 0.0000001)
require.InDelta(t, 2.5, result[0].YesterdayCost, 0.0000001)
require.NoError(t, mock.ExpectationsWereMet())
}
@@ -715,39 +715,9 @@ func (r *usageLogRepository) GetUserBreakdownStats(ctx context.Context, startTim
return results, nil
}
// GetAllGroupUsageSummary returns today's and cumulative actual_cost for every group.
// todayStart is the start-of-day in the caller's timezone (UTC-based).
// TODO(perf): This query scans ALL usage_logs rows for total_cost aggregation.
// When usage_logs exceeds ~1M rows, consider adding a short-lived cache (30s)
// or a materialized view / pre-aggregation table for cumulative costs.
// GetAllGroupUsageSummary 返回所有分组在服务端配置时区内的今日、昨日与当前保留记录累计金额。
func (r *usageLogRepository) GetAllGroupUsageSummary(ctx context.Context, todayStart time.Time) ([]usagestats.GroupUsageSummary, error) {
query := `
SELECT
g.id AS group_id,
COALESCE(SUM(ul.actual_cost), 0) AS total_cost,
COALESCE(SUM(CASE WHEN ul.created_at >= $1 THEN ul.actual_cost ELSE 0 END), 0) AS today_cost
FROM groups g
LEFT JOIN usage_logs ul ON ul.group_id = g.id
GROUP BY g.id
`
rows, err := r.sql.QueryContext(ctx, query, todayStart)
if err != nil {
return nil, err
}
defer func() { _ = rows.Close() }()
var results []usagestats.GroupUsageSummary
for rows.Next() {
var row usagestats.GroupUsageSummary
if err := rows.Scan(&row.GroupID, &row.TotalCost, &row.TodayCost); err != nil {
return nil, err
}
results = append(results, row)
}
if err := rows.Err(); err != nil {
return nil, err
}
return results, nil
return r.getAllGroupUsageSummaryFromRollups(ctx, todayStart)
}
// resolveModelDimensionExpression maps model source type to a safe SQL expression.
@@ -98,6 +98,35 @@ func TestUserPlatformQuotaRepository_BulkInsertInitial_GrokAllowed(t *testing.T)
require.InDelta(t, 9.0, *rec.DailyLimitUSD, 1e-9)
}
// TestUserPlatformQuotaRepository_BulkInsertInitial_CNProvidersAllowed 回归迁移 224:
// kimi/zhipu/deepseek 平台必须能写入 user_platform_quotas(CHECK 约束已含国产供应商)。
// 历史 bug:三个平台不在约束内 → 注册预填充 8 平台默认配额时整条多行 INSERT 中止 →
// fail-open 吞错 → 新用户拿到零条配额记录(缺失配额行 = 无限额)。
func TestUserPlatformQuotaRepository_BulkInsertInitial_CNProvidersAllowed(t *testing.T) {
ctx := context.Background()
tx := testEntTx(t)
txCtx := dbent.NewTxContext(ctx, tx)
client := tx.Client()
userID := mustCreateUserForQuota(t, client)
repo := NewUserPlatformQuotaRepository(client)
daily := 12.0
records := []UserPlatformQuotaRecord{
{UserID: userID, Platform: "kimi", DailyLimitUSD: &daily},
{UserID: userID, Platform: "zhipu"},
{UserID: userID, Platform: "deepseek"},
}
require.NoError(t, repo.BulkInsertInitial(txCtx, records),
"kimi/zhipu/deepseek 平台应可写入(迁移 224 后 CHECK 约束已含国产供应商)")
for _, platform := range []string{"kimi", "zhipu", "deepseek"} {
rec, err := repo.GetByUserPlatform(txCtx, userID, platform)
require.NoError(t, err)
require.NotNil(t, rec, "%s 配额行应已写入", platform)
}
}
func TestUserPlatformQuotaRepository_GetByUserPlatform(t *testing.T) {
ctx := context.Background()
tx := testEntTx(t)
+4 -3
View File
@@ -364,6 +364,7 @@ func TestAPIContracts(t *testing.T) {
"daily_limit_usd": null,
"weekly_limit_usd": null,
"monthly_limit_usd": null,
"long_context_pricing_enabled": false,
"image_price_1k": null,
"image_price_2k": null,
"image_price_4k": null,
@@ -860,7 +861,7 @@ func TestAPIContracts(t *testing.T) {
"force_email_on_third_party_signup": false,
"default_concurrency": 5,
"default_balance": 1.25,
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null}},
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"deepseek":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"kimi":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null},"zhipu":{"daily":null,"weekly":null,"monthly":null}},
"auth_source_default_email_platform_quotas": null,
"auth_source_default_github_platform_quotas": null,
"auth_source_default_google_platform_quotas": null,
@@ -1174,7 +1175,7 @@ func TestAPIContracts(t *testing.T) {
"purchase_subscription_url": "",
"table_default_page_size": 20,
"table_page_size_options": [10, 20, 50],
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null}},
"default_platform_quotas": {"anthropic":{"daily":null,"weekly":null,"monthly":null},"antigravity":{"daily":null,"weekly":null,"monthly":null},"deepseek":{"daily":null,"weekly":null,"monthly":null},"gemini":{"daily":null,"weekly":null,"monthly":null},"grok":{"daily":null,"weekly":null,"monthly":null},"kimi":{"daily":null,"weekly":null,"monthly":null},"openai":{"daily":null,"weekly":null,"monthly":null},"zhipu":{"daily":null,"weekly":null,"monthly":null}},
"auth_source_default_email_platform_quotas": null,
"auth_source_default_github_platform_quotas": null,
"auth_source_default_google_platform_quotas": null,
@@ -1467,7 +1468,7 @@ func newContractDeps(t *testing.T) *contractDeps {
settingRepo := newStubSettingRepo()
settingService := service.NewSettingService(settingRepo, cfg)
adminService := service.NewAdminService(userRepo, groupRepo, &accountRepo, proxyRepo, apiKeyRepo, redeemRepo, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
adminService := service.NewAdminService(userRepo, groupRepo, &accountRepo, proxyRepo, apiKeyRepo, redeemRepo, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
authHandler := handler.NewAuthHandler(cfg, nil, userService, settingService, nil, redeemService, nil, nil)
apiKeyHandler := handler.NewAPIKeyHandler(apiKeyService)
usageHandler := handler.NewUsageHandler(usageService, apiKeyService, nil, nil)
+14
View File
@@ -58,6 +58,9 @@ func RegisterAdminRoutes(
// Grok OAuth
registerGrokOAuthRoutes(admin, h)
// 国产供应商(kimi/zhipu/deepseek)额度与余额
registerCNProviderRoutes(admin, h)
// 代理管理
registerProxyRoutes(admin, h, stepUpAuth)
@@ -482,6 +485,17 @@ func registerGrokOAuthRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
}
}
// registerCNProviderRoutes 注册国产供应商(kimi/zhipu/deepseek)的额度与余额查询端点。
func registerCNProviderRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
cn := admin.Group("/cn-providers")
{
// Coding Plan 滚动窗口用量(kimi/zhipu coding 账号)。
cn.GET("/accounts/:id/quota", h.Admin.CNProvider.QueryQuota)
// payg 账号余额(kimi/deepseek;zhipu 无余额端点)。
cn.GET("/accounts/:id/balance", h.Admin.CNProvider.QueryBalance)
}
}
func registerProxyRoutes(admin *gin.RouterGroup, h *handler.Handlers, stepUpAuth middleware.StepUpAuthMiddleware) {
proxies := admin.Group("/proxies")
{
+20 -2
View File
@@ -47,7 +47,9 @@ func RegisterGatewayRoutes(
isOpenAIResponsesCompatibleGatewayPlatform := func(c *gin.Context) bool {
switch getGroupPlatform(c) {
case service.PlatformOpenAI, service.PlatformGrok:
case service.PlatformOpenAI, service.PlatformGrok,
service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek:
// 国产 OpenAI 兼容供应商(kimi/zhipu/deepseek)与 openai/grok 一样经 OpenAI 网关转发。
return true
default:
return false
@@ -58,7 +60,7 @@ func RegisterGatewayRoutes(
}
countTokensHandler := func(c *gin.Context) {
switch getGroupPlatform(c) {
case service.PlatformOpenAI:
case service.PlatformOpenAI, service.PlatformKimi, service.PlatformZhipu, service.PlatformDeepseek:
h.OpenAIGateway.CountTokens(c)
case service.PlatformGrok:
h.OpenAIGateway.GrokCountTokens(c)
@@ -315,6 +317,14 @@ func RegisterGatewayRoutes(
}
h.Gateway.WebSearch(c)
})
gateway.POST("/x_search", func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "X Search API is not supported for this platform"}})
return
}
h.Gateway.XSearch(c)
})
}
// Gemini 原生 API 兼容层(Gemini SDK/CLI 直连)
@@ -443,6 +453,14 @@ func RegisterGatewayRoutes(
}
h.Gateway.WebSearch(c)
})
r.POST("/x_search", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
if getGroupPlatform(c) != service.PlatformGrok {
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "X Search API is not supported for this platform"}})
return
}
h.Gateway.XSearch(c)
})
// Antigravity 模型列表
r.GET("/antigravity/models", gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, h.Gateway.AntigravityModels)
@@ -48,6 +48,7 @@ func TestEveryGatewayPOSTRouteIsClassifiedForPromptAuditCoverage(t *testing.T) {
"/models/*modelAction": {"gemini_v1beta_handler.go"},
"/tts": {"grok_audio.go"},
"/web_search": {"gateway_web_search.go"},
"/x_search": {"gateway_web_search.go"},
}
excluded := map[string]string{
"/messages/count_tokens": "tokenization only; it does not execute a model request",
+189 -6
View File
@@ -271,8 +271,30 @@ func (a *Account) IsGrokOAuth() bool {
return a.IsGrok() && a.Type == AccountTypeOAuth
}
// IsKimi / IsZhipu / IsDeepseek 标识国产 OpenAI 兼容供应商账号。
func (a *Account) IsKimi() bool {
return a.Platform == PlatformKimi
}
func (a *Account) IsZhipu() bool {
return a.Platform == PlatformZhipu
}
func (a *Account) IsDeepseek() bool {
return a.Platform == PlatformDeepseek
}
// IsCNProvider 报告是否为国产 OpenAI 兼容供应商(kimi/zhipu/deepseek)。
func (a *Account) IsCNProvider() bool {
return a != nil && IsCNProvider(a.Platform)
}
// IsOpenAICompatible 报告账号是否走 OpenAI 网关(OpenAI 协议族)。
// openai/grok 原生走 OpenAI 网关;kimi/zhipu/deepseek 同为 OpenAI Chat Completions
// 兼容上游,也经 OpenAI 网关转发。
func (a *Account) IsOpenAICompatible() bool {
return a != nil && (a.Platform == PlatformOpenAI || a.Platform == PlatformGrok)
return a != nil && (a.Platform == PlatformOpenAI || a.Platform == PlatformGrok ||
a.Platform == PlatformKimi || a.Platform == PlatformZhipu || a.Platform == PlatformDeepseek)
}
func (a *Account) GeminiOAuthType() string {
@@ -1281,17 +1303,161 @@ func (a *Account) IsOpenAIApiKey() bool {
return a.IsOpenAI() && a.Type == AccountTypeAPIKey
}
// GetOpenAIBaseURL 解析 OpenAI 协议族账号的上游 base_url。
// 适用 openai 与国产 OpenAI 兼容供应商(kimi/zhipu/deepseek);grok 走 GetGrokBaseURL,
// 此处对 grok 返回 "" 以保持原有行为。
func (a *Account) GetOpenAIBaseURL() string {
if !a.IsOpenAI() {
if !a.IsOpenAI() && !a.IsCNProvider() {
return ""
}
if a.Type == AccountTypeAPIKey {
baseURL := a.GetCredential("base_url")
if baseURL != "" {
if a.Type == AccountTypeAPIKey || a.Type == AccountTypeUpstream {
if baseURL := strings.TrimSpace(a.GetCredential("base_url")); baseURL != "" {
return baseURL
}
}
return "https://api.openai.com"
// 平台默认 base_url:CN 供应商按 account_mode 选择 payg / coding 默认值。
switch a.Platform {
case PlatformKimi:
if a.GetAccountMode() == AccountModeCoding {
return DefaultKimiCodingBaseURL
}
return DefaultKimiPayGBaseURL
case PlatformZhipu:
if a.GetAccountMode() == AccountModeCoding {
return DefaultZhipuCodingBaseURL
}
return DefaultZhipuPayGBaseURL
case PlatformDeepseek:
return DefaultDeepseekBaseURL
default:
return "https://api.openai.com"
}
}
// GetAccountMode 返回国产供应商账号的接入模式(payg / coding);非国产供应商或未设置时
// 返回空串。存储于 credentials["account_mode"]。
func (a *Account) GetAccountMode() string {
if a == nil {
return ""
}
mode := strings.TrimSpace(a.GetCredential("account_mode"))
if mode == AccountModePayG || mode == AccountModeCoding {
return mode
}
return ""
}
// IsCodingPlan 报告账号是否为 Coding Plan 模式(用于滚动用量窗口冷却)。
func (a *Account) IsCodingPlan() bool {
return a.GetAccountMode() == AccountModeCoding
}
// GetAPIProtocol 返回国产供应商账号的上游 API 协议。存储于
// credentials["api_protocol"];缺失或与平台不匹配时回退 chat_completions
// (与既有行为完全一致)。responses 协议仅 deepseek 支持(官方原生 /responses
// 端点,适配 Codex);kimi/zhipu 无此端点。
func (a *Account) GetAPIProtocol() string {
if a == nil || !a.IsCNProvider() {
return APIProtocolChatCompletions
}
switch strings.TrimSpace(a.GetCredential("api_protocol")) {
case APIProtocolAnthropic:
return APIProtocolAnthropic
case APIProtocolResponses:
if a.Platform == PlatformDeepseek {
return APIProtocolResponses
}
case APIProtocolChatCompletions:
return APIProtocolChatCompletions
}
return APIProtocolChatCompletions
}
// IsAnthropicProtocol 报告账号是否以原生 Anthropic 协议接入上游
// (/v1/messages 直通,适配 Claude Code 等客户端)。
func (a *Account) IsAnthropicProtocol() bool {
return a.GetAPIProtocol() == APIProtocolAnthropic
}
// GetAnthropicProtocolBaseURL 返回 Anthropic 协议账号的上游 base_url
// (上游路径为 {base}/v1/messages)。优先取凭证 base_url,缺失时按
// 供应商 × 接入模式返回默认端点。非 Anthropic 协议账号返回空串。
func (a *Account) GetAnthropicProtocolBaseURL() string {
if a == nil || !a.IsAnthropicProtocol() {
return ""
}
if a.Type == AccountTypeAPIKey || a.Type == AccountTypeUpstream {
if baseURL := strings.TrimSpace(a.GetCredential("base_url")); baseURL != "" {
return baseURL
}
}
switch a.Platform {
case PlatformKimi:
if a.GetAccountMode() == AccountModeCoding {
return DefaultKimiCodingAnthropicBaseURL
}
return DefaultKimiPayGAnthropicBaseURL
case PlatformZhipu:
return DefaultZhipuAnthropicBaseURL
case PlatformDeepseek:
return DefaultDeepseekAnthropicBaseURL
default:
return ""
}
}
// GetOpenAIFormatBaseURL 返回供 OpenAI 格式端点(/v1/models、/v1/chat/completions
// 等)使用的 base。chat_completions / responses 协议下与 GetOpenAIBaseURL
// 一致(凭证 base_url 或平台默认);anthropic 协议下凭证 base_url 指向 Anthropic
// 端点,不能拿来拼 OpenAI 路径,此时返回该供应商 × 模式的 Chat Completions
// 默认 base(模型同步等协议族共用路径仍可用)。
func (a *Account) GetOpenAIFormatBaseURL() string {
if a == nil || !a.IsAnthropicProtocol() {
return a.GetOpenAIBaseURL()
}
switch a.Platform {
case PlatformKimi:
if a.GetAccountMode() == AccountModeCoding {
return DefaultKimiCodingBaseURL
}
return DefaultKimiPayGBaseURL
case PlatformZhipu:
if a.GetAccountMode() == AccountModeCoding {
return DefaultZhipuCodingBaseURL
}
return DefaultZhipuPayGBaseURL
case PlatformDeepseek:
return DefaultDeepseekBaseURL
default:
return a.GetOpenAIBaseURL()
}
}
// GetCNAPIKey 返回国产 OpenAI 兼容供应商账号的 api_key 凭据(kimi/zhipu/deepseek)。
// 与 openai 的 GetOpenAIApiKey 区分:后者仅对 openai 平台返回。
func (a *Account) GetCNAPIKey() string {
if a == nil || !a.IsCNProvider() {
return ""
}
return a.GetCredential("api_key")
}
// GetCodingPlanProvider 根据 base_url 识别 Coding Plan 供应商(kimi / zhipu),
// 用于路由到对应的额度查询端点。非 coding 模式或无法识别时返回空串。
// 判定规则与 cc-switch coding_plan.rs::detect_provider 保持一致。
func (a *Account) GetCodingPlanProvider() string {
if a == nil || a.GetAccountMode() != AccountModeCoding {
return ""
}
baseURL := strings.ToLower(a.GetOpenAIBaseURL())
switch {
case strings.Contains(baseURL, "api.kimi.com/coding"):
return PlatformKimi
case strings.Contains(baseURL, "bigmodel.cn"), strings.Contains(baseURL, "api.z.ai"):
return PlatformZhipu
default:
return ""
}
}
func (a *Account) GetOpenAIAccessToken() string {
@@ -1404,6 +1570,23 @@ func (a *Account) GetOpenAIApiKey() string {
return a.GetCredential("api_key")
}
// GetOpenAIProtocolAPIKey 返回 OpenAI 协议族 APIKey 账号的密钥。
// 覆盖 openai 原生账号与国产 OpenAI 兼容供应商(kimi/zhipu/deepseek)账号,
// 供转发鉴权、模型列表同步等协议族共用路径使用。注意 IsOpenAIApiKey 语义上
// 仅指 openai 平台账号,调度倍率/WS 能力门控继续以其为准,不受本方法影响。
func (a *Account) GetOpenAIProtocolAPIKey() string {
if a == nil {
return ""
}
if a.IsCNProvider() {
if a.Type != AccountTypeAPIKey {
return ""
}
return a.GetCredential("api_key")
}
return a.GetOpenAIApiKey()
}
func (a *Account) GetOpenAIUserAgent() string {
if !a.IsOpenAI() {
return ""
@@ -58,6 +58,10 @@ func EvaluateAccountSchedulingThreshold(account *Account, thresholds map[string]
winner = pickLatestResetSchedulingCandidate(anthropicThresholdCandidates(account), threshold, now)
case PlatformGrok:
winner = pickLatestResetSchedulingCandidate(grokThresholdCandidates(account), threshold, now)
case PlatformKimi:
winner = pickLatestResetSchedulingCandidate(cnProviderThresholdCandidates(account, PlatformKimi), threshold, now)
case PlatformZhipu:
winner = pickLatestResetSchedulingCandidate(cnProviderThresholdCandidates(account, PlatformZhipu), threshold, now)
default:
return decision
}
@@ -303,6 +307,45 @@ func grokThresholdCandidates(account *Account) []*accountSchedulingThresholdCand
}
}
// cnProviderThresholdCandidates 读取国产供应商 Coding Plan 账号的 5h / weekly 滚动窗口
// 用量快照(由 CNProviderQuotaService 写入 account.Extra,键形如
// <provider>_5h_used_percent / <provider>_weekly_reset_at)。payg 账号无此快照,
// 候选为空 → 不触发阈值停调(余额型走余额检测)。与 openai 的快照驱动停调一致:
// 仅当用量超阈值且窗口尚未重置时才停调。
func cnProviderThresholdCandidates(account *Account, provider string) []*accountSchedulingThresholdCandidate {
if account == nil || len(account.Extra) == 0 {
return nil
}
return []*accountSchedulingThresholdCandidate{
cnThresholdCandidate(account.Extra, provider, "5h"),
cnThresholdCandidate(account.Extra, provider, "weekly"),
}
}
func cnThresholdCandidate(extra map[string]any, provider, window string) *accountSchedulingThresholdCandidate {
var usedKey, resetKey string
switch window {
case "5h":
usedKey = cnExtraKey(provider, cnExtraSuffix5hUsed)
resetKey = cnExtraKey(provider, cnExtraSuffix5hReset)
case "weekly":
usedKey = cnExtraKey(provider, cnExtraSuffixWeeklyUsed)
resetKey = cnExtraKey(provider, cnExtraSuffixWeeklyReset)
default:
return nil
}
usedPercent, ok := extra[usedKey]
if !ok {
return nil
}
return &accountSchedulingThresholdCandidate{
window: window,
scope: provider,
usedPercent: schedulingPercentValue(usedPercent),
until: parseSchedulingResetAt(extra[resetKey]),
}
}
func pickLatestResetSchedulingCandidate(candidates []*accountSchedulingThresholdCandidate, threshold int, now time.Time) *accountSchedulingThresholdCandidate {
var winner *accountSchedulingThresholdCandidate
for _, candidate := range candidates {
@@ -509,6 +509,9 @@ func (s *AccountService) TestCredentials(ctx context.Context, id int64) error {
case PlatformGrok:
// Grok OAuth credentials are validated via token exchange/refresh and request-path probes.
return nil
case PlatformKimi, PlatformZhipu, PlatformDeepseek:
// 国产 OpenAI 兼容供应商:凭证为 API Key,实际可用性经余额/额度探测与转发路径验证。
return nil
default:
return fmt.Errorf("unsupported platform: %s", account.Platform)
}
@@ -610,10 +610,11 @@ func (s *AccountTestService) testOpenAIAccountConnection(c *gin.Context, account
}
// Align test routing with gateway behavior: OpenAI accounts apply normal
// account model mapping, and compact mode applies compact-only mapping on top.
// account model mapping. Native remote compaction v2 rides the ordinary
// /responses wire and does NOT apply the legacy compact-only mapping
// (post-#5641 semantics: compact_model_mapping is /responses/compact-only).
testModelID = account.GetMappedModel(testModelID)
if mode == AccountTestModeCompact {
testModelID = resolveOpenAICompactForwardModel(account, testModelID)
return s.testOpenAICompactConnection(c, account, testModelID)
}
@@ -969,6 +970,7 @@ func (s *AccountTestService) observeGrokTestResponse(ctx context.Context, accoun
resp.Body = io.NopCloser(bytes.NewReader(responseBody))
}
snapshot := parseGrokQuotaSnapshot(resp.Header, resp.StatusCode, now)
stampGrokQuotaSnapshotForPlan(account, snapshot, grokRequestedModelFromCtx(ctx))
if snapshot != nil && s.accountRepo != nil {
resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now)
if limited {
@@ -1073,7 +1075,7 @@ func (s *AccountTestService) testGrokResponsesConnection(c *gin.Context, ctx con
}
defer func() { _ = resp.Body.Close() }()
s.observeGrokTestResponse(ctx, account, resp)
s.observeGrokTestResponse(withGrokTeamRateLimitModel(ctx, testModelID), account, resp)
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
@@ -1421,7 +1423,7 @@ User query:
return s.sendErrorAndEnd(c, fmt.Sprintf("standalone web_search probe failed: %s", err.Error()))
}
defer func() { _ = resp.Body.Close() }()
s.observeGrokTestResponse(ctx, account, resp)
s.observeGrokTestResponse(withGrokTeamRateLimitModel(ctx, grokDefaultResponsesModel), account, resp)
body, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK {
@@ -1975,8 +1977,10 @@ func (s *AccountTestService) testOpenAIChatCompletionsConnection(
return s.processOpenAIChatCompletionsStream(c, resp.Body)
}
// testOpenAICompactConnection probes /responses/compact and persists the
// resulting capability state on the account.
// testOpenAICompactConnection probes native remote compaction v2 (streaming
// /responses with a compaction_trigger input item) and persists the resulting
// capability state on the account. The legacy unary /responses/compact
// endpoint has been sunset upstream (404, #5598/#5624) and is no longer probed.
func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account *Account, testModelID string) error {
ctx := c.Request.Context()
credentialAccount := account
@@ -2001,7 +2005,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
if authToken == "" && !credentialAccount.IsOpenAIAgentIdentity() {
return s.sendErrorAndEnd(c, "No access token available")
}
apiURL = chatgptCodexAPIURL + "/compact"
apiURL = chatgptCodexAPIURL
case account.Type == AccountTypeAPIKey:
authToken = account.GetOpenAIApiKey()
if authToken == "" {
@@ -2015,7 +2019,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
if err != nil {
return s.sendErrorAndEnd(c, fmt.Sprintf("Invalid base URL: %s", err.Error()))
}
apiURL = appendOpenAIResponsesRequestPathSuffix(buildOpenAIResponsesURL(normalizedBaseURL), "/compact")
apiURL = buildOpenAIResponsesURL(normalizedBaseURL)
default:
return s.sendErrorAndEnd(c, fmt.Sprintf("Unsupported account type: %s", account.Type))
}
@@ -2026,7 +2030,11 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
c.Writer.Header().Set("X-Accel-Buffering", "no")
c.Writer.Flush()
payloadBytes, _ := json.Marshal(createOpenAICompactProbePayload(testModelID))
// 原生 v2 走普通 /responses 线:OAuth 与真实转发一致做上游模型归一化。
if isOAuth {
testModelID = normalizeOpenAIModelForUpstream(credentialAccount, testModelID)
}
payloadBytes, _ := json.Marshal(createOpenAICompactProbePayload(testModelID, isOAuth))
if !agentIdentityTaskRecoveryWasTried(ctx) {
s.sendEvent(c, TestEvent{Type: "test_start", Model: testModelID})
}
@@ -2038,7 +2046,9 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileOpenAI))
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "application/json")
// v2 探测是流式请求;同时补注协商头,与真实 codex 出站线型一致。
req.Header.Set("Accept", "text/event-stream")
ensureOpenAIRemoteCompactionV2BetaFeature(req.Header)
if credentialAccount.IsOpenAIAgentIdentity() {
authHeaders, authErr := buildAgentIdentityAuthenticationHeaders(ctx, s.accountRepo, s.agentIdentityWS, &s.agentIdentityTaskMu, credentialAccount)
if authErr != nil {
@@ -2060,6 +2070,12 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
if isOAuth {
req.Host = "chatgpt.com"
setOpenAIChatGPTAccountHeaders(req.Header, credentialAccount)
// 指纹收敛:探测与真实转发走同一个 /responses 端点,身份也必须同构,
// 否则探测流量会以「缺 x-codex-installation-id + 非收敛 session」的
// 形态暴露在上游眼里。账号关闭收敛(off)时返回 nil,探测保持原样。
if fpIDs := resolveCodexFingerprintIDsFromRequest(account, req.Header); fpIDs != nil {
applyCodexFingerprintHeaders(req.Header, fpIDs)
}
}
// 账号级请求头覆写:测试请求与真实转发保持一致的最终头
@@ -2073,7 +2089,7 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
resp, err := s.httpUpstream.DoWithTLS(req, proxyURL, account.ID, account.Concurrency, s.tlsFPProfileService.ResolveTLSProfile(account))
if err != nil {
if s.accountRepo != nil {
updates := buildOpenAICompactProbeExtraUpdates(nil, nil, err, time.Now())
updates := buildOpenAICompactProbeExtraUpdates(nil, nil, err, false, time.Now())
_ = s.accountRepo.UpdateExtra(ctx, account.ID, updates)
mergeAccountExtra(account, updates)
}
@@ -2092,8 +2108,9 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
return s.testOpenAICompactConnection(c, account, testModelID)
}
compactionFound := openAICompactProbeFoundCompactionItem(body)
if s.accountRepo != nil {
updates := buildOpenAICompactProbeExtraUpdates(resp, body, nil, time.Now())
updates := buildOpenAICompactProbeExtraUpdates(resp, body, nil, compactionFound, time.Now())
if codexUpdates, err := extractOpenAICodexProbeUpdates(resp); err == nil && len(codexUpdates) > 0 {
updates = mergeExtraUpdates(updates, codexUpdates)
}
@@ -2115,7 +2132,11 @@ func (s *AccountTestService) testOpenAICompactConnection(c *gin.Context, account
return s.sendErrorAndEnd(c, fmt.Sprintf("API returned %d: %s", resp.StatusCode, string(body)))
}
s.sendEvent(c, TestEvent{Type: "content", Text: "Compact probe succeeded"})
if !compactionFound {
return s.sendErrorAndEnd(c, "Upstream returned 2xx without a compaction output item (native remote compaction v2 unsupported on this chain)")
}
s.sendEvent(c, TestEvent{Type: "content", Text: "Compact probe succeeded (native remote compaction v2)"})
s.sendEvent(c, TestEvent{Type: "test_complete", Success: true})
return nil
}
@@ -10,10 +10,16 @@ import (
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
// compactProbeSSESuccessBody 是原生 v2 压缩成功的最小 SSE 形态:
// output_item.done 携带 compaction item + response.completed。
const compactProbeSSESuccessBody = "data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"compaction\",\"id\":\"cmp_probe\",\"encrypted_content\":\"blob\"}}\n\n" +
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_probe\",\"output\":[]}}\n\n"
func TestAccountTestService_TestAccountConnection_OpenAICompactOAuthSuccessPersistsSupport(t *testing.T) {
gin.SetMode(gin.TestMode)
@@ -38,8 +44,8 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactOAuthSuccessPersi
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}, "x-request-id": []string{"rid-probe"}},
Body: io.NopCloser(strings.NewReader(`{"id":"cmp_probe","status":"completed"}`)),
Header: http.Header{"Content-Type": []string{"text/event-stream"}, "x-request-id": []string{"rid-probe"}},
Body: io.NopCloser(strings.NewReader(compactProbeSSESuccessBody)),
}}
svc := &AccountTestService{
accountRepo: repo,
@@ -53,16 +59,22 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactOAuthSuccessPersi
err := svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact)
require.NoError(t, err)
require.Equal(t, chatgptCodexAPIURL+"/compact", upstream.lastReq.URL.String())
// 原生 v2:探测普通 /responses 线,不再打已下线的 /responses/compact。
require.Equal(t, chatgptCodexAPIURL, upstream.lastReq.URL.String())
require.Equal(t, "chatgpt.com", upstream.lastReq.Host)
require.Equal(t, "application/json", upstream.lastReq.Header.Get("Accept"))
require.Equal(t, codexCLIVersion, upstream.lastReq.Header.Get("Version"))
require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept"))
require.Contains(t, upstream.lastReq.Header.Get("x-codex-beta-features"), "remote_compaction_v2")
require.NotEmpty(t, upstream.lastReq.Header.Get("Session_Id"))
require.Equal(t, HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileFromContext(upstream.lastReq.Context()))
require.Equal(t, codexCLIUserAgent, upstream.lastReq.Header.Get("User-Agent"))
require.Equal(t, "chatgpt-acc", upstream.lastReq.Header.Get("chatgpt-account-id"))
require.Equal(t, "true", upstream.lastReq.Header.Get("x-openai-fedramp"))
require.Equal(t, "gpt-5.4", gjson.GetBytes(upstream.lastBody, "model").String())
require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
require.False(t, gjson.GetBytes(upstream.lastBody, "store").Bool())
inputItems := gjson.GetBytes(upstream.lastBody, "input").Array()
require.NotEmpty(t, inputItems)
require.Equal(t, "compaction_trigger", inputItems[len(inputItems)-1].Get("type").String())
updates := <-updateCalls
require.Equal(t, true, updates["openai_compact_supported"])
@@ -114,7 +126,7 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactOAuth404MarksUnsu
require.Contains(t, rec.Body.String(), `"type":"error"`)
}
func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyUsesCompactPath(t *testing.T) {
func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyUsesNativeResponsesPath(t *testing.T) {
gin.SetMode(gin.TestMode)
updateCalls := make(chan map[string]any, 1)
@@ -127,8 +139,10 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyUsesCompact
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"api_key": "sk-test",
"base_url": "https://example.com/v1",
"api_key": "sk-test",
"base_url": "https://example.com/v1",
// post-#5641:compact_model_mapping 仅作用于 legacy /responses/compact,
// 原生 v2 探测不应用它。
"compact_model_mapping": map[string]any{"gpt-5.4": "gpt-5.4-openai-compact"},
},
}
@@ -138,8 +152,8 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyUsesCompact
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"id":"cmp_probe_apikey","status":"completed"}`)),
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(compactProbeSSESuccessBody)),
}}
svc := &AccountTestService{
accountRepo: repo,
@@ -154,14 +168,16 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyUsesCompact
err := svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact)
require.NoError(t, err)
require.Equal(t, "https://example.com/v1/responses/compact", upstream.lastReq.URL.String())
require.Equal(t, "https://example.com/v1/responses", upstream.lastReq.URL.String())
requireOpenAICodexProbeHeaders(t, upstream.lastReq.Header)
require.Equal(t, "gpt-5.4-openai-compact", gjson.GetBytes(upstream.lastBody, "model").String())
require.Contains(t, upstream.lastReq.Header.Get("x-codex-beta-features"), "remote_compaction_v2")
require.Equal(t, "gpt-5.4", gjson.GetBytes(upstream.lastBody, "model").String(),
"原生 v2 探测不应用 compact_model_mapping")
updates := <-updateCalls
require.Equal(t, true, updates["openai_compact_supported"])
}
func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyDefaultBaseURLUsesV1Path(t *testing.T) {
func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyDefaultBaseURLUsesResponsesPath(t *testing.T) {
gin.SetMode(gin.TestMode)
updateCalls := make(chan map[string]any, 1)
@@ -183,8 +199,8 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyDefaultBase
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(`{"id":"cmp_probe_apikey_default","status":"completed"}`)),
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(compactProbeSSESuccessBody)),
}}
svc := &AccountTestService{
accountRepo: repo,
@@ -198,6 +214,112 @@ func TestAccountTestService_TestAccountConnection_OpenAICompactAPIKeyDefaultBase
err := svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact)
require.NoError(t, err)
require.Equal(t, "https://api.openai.com/v1/responses/compact", upstream.lastReq.URL.String())
require.Equal(t, "https://api.openai.com/v1/responses", upstream.lastReq.URL.String())
<-updateCalls
}
func TestAccountTestService_TestAccountConnection_OpenAICompact2xxWithoutItemMarksUnsupported(t *testing.T) {
gin.SetMode(gin.TestMode)
updateCalls := make(chan map[string]any, 1)
account := Account{
ID: 5,
Name: "openai-oauth-no-item",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "oauth-token",
"chatgpt_account_id": "chatgpt-acc",
},
}
repo := &snapshotUpdateAccountRepo{
stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}},
updateExtraCalls: updateCalls,
}
// 200 但流里没有 compaction item:链路吞掉了 compaction_trigger 的形态
//(#5478 的 "got 0 items"),必须判定为不支持。
noItemBody := "data: {\"type\":\"response.output_item.done\",\"item\":{\"type\":\"message\",\"id\":\"msg_1\"}}\n\n" +
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_x\",\"output\":[]}}\n\n"
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(noItemBody)),
}}
svc := &AccountTestService{
accountRepo: repo,
httpUpstream: upstream,
}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/5/test", bytes.NewReader(nil))
err := svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact)
require.Error(t, err)
updates := <-updateCalls
require.Equal(t, false, updates["openai_compact_supported"])
require.Contains(t, rec.Body.String(), `"type":"error"`)
}
// 探测与真实转发走同一 /responses 端点,出站身份必须与真实 Codex 同构:
// session/thread 为 UUID、携带 x-codex-installation-id(收敛账号用收敛值)。
func TestAccountTestService_TestAccountConnection_OpenAICompactProbeIdentityMatchesRealTraffic(t *testing.T) {
gin.SetMode(gin.TestMode)
updateCalls := make(chan map[string]any, 1)
account := Account{
ID: 6,
Name: "openai-oauth-identity",
Platform: PlatformOpenAI,
Type: AccountTypeOAuth,
Status: StatusActive,
Schedulable: true,
Concurrency: 1,
Credentials: map[string]any{
"access_token": "oauth-token",
"chatgpt_account_id": "chatgpt-acc",
},
// 收敛是显式 opt-in(#5610),这里显式开启以验证探测身份与真实流量同构。
Extra: map[string]any{"codex_fingerprint_mode": "session"},
}
repo := &snapshotUpdateAccountRepo{
stubOpenAIAccountRepo: stubOpenAIAccountRepo{accounts: []Account{account}},
updateExtraCalls: updateCalls,
}
upstream := &httpUpstreamRecorder{resp: &http.Response{
StatusCode: http.StatusOK,
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
Body: io.NopCloser(strings.NewReader(compactProbeSSESuccessBody)),
}}
svc := &AccountTestService{accountRepo: repo, httpUpstream: upstream}
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/6/test", bytes.NewReader(nil))
require.NoError(t, svc.TestAccountConnection(c, account.ID, "gpt-5.4", "", AccountTestModeCompact))
// 显式 session 收敛模式:出站身份 = 账号级收敛值
converged := resolveConvergedSessionID(&account)
require.Equal(t, converged, upstream.lastReq.Header.Get("session-id"))
require.Equal(t, converged, upstream.lastReq.Header.Get("session_id"))
require.Equal(t, resolveConvergedInstallationID(&account), upstream.lastReq.Header.Get("x-codex-installation-id"),
"真实 Codex 每个请求必带 installation-id,探测不得缺失")
require.NotContains(t, upstream.lastReq.Header.Get("session-id"), "probe_compact",
"探测标识不得是可被上游一眼识别的字面量")
<-updateCalls
}
func TestCompactProbeSessionID_IsUUIDShaped(t *testing.T) {
for _, id := range []int64{0, 1, 987654} {
got := compactProbeSessionID(id)
_, err := uuid.Parse(got)
require.NoError(t, err, "探测会话标识必须是 UUID 形态: %s", got)
}
require.Equal(t, compactProbeSessionID(7), compactProbeSessionID(7), "同账号应稳定复用同一会话")
require.NotEqual(t, compactProbeSessionID(7), compactProbeSessionID(8))
}
+48
View File
@@ -301,6 +301,10 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
}
platform := NormalizeGroupPlatform(input.Platform)
modelPricing, err := normalizeGroupModelPricing(platform, input.ModelPricing)
if err != nil {
return nil, err
}
maxReasoningEffort, err := normalizeMaxReasoningEffortForPlatform(platform, input.MaxReasoningEffort)
if err != nil {
return nil, infraerrors.Newf(http.StatusBadRequest, "INVALID_MAX_REASONING_EFFORT", "%v", err)
@@ -459,6 +463,8 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
DailyLimitUSD: dailyLimit,
WeeklyLimitUSD: weeklyLimit,
MonthlyLimitUSD: monthlyLimit,
LongContextPricingEnabled: input.LongContextPricingEnabled,
ModelPricing: modelPricing,
AllowImageGeneration: allowImageGeneration,
AllowBatchImageGeneration: allowBatchImageGeneration,
ImageRateIndependent: input.ImageRateIndependent,
@@ -635,6 +641,9 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
return nil, err
}
// 渠道缓存里存了 groupID → platform 的映射,改了平台要让它失效(见函数末尾)
previousPlatform := group.Platform
if input.Name != "" {
group.Name = input.Name
}
@@ -656,6 +665,16 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
if input.Status != "" {
group.Status = input.Status
}
if input.LongContextPricingEnabled != nil {
group.LongContextPricingEnabled = *input.LongContextPricingEnabled
}
if input.ModelPricing != nil {
modelPricing, normalizeErr := normalizeGroupModelPricing(group.Platform, *input.ModelPricing)
if normalizeErr != nil {
return nil, normalizeErr
}
group.ModelPricing = modelPricing
}
// 订阅相关字段
if input.SubscriptionType != "" {
@@ -884,6 +903,13 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
s.authCacheInvalidator.InvalidateAuthCacheByGroupID(ctx, id)
}
// 平台变了就失效渠道缓存:该缓存持有 groupID → platform,而渠道定价 / 模型映射 /
// 模型白名单都按平台严格隔离。不失效的话,缓存最长 10 分钟仍按旧平台匹配,
// 期间定价查不到会静默回落到 LiteLLM 价格表、映射与白名单也不生效。
if group.Platform != previousPlatform && s.channelCacheInvalidator != nil {
s.channelCacheInvalidator.InvalidateCache()
}
// 如果指定了复制账号的源分组,同步绑定(替换当前分组的账号)
if len(input.CopyAccountsFromGroupIDs) > 0 {
// 去重源分组 IDs
@@ -955,6 +981,28 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
return group, nil
}
func normalizeGroupModelPricing(platform string, pricing []ChannelModelPricing) ([]ChannelModelPricing, error) {
out := make([]ChannelModelPricing, len(pricing))
for i := range pricing {
out[i] = pricing[i].Clone()
out[i].ID = 0
out[i].ChannelID = 0
if strings.TrimSpace(out[i].Platform) == "" {
out[i].Platform = platform
}
for j := range out[i].Models {
out[i].Models[j] = strings.TrimSpace(out[i].Models[j])
}
if len(out[i].Models) == 0 {
return nil, infraerrors.New(http.StatusBadRequest, "GROUP_MODEL_PRICING_MODELS_REQUIRED", "group model pricing entry requires at least one model")
}
}
if err := validatePricingEntries(out); err != nil {
return nil, err
}
return out, nil
}
func (s *adminServiceImpl) DeleteGroup(ctx context.Context, id int64) error {
var groupKeys []string
if s.authCacheInvalidator != nil {
@@ -0,0 +1,87 @@
//go:build unit
package service
import (
"context"
"testing"
"github.com/stretchr/testify/require"
)
// groupPlatformRepoStub 只实现 UpdateGroup 走到的两个方法,其余靠内嵌接口占位。
type groupPlatformRepoStub struct {
GroupRepository
group *Group
updated *Group
}
func (r *groupPlatformRepoStub) GetByID(_ context.Context, _ int64) (*Group, error) {
cloned := *r.group
return &cloned, nil
}
func (r *groupPlatformRepoStub) Update(_ context.Context, group *Group) error {
r.updated = group
return nil
}
type channelCacheInvalidatorSpy struct {
calls int
}
func (s *channelCacheInvalidatorSpy) InvalidateCache() { s.calls++ }
// 渠道缓存持有 groupID → platform,而渠道定价/模型映射/模型白名单都按平台严格隔离。
// 改了分组平台却不失效缓存,最长 10 分钟内这些查找仍按旧平台匹配(静默走错价)。
func TestUpdateGroupInvalidatesChannelCacheOnPlatformChange(t *testing.T) {
tests := []struct {
name string
fromPlatform string
inputPlatform string
wantCalls int
}{
{
name: "platform changed invalidates",
fromPlatform: PlatformAnthropic,
inputPlatform: PlatformOpenAI,
wantCalls: 1,
},
{
name: "same platform does not invalidate",
fromPlatform: PlatformAnthropic,
inputPlatform: PlatformAnthropic,
wantCalls: 0,
},
{
// 请求里不带 platform 字段时不应该动缓存
name: "platform omitted does not invalidate",
fromPlatform: PlatformAnthropic,
inputPlatform: "",
wantCalls: 0,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
repo := &groupPlatformRepoStub{group: &Group{ID: 7, Name: "g", Platform: tt.fromPlatform}}
spy := &channelCacheInvalidatorSpy{}
svc := &adminServiceImpl{groupRepo: repo, channelCacheInvalidator: spy}
got, err := svc.UpdateGroup(context.Background(), 7, &UpdateGroupInput{Platform: tt.inputPlatform})
require.NoError(t, err)
require.NotNil(t, got)
require.Equal(t, tt.wantCalls, spy.calls)
})
}
}
// 依赖可以不注入(例如测试或裁剪构建),此时不应 panic——缓存靠 TTL 自然重建。
func TestUpdateGroupWithoutChannelCacheInvalidator(t *testing.T) {
repo := &groupPlatformRepoStub{group: &Group{ID: 7, Name: "g", Platform: PlatformAnthropic}}
svc := &adminServiceImpl{groupRepo: repo}
got, err := svc.UpdateGroup(context.Background(), 7, &UpdateGroupInput{Platform: PlatformOpenAI})
require.NoError(t, err)
require.Equal(t, PlatformOpenAI, got.Platform)
}
+34 -19
View File
@@ -209,15 +209,17 @@ type AdminBoundAuthIdentityChannel struct {
}
type CreateGroupInput struct {
Name string
Description string
Platform string
RateMultiplier float64
IsExclusive bool
SubscriptionType string // standard/subscription
DailyLimitUSD *float64 // 日限额 (USD)
WeeklyLimitUSD *float64 // 周限额 (USD)
MonthlyLimitUSD *float64 // 月限额 (USD)
Name string
Description string
Platform string
RateMultiplier float64
IsExclusive bool
SubscriptionType string // standard/subscription
DailyLimitUSD *float64 // 日限额 (USD)
WeeklyLimitUSD *float64 // 周限额 (USD)
MonthlyLimitUSD *float64 // 月限额 (USD)
LongContextPricingEnabled bool
ModelPricing []ChannelModelPricing
// 图片生成计费配置(仅 antigravity 平台使用)
AllowImageGeneration bool
AllowBatchImageGeneration bool
@@ -281,16 +283,18 @@ type CreateGroupInput struct {
}
type UpdateGroupInput struct {
Name string
Description *string
Platform string
RateMultiplier *float64 // 使用指针以支持设置为0
IsExclusive *bool
Status string
SubscriptionType string // standard/subscription
DailyLimitUSD *float64 // 日限额 (USD)
WeeklyLimitUSD *float64 // 周限额 (USD)
MonthlyLimitUSD *float64 // 月限额 (USD)
Name string
Description *string
Platform string
RateMultiplier *float64 // 使用指针以支持设置为0
IsExclusive *bool
Status string
SubscriptionType string // standard/subscription
DailyLimitUSD *float64 // 日限额 (USD)
WeeklyLimitUSD *float64 // 周限额 (USD)
MonthlyLimitUSD *float64 // 月限额 (USD)
LongContextPricingEnabled *bool
ModelPricing *[]ChannelModelPricing
// 图片生成计费配置(仅 antigravity 平台使用)
AllowImageGeneration *bool
AllowBatchImageGeneration *bool
@@ -662,6 +666,14 @@ type adminServiceImpl struct {
affiliateService adminRechargeAffiliateAccruer
compositeRouteRepo CompositeModelRouteRepository
compositeResolver *CompositeRouteResolver
// 分组平台变更后用来失效渠道缓存;可为 nil(缓存会在 TTL 到期后自然重建)
channelCacheInvalidator ChannelCacheInvalidator
}
// ChannelCacheInvalidator 失效渠道缓存。
// 窄接口,避免 admin 服务依赖整个 ChannelService——与 APIKeyAuthCacheInvalidator 同一思路。
type ChannelCacheInvalidator interface {
InvalidateCache()
}
type adminRechargeAffiliateAccruer interface {
@@ -695,6 +707,7 @@ func NewAdminService(
affiliateService *AffiliateService,
compositeRouteRepo CompositeModelRouteRepository,
compositeResolver *CompositeRouteResolver,
channelCacheInvalidator ChannelCacheInvalidator,
) AdminService {
return &adminServiceImpl{
userRepo: userRepo,
@@ -721,5 +734,7 @@ func NewAdminService(
affiliateService: affiliateService,
compositeRouteRepo: compositeRouteRepo,
compositeResolver: compositeResolver,
channelCacheInvalidator: channelCacheInvalidator,
}
}
@@ -14,9 +14,14 @@ const (
// GetAnthropicAPIKeyAuthScheme returns the upstream authentication scheme for
// Anthropic API-key accounts. Missing or invalid values keep the historical
// x-api-key behavior.
// x-api-key behavior. CN providers using their native Anthropic endpoints
// (api_protocol=anthropic) share the same override knob — Kimi/DeepSeek default
// to x-api-key, Zhipu can opt into Authorization: Bearer.
func (a *Account) GetAnthropicAPIKeyAuthScheme() string {
if a == nil || a.Platform != PlatformAnthropic || a.Type != AccountTypeAPIKey {
if a == nil || a.Type != AccountTypeAPIKey {
return AnthropicAPIKeyAuthSchemeXAPIKey
}
if a.Platform != PlatformAnthropic && !a.IsCNProvider() {
return AnthropicAPIKeyAuthSchemeXAPIKey
}
@@ -4,6 +4,7 @@ import (
"compress/gzip"
"context"
"crypto/sha256"
"database/sql"
"encoding/hex"
"encoding/json"
"errors"
@@ -31,6 +32,21 @@ const (
maxBackupRecords = 100
backupObjectCleanupTimeout = 2 * time.Minute
// backupScheduledLeaderLockKey gates the scheduled full-database backup so
// that only one instance in a clustered deployment performs the
// dump-and-upload each cycle. Without it every instance runs the cron
// independently, producing N concurrent pg_dumps against the same database,
// N× peak memory while the archive is uploaded, and N identical objects that
// overwrite the same timestamped key. Every other periodic job in this
// package is already gated the same way; the scheduled backup was the last
// one that still fanned out across every instance.
backupScheduledLeaderLockKey = "backup:scheduled:leader"
// backupScheduledLeaderLockTTL bounds crash recovery only; the lock is
// released as soon as the backup finishes. It must exceed the job's
// worst-case runtime (the scheduled backup context is bounded at 30m) so the
// lock cannot expire mid-dump and let a peer start a second backup.
backupScheduledLeaderLockTTL = 35 * time.Minute
)
var (
@@ -161,6 +177,14 @@ type BackupService struct {
cronSched *cron.Cron
cronEntryID cron.EntryID
// lockCache/db elect a single leader for the scheduled backup across
// instances; instanceID identifies this process as the lock owner. Injected
// via SetLeaderLock — when both are nil the backup runs ungated
// (single-instance / test behavior).
lockCache LeaderLockCache
db *sql.DB
instanceID string
wg sync.WaitGroup // 追踪活跃的备份/恢复 goroutine
shuttingDown atomic.Bool // 阻止新备份启动
bgCtx context.Context // 所有后台操作的 parent context
@@ -186,9 +210,21 @@ func NewBackupService(
bgCtx: bgCtx,
bgCancel: bgCancel,
partSizeBytes: defaultBackupPartSizeBytes,
instanceID: uuid.NewString(),
}
}
// SetLeaderLock injects the leader-lock cache and DB used to elect a single
// instance for the scheduled backup. When both are nil the scheduled backup runs
// ungated (single-instance / test behavior).
func (s *BackupService) SetLeaderLock(lockCache LeaderLockCache, db *sql.DB) {
if s == nil {
return
}
s.lockCache = lockCache
s.db = db
}
// Start 启动定时备份调度器并清理孤立记录
func (s *BackupService) Start() {
s.cronSched = cron.New()
@@ -467,6 +503,16 @@ func (s *BackupService) runScheduledBackup() {
ctx, cancel := context.WithTimeout(s.bgCtx, 30*time.Minute)
defer cancel()
// 多实例保护: 集群部署时只让 leader 执行定时备份, 避免每个实例各自对同一个
// 数据库跑一次全量 dump、上传时峰值内存翻倍、以及多份同名对象互相覆盖。
// 手动触发的备份 (CreateBackup/StartBackup) 不受此限, 运维仍可随时在任一节点强制备份。
release, ok := tryAcquireSingletonLeaderLock(ctx, s.lockCache, s.db, backupScheduledLeaderLockKey, s.instanceID, backupScheduledLeaderLockTTL)
if !ok {
logger.LegacyPrintf("service.backup", "[Backup] 定时备份跳过: 本实例非 leader")
return
}
defer release()
// 读取定时备份配置中的过期天数
schedule, _ := s.GetSchedule(ctx)
expireDays := 14 // 默认14天过期
@@ -629,6 +629,56 @@ func TestBackupService_CreateBackup_ConcurrentBlocked(t *testing.T) {
require.ErrorIs(t, err, ErrBackupInProgress)
}
// TestBackupService_RunScheduledBackup_LeaderElection verifies the scheduled
// backup is gated by a cross-instance leader lock: a non-leader instance skips
// the dump entirely so a clustered deployment does not run N identical backups
// against the same database, while the leader runs it and releases the lock
// afterward. Manual backups (CreateBackup/StartBackup) are intentionally left
// ungated and are covered by the other tests.
func TestBackupService_RunScheduledBackup_LeaderElection(t *testing.T) {
t.Run("non-leader skips", func(t *testing.T) {
repo := newMockSettingRepo()
seedS3Config(t, repo)
store := newMockObjectStore()
svc := newTestBackupService(repo, &mockDumper{dumpData: []byte("data")}, store)
// A peer already owns the lock, so this instance is not the leader.
cache := &fakeLeaderLockCache{}
peerRelease, ok := tryAcquireSingletonLeaderLock(context.Background(), cache, nil, backupScheduledLeaderLockKey, "peer", time.Minute)
require.True(t, ok)
defer peerRelease()
svc.SetLeaderLock(cache, nil)
svc.runScheduledBackup()
store.mu.Lock()
require.Empty(t, store.objects, "non-leader must not upload a backup")
store.mu.Unlock()
records, err := svc.ListBackups(context.Background())
require.NoError(t, err)
require.Empty(t, records, "non-leader must not create a backup record")
require.Equal(t, "peer", cache.heldBy(backupScheduledLeaderLockKey), "peer keeps the lock")
})
t.Run("leader runs and releases", func(t *testing.T) {
repo := newMockSettingRepo()
seedS3Config(t, repo)
store := newMockObjectStore()
svc := newTestBackupService(repo, &mockDumper{dumpData: []byte("-- dump\n")}, store)
cache := &fakeLeaderLockCache{}
svc.SetLeaderLock(cache, nil)
svc.runScheduledBackup()
records, err := svc.ListBackups(context.Background())
require.NoError(t, err)
require.Len(t, records, 1, "leader creates exactly one backup record")
require.Equal(t, "completed", records[0].Status)
require.Empty(t, cache.heldBy(backupScheduledLeaderLockKey), "leader releases the lock when done")
})
}
func TestBackupService_RestoreBackup_Streaming(t *testing.T) {
repo := newMockSettingRepo()
seedS3Config(t, repo)
+1 -1
View File
@@ -139,7 +139,7 @@ func (s *GatewayService) handleBedrockStreamingResponse(
sseData = transformBedrockInvocationMetrics(sseData)
// 解析 SSE 事件数据提取 usage
s.parseSSEUsagePassthrough(string(sseData), usage)
parseSSEUsagePassthrough(string(sseData), usage)
// 确定 SSE event type
eventType := gjson.GetBytes(sseData, "type").String()
@@ -10,8 +10,8 @@ func TestCalculateSearchCost(t *testing.T) {
t.Parallel()
s := &BillingService{}
require.Equal(t, 0.0, s.CalculateSearchCost(0, floatPtr(10), 1).ActualCost)
// nil price → default $10/1k: 5 calls = 0.05
require.InDelta(t, 0.05, s.CalculateSearchCost(5, nil, 1).ActualCost, 1e-9)
// nil price -> official xAI default $5/1k: 5 calls = 0.025
require.InDelta(t, 0.025, s.CalculateSearchCost(5, nil, 1).ActualCost, 1e-9)
// explicit 0 → free
require.Equal(t, 0.0, s.CalculateSearchCost(5, floatPtr(0), 1).ActualCost)
price := 10.0
@@ -30,10 +30,10 @@ func TestCalculateAudioCost(t *testing.T) {
require.InDelta(t, 1.5, s.CalculateAudioCost("tts", 0.1, cfg, 1).ActualCost, 1e-9)
require.InDelta(t, 0.25, s.CalculateAudioCost("stt", 0.5, cfg, 1).ActualCost, 1e-9)
require.Equal(t, 0.0, s.CalculateAudioCost("unknown", 1, cfg, 1).ActualCost)
// nil config → defaults (realtime $0.10/min, tts $15/M, stt $0.36/hr)
require.InDelta(t, 0.10, s.CalculateAudioCost("realtime", 1, nil, 1).ActualCost, 1e-9)
// nil config -> official defaults (think-fast-1 $0.05/min, TTS $15/M, REST STT $0.10/hr)
require.InDelta(t, 0.05, s.CalculateAudioCost("realtime", 1, nil, 1).ActualCost, 1e-9)
require.InDelta(t, 15.0, s.CalculateAudioCost("tts", 1, nil, 1).ActualCost, 1e-9)
require.InDelta(t, 0.36, s.CalculateAudioCost("stt", 1, nil, 1).ActualCost, 1e-9)
require.InDelta(t, 0.10, s.CalculateAudioCost("stt", 1, nil, 1).ActualCost, 1e-9)
// explicit 0 → free
zero := 0.0
require.Equal(t, 0.0, s.CalculateAudioCost("realtime", 1, &audioPriceConfig{RealtimePerMin: &zero}, 1).ActualCost)
+124 -32
View File
@@ -10,6 +10,7 @@ import (
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
)
// APIKeyRateLimitCacheData holds rate limit usage data cached in Redis.
@@ -104,6 +105,7 @@ type ModelPricing struct {
CacheCreation1hPrice float64 // 1小时缓存创建每token价格 (USD)
SupportsCacheBreakdown bool // 是否支持详细的缓存分类
LongContextInputThreshold int // 超过阈值后按整次会话提升输入价格
LongContextThresholdInclusive bool // 达到阈值即应用(xAI);默认保持严格大于以兼容既有模型
LongContextInputMultiplier float64 // 长上下文整次会话输入倍率
LongContextOutputMultiplier float64 // 长上下文整次会话输出倍率
ImageOutputPricePerToken float64 // 图片输出 token 价格 (USD)
@@ -585,32 +587,56 @@ func (s *BillingService) initFallbackPricing() {
SupportsCacheBreakdown: false,
}
// xAI Grok 4.5 (official docs: $2 input / $0.50 cached input / $6 output per MTok)
// xAI Grok 4.5: $2 input / $0.30 cached input / $6 output below 200k.
s.fallbackPrices["grok-4.5"] = &ModelPricing{
InputPricePerToken: 2e-6,
OutputPricePerToken: 6e-6,
CacheReadPricePerToken: 0.5e-6,
SupportsCacheBreakdown: false,
InputPricePerToken: 2e-6,
OutputPricePerToken: 6e-6,
CacheReadPricePerToken: 0.3e-6,
SupportsCacheBreakdown: false,
LongContextInputThreshold: 200000,
LongContextThresholdInclusive: true,
LongContextInputMultiplier: 2,
LongContextOutputMultiplier: 2,
}
// xAI Grok 4.3 (official docs: $1.25 input / $2.50 output per MTok)
// xAI Grok 4.6 (docs.x.ai/developers/models: $2 input / $0.50 cached input /
// $6 output per MTok under 200k prompt tokens; ≥200k is 2× on input,
// cached input, and output).
s.fallbackPrices["grok-4.6"] = &ModelPricing{
InputPricePerToken: 2e-6,
OutputPricePerToken: 6e-6,
CacheReadPricePerToken: 0.5e-6,
SupportsCacheBreakdown: false,
LongContextInputThreshold: 200000,
LongContextThresholdInclusive: true,
LongContextInputMultiplier: 2,
LongContextOutputMultiplier: 2,
}
// xAI Grok 4.3: $1.25 input / $0.20 cached / $2.50 output below 200k.
s.fallbackPrices["grok-4.3"] = &ModelPricing{
InputPricePerToken: 1.25e-6,
OutputPricePerToken: 2.5e-6,
CacheReadPricePerToken: 0.2e-6,
SupportsCacheBreakdown: false,
LongContextInputThreshold: 1000000,
LongContextInputMultiplier: 1,
InputPricePerToken: 1.25e-6,
OutputPricePerToken: 2.5e-6,
CacheReadPricePerToken: 0.2e-6,
SupportsCacheBreakdown: false,
LongContextInputThreshold: 200000,
LongContextThresholdInclusive: true,
LongContextInputMultiplier: 2,
LongContextOutputMultiplier: 2,
}
// xAI Grok Build 0.1 (official docs: $1 input / $0.20 cached input /
// $2 output per MTok). Composer is available only through Grok Build and
// has no standalone public API rate card, so its aliases use this coding
// model rate instead of silently billing at zero.
s.fallbackPrices["grok-build-0.1"] = &ModelPricing{
InputPricePerToken: 1e-6,
OutputPricePerToken: 2e-6,
CacheReadPricePerToken: 0.2e-6,
SupportsCacheBreakdown: false,
InputPricePerToken: 1e-6,
OutputPricePerToken: 2e-6,
CacheReadPricePerToken: 0.2e-6,
SupportsCacheBreakdown: false,
LongContextInputThreshold: 200000,
LongContextThresholdInclusive: true,
LongContextInputMultiplier: 2,
LongContextOutputMultiplier: 2,
}
}
@@ -803,8 +829,10 @@ func (s *BillingService) getFallbackPricing(model string) *ModelPricing {
}
switch modelLower {
case "grok", "grok-latest", "grok-4.5", "grok-4.5-latest", "grok-build-latest":
case "grok", "grok-latest", "grok-4.5", "grok-4.5-latest":
return s.fallbackPrices["grok-4.5"]
case "grok-4.6", "grok-4.6-latest":
return s.fallbackPrices["grok-4.6"]
case "grok-4.3",
"grok-4.20-0309-reasoning",
"grok-4.20-0309-non-reasoning",
@@ -812,13 +840,59 @@ func (s *BillingService) getFallbackPricing(model string) *ModelPricing {
"grok-4.20-reasoning",
"grok-4.20-non-reasoning":
return s.fallbackPrices["grok-4.3"]
case "grok-build", "grok-build-0.1", "grok-composer", "grok-composer-2.5-fast", "composer-2.5":
case "grok-build", "grok-build-latest", "grok-build-0.1", "grok-composer", "grok-composer-2.5-fast", "composer-2.5":
return s.fallbackPrices["grok-build-0.1"]
}
// Unknown Grok text IDs (grok-5, dated snapshots, provider-prefixed) inherit
// the current default text card so a new model cannot ship unbilled.
if pricing := s.grokUnknownTextFamilyFallback(modelLower); pricing != nil {
return pricing
}
return nil
}
func (s *BillingService) grokUnknownTextFamilyFallback(model string) *ModelPricing {
if s == nil || !isGrokUnknownTextFamilyModel(model) {
return nil
}
return s.fallbackPrices["grok-4.5"]
}
func isGrokUnknownTextFamilyModel(model string) bool {
native := strings.ToLower(strings.TrimSpace(xai.StripGrokProviderPrefix(model)))
if isGrokMediaFamilyModel(native) {
return false
}
switch {
case native == "grok", native == "grok-latest":
return true
case strings.HasPrefix(native, "grok-build"),
strings.HasPrefix(native, "grok-composer"),
strings.HasPrefix(native, "composer-"):
return true
case len(native) > 5 && strings.HasPrefix(native, "grok-"):
rest := native[len("grok-"):]
return rest[0] >= '0' && rest[0] <= '9'
default:
return false
}
}
// isGrokMediaFamilyModel matches ids that are billed per image/video/audio unit
// rather than per token, so version-numbered media ids (grok-2-image-1212,
// grok-5-video) cannot slip into the unknown-text fallback and pick up a token
// card. "vision" is deliberately absent: multimodal chat models are token billed.
func isGrokMediaFamilyModel(native string) bool {
for _, marker := range []string{"imagine", "image", "video", "audio", "speech", "tts", "transcribe", "realtime"} {
if strings.Contains(native, marker) {
return true
}
}
return false
}
// HasIdentifiedTokenPricing 判断模型能否在价格表中被"确定性识别"出 token 价格。
//
// 与 GetModelPricing 的关键区别:本函数拒绝按子串猜系列的兜底。GetModelPricing 会
@@ -950,9 +1024,11 @@ type CostInput struct {
Ctx context.Context
Model string
GroupID *int64 // 用于渠道定价查找
Group *Group
Tokens UsageTokens
RequestCount int // 按次计费时使用
SizeTier string // 按次/图片模式的层级标签("1K","2K","4K","HD" 等)
RequestCount int // 按次计费时使用
UsageUnits float64 // 音频等连续计量单位(分钟/小时/百万字符)
SizeTier string // 按次/图片模式的层级标签("1K","2K","4K","HD" 等)
RateMultiplier float64
ServiceTier string // "priority","flex","" 等
Resolver *ModelPricingResolver // 定价解析器
@@ -985,6 +1061,7 @@ func (s *BillingService) CalculateCostUnified(input CostInput) (*CostBreakdown,
resolved = input.Resolver.Resolve(input.Ctx, PricingInput{
Model: input.Model,
GroupID: input.GroupID,
Group: input.Group,
})
}
@@ -996,7 +1073,7 @@ func (s *BillingService) CalculateCostUnified(input CostInput) (*CostBreakdown,
var breakdown *CostBreakdown
var err error
switch resolved.Mode {
case BillingModePerRequest, BillingModeImage:
case BillingModePerRequest, BillingModeImage, BillingModeVideo:
breakdown, err = s.calculatePerRequestCost(resolved, input)
default: // BillingModeToken
breakdown, err = s.calculateTokenCost(resolved, input)
@@ -1022,7 +1099,7 @@ func (s *BillingService) calculateTokenCost(resolved *ResolvedPricing, input Cos
pricing = s.applyModelSpecificPricingPolicy(input.Model, pricing)
// 长上下文定价仅在无区间定价时应用(区间定价已包含上下文分层)
applyLongCtx := len(resolved.Intervals) == 0
applyLongCtx := len(resolved.Intervals) == 0 && resolved.longContextPricingEnabled
if input.LongContextBillingEnabled != nil {
applyLongCtx = applyLongCtx && *input.LongContextBillingEnabled
}
@@ -1157,9 +1234,13 @@ func (s *BillingService) computeCacheCreationCost(pricing *ModelPricing, tokens
// calculatePerRequestCost 按次/图片计费
func (s *BillingService) calculatePerRequestCost(resolved *ResolvedPricing, input CostInput) (*CostBreakdown, error) {
count := input.RequestCount
if count <= 0 {
count = 1
units := input.UsageUnits
if units <= 0 {
count := input.RequestCount
if count <= 0 {
count = 1
}
units = float64(count)
}
var unitPrice float64
@@ -1178,7 +1259,7 @@ func (s *BillingService) calculatePerRequestCost(resolved *ResolvedPricing, inpu
unitPrice = resolved.DefaultPerRequestPrice
}
totalCost := unitPrice * float64(count)
totalCost := unitPrice * units
actualCost := totalCost * input.RateMultiplier
return &CostBreakdown{
@@ -1280,6 +1361,9 @@ func (s *BillingService) shouldApplySessionLongContextPricing(tokens UsageTokens
return false
}
totalInputTokens := tokens.InputTokens + tokens.CacheCreationTokens + tokens.CacheReadTokens
if pricing.LongContextThresholdInclusive {
return totalInputTokens >= pricing.LongContextInputThreshold
}
return totalInputTokens > pricing.LongContextInputThreshold
}
@@ -1451,6 +1535,8 @@ const (
defaultGrokImagineImagePrice2K = 0.02
defaultGrokImagineImageQualityPrice1K = 0.05
defaultGrokImagineImageQualityPrice2K = 0.07
defaultGrokImagineImage20Price1K = 0.06 // default quality is Medium
defaultGrokImagineImage20Price2K = 0.08
// 视频默认价为 xAI 官方**每秒**输出价格(USD/s),总价 = 每秒价 × 时长(秒)。
defaultGrokImagineVideoPrice480P = 0.05
@@ -1462,14 +1548,14 @@ const (
// Codex alpha/search 网页搜索单次默认价:OpenAI 官方 web search 定价 $10/1000 次。
defaultWebSearchPricePerCall = 0.01
// Grok /v1/web_search 与 SearchCount 附加费:与 Codex 对齐 $10/1000 次(按 1k 计价字段存储)。
defaultSearchPricePer1k = 10.0
// xAI server-side web/X search and code execution are $5/1000 calls.
defaultSearchPricePer1k = 5.0
// Grok Voice 默认价(分组列 NULL 时使用;显式配 0 表示免费)。
// 保守运营占位,运维可通过 groups.audio_* 覆盖。
defaultAudioRealtimePricePerMin = 0.10
// Generic realtime defaults to think-fast-1.0; think-fast-2.0 can be
// configured independently through per-model group/channel pricing.
defaultAudioRealtimePricePerMin = 0.05
defaultAudioTTSPricePerMillionChars = 15.0
defaultAudioSTTPricePerHour = 0.36
defaultAudioSTTPricePerHour = 0.10
)
// CalculateWebSearchCost 计算 Codex alpha/search 网页搜索按次费用。
@@ -1727,6 +1813,12 @@ func (s *BillingService) getDefaultVideoPrice(model string, resolution string) f
func getDefaultGrokImagineImagePrice(model string, imageSize string) (float64, bool) {
model = strings.ToLower(strings.TrimSpace(model))
switch model {
case "grok-imagine-image-2.0":
return getGrokImagineImageTierPrice(
imageSize,
defaultGrokImagineImage20Price1K,
defaultGrokImagineImage20Price2K,
), true
case "grok-imagine-image-quality":
return getGrokImagineImageTierPrice(
imageSize,
@@ -871,8 +871,8 @@ func TestComputeTokenBreakdown_GptImage2ImageEditIssue4386(t *testing.T) {
cost := svc.computeTokenBreakdown(pricing, tokens, 1.0, "", false)
wantTextInput := float64(19) * 5e-6 // 0.000095
wantImageInput := float64(352) * 8e-6 // 0.002816
wantTextInput := float64(19) * 5e-6 // 0.000095
wantImageInput := float64(352) * 8e-6 // 0.002816
wantImageOutput := float64(439) * 30e-6 // 0.013170
require.InDelta(t, wantTextInput, cost.InputCost, 1e-15, "InputCost 仅含文本输入")
require.InDelta(t, wantImageInput, cost.ImageInputCost, 1e-15, "图片输入按 $8/1M 独立计费")
@@ -1153,7 +1153,23 @@ func TestCalculateCostWithLongContext_PropagatesError(t *testing.T) {
func TestGetModelPricing_Grok45OfficialFallback(t *testing.T) {
svc := newTestBillingService()
for _, model := range []string{"grok", "grok-latest", "grok-4.5", "grok-4.5-latest", "grok-build-latest"} {
for _, model := range []string{"grok", "grok-latest", "grok-4.5", "grok-4.5-latest"} {
model := model
t.Run(model, func(t *testing.T) {
pricing, err := svc.GetModelPricing(model)
require.NoError(t, err)
require.InDelta(t, 2e-6, pricing.InputPricePerToken, 1e-12)
require.InDelta(t, 6e-6, pricing.OutputPricePerToken, 1e-12)
require.InDelta(t, 0.3e-6, pricing.CacheReadPricePerToken, 1e-12)
require.False(t, pricing.SupportsCacheBreakdown)
})
}
}
func TestGetModelPricing_Grok46OfficialFallback(t *testing.T) {
svc := newTestBillingService()
for _, model := range []string{"grok-4.6", "grok-4.6-latest"} {
model := model
t.Run(model, func(t *testing.T) {
pricing, err := svc.GetModelPricing(model)
@@ -1161,11 +1177,77 @@ func TestGetModelPricing_Grok45OfficialFallback(t *testing.T) {
require.InDelta(t, 2e-6, pricing.InputPricePerToken, 1e-12)
require.InDelta(t, 6e-6, pricing.OutputPricePerToken, 1e-12)
require.InDelta(t, 0.5e-6, pricing.CacheReadPricePerToken, 1e-12)
require.Equal(t, 200000, pricing.LongContextInputThreshold)
require.InDelta(t, 2.0, pricing.LongContextInputMultiplier, 1e-12)
require.InDelta(t, 2.0, pricing.LongContextOutputMultiplier, 1e-12)
require.False(t, pricing.SupportsCacheBreakdown)
})
}
}
func TestCalculateCostUnified_GroupLongContextToggleUsesPresetLadder(t *testing.T) {
svc := newTestBillingService()
resolver := NewModelPricingResolver(nil, svc)
tokens := UsageTokens{InputTokens: 250000, OutputTokens: 1000}
off := &Group{LongContextPricingEnabled: false}
disabled, err := svc.CalculateCostUnified(CostInput{
Model: "grok-4.5", Group: off, Tokens: tokens, RateMultiplier: 1, Resolver: resolver,
})
require.NoError(t, err)
on := &Group{LongContextPricingEnabled: true}
enabled, err := svc.CalculateCostUnified(CostInput{
Model: "grok-4.5", Group: on, Tokens: tokens, RateMultiplier: 1, Resolver: resolver,
})
require.NoError(t, err)
require.False(t, disabled.LongContextBillingApplied)
require.True(t, enabled.LongContextBillingApplied)
require.InDelta(t, disabled.InputCost*2, enabled.InputCost, 1e-12)
require.InDelta(t, disabled.OutputCost*2, enabled.OutputCost, 1e-12)
}
func TestGetModelPricing_UnknownGrokTextFallsBackToGrok45(t *testing.T) {
svc := newTestBillingService()
baseline, err := svc.GetModelPricing("grok-4.5")
require.NoError(t, err)
for _, model := range []string{"grok-5", "grok-5-latest", "x-ai/grok-7", "grok-4.7-beta"} {
pricing, err := svc.GetModelPricing(model)
require.NoError(t, err, "model %s", model)
require.InDelta(t, baseline.InputPricePerToken, pricing.InputPricePerToken, 1e-12, model)
require.InDelta(t, baseline.OutputPricePerToken, pricing.OutputPricePerToken, 1e-12, model)
require.InDelta(t, baseline.CacheReadPricePerToken, pricing.CacheReadPricePerToken, 1e-12, model)
}
// Per-unit media ids must not inherit the text card just because they carry
// a version number; they are billed by the image/video/audio paths instead.
for _, model := range []string{"grok-2-image-1212", "grok-2-audio", "grok-5-video", "x-ai/grok-6-image"} {
require.False(t, isGrokUnknownTextFamilyModel(model), "model %s", model)
}
// Multimodal chat models stay token billed.
require.True(t, isGrokUnknownTextFamilyModel("grok-2-vision-1212"))
for _, model := range []string{
"grok-imagine-image-3.0",
"grok-imagine-video-2",
"grok-voice-latest",
"grok-web-search",
"grok-x-search",
"grok-speech-1",
} {
_, err := svc.GetModelPricing(model)
require.Error(t, err, "non-text grok family %s must not inherit grok-4.5 token rates", model)
require.ErrorIs(t, err, ErrModelPricingUnavailable)
}
// Known cards stay on their own rate, not the 4.5 family floor.
build, err := svc.GetModelPricing("grok-build-0.1")
require.NoError(t, err)
require.InDelta(t, 1e-6, build.InputPricePerToken, 1e-12)
}
func TestGetModelPricing_GrokCatalogFallbacks(t *testing.T) {
svc := newTestBillingService()
+30 -30
View File
@@ -20,7 +20,7 @@ const (
// IsValid 检查 BillingMode 是否为合法值
func (m BillingMode) IsValid() bool {
switch m {
case BillingModeToken, BillingModePerRequest, BillingModeImage, "":
case BillingModeToken, BillingModePerRequest, BillingModeImage, BillingModeVideo, "":
return true
}
return false
@@ -87,38 +87,38 @@ type AccountStatsPricingRule struct {
// ChannelModelPricing 渠道模型定价条目
type ChannelModelPricing struct {
ID int64
ChannelID int64
Platform string // 所属平台(anthropic/openai/gemini/...)
Models []string // 绑定的模型列表
BillingMode BillingMode // 计费模式
InputPrice *float64 // 每 token 输入价格(USD)— 向后兼容 flat 定价
OutputPrice *float64 // 每 token 输出价格(USD)
CacheWritePrice *float64 // 缓存写入价格
CacheReadPrice *float64 // 缓存读取价格
ImageInputPrice *float64 // 图片输入 token 价格(如 gpt-image-2 图片编辑);未配置时回退文本输入价
ImageOutputPrice *float64 // 图片输出价格(向后兼容)
PerRequestPrice *float64 // 默认按次计费价格(USD)
Intervals []PricingInterval // 区间定价列表
CreatedAt time.Time
UpdatedAt time.Time
ID int64 `json:"id,omitempty"`
ChannelID int64 `json:"channel_id,omitempty"`
Platform string `json:"platform"` // 所属平台(anthropic/openai/gemini/...)
Models []string `json:"models"`
BillingMode BillingMode `json:"billing_mode"`
InputPrice *float64 `json:"input_price"`
OutputPrice *float64 `json:"output_price"`
CacheWritePrice *float64 `json:"cache_write_price"`
CacheReadPrice *float64 `json:"cache_read_price"`
ImageInputPrice *float64 `json:"image_input_price"`
ImageOutputPrice *float64 `json:"image_output_price"`
PerRequestPrice *float64 `json:"per_request_price"`
Intervals []PricingInterval `json:"intervals"`
CreatedAt time.Time `json:"created_at,omitempty"`
UpdatedAt time.Time `json:"updated_at,omitempty"`
}
// PricingInterval 定价区间(token 区间 / 按次分层 / 图片分辨率分层)
type PricingInterval struct {
ID int64
PricingID int64
MinTokens int // 区间下界(含)
MaxTokens *int // 区间上界(不含),nil = 无上限
TierLabel string // 层级标签(按次/图片模式:1K, 2K, 4K, HD 等)
InputPrice *float64 // token 模式:每 token 输入价
OutputPrice *float64 // token 模式:每 token 输出价
CacheWritePrice *float64 // token 模式:缓存写入价
CacheReadPrice *float64 // token 模式:缓存读取价
PerRequestPrice *float64 // 按次/图片模式:每次请求价格
SortOrder int
CreatedAt time.Time
UpdatedAt time.Time
ID int64 `json:"id,omitempty"`
PricingID int64 `json:"pricing_id,omitempty"`
MinTokens int `json:"min_tokens"`
MaxTokens *int `json:"max_tokens"`
TierLabel string `json:"tier_label"`
InputPrice *float64 `json:"input_price"`
OutputPrice *float64 `json:"output_price"`
CacheWritePrice *float64 `json:"cache_write_price"`
CacheReadPrice *float64 `json:"cache_read_price"`
PerRequestPrice *float64 `json:"per_request_price"`
SortOrder int `json:"sort_order"`
CreatedAt time.Time `json:"created_at,omitempty"`
UpdatedAt time.Time `json:"updated_at,omitempty"`
}
// IsActive 判断渠道是否启用
@@ -315,7 +315,7 @@ func ValidateIntervals(intervals []PricingInterval, mode BillingMode) error {
}
// per_request / image 模式按 tier_label 匹配,不做 token 区间重叠校验
if mode == BillingModePerRequest || mode == BillingModeImage {
if mode == BillingModePerRequest || mode == BillingModeImage || mode == BillingModeVideo {
return nil
}
return validateIntervalOverlap(sorted)
+26 -3
View File
@@ -375,6 +375,13 @@ func channelLookupPlatform(ctx context.Context, groupPlatform string) string {
}
return groupPlatform
}
// InvalidateCache 失效并重建渠道缓存。
// 供渠道以外、但会影响渠道缓存内容的变更调用(如分组平台变更)。
func (s *ChannelService) InvalidateCache() {
s.invalidateCache()
}
func (s *ChannelService) invalidateCache() {
s.cache.Store((*channelCache)(nil))
s.cacheSF.Forget("channel_cache")
@@ -659,7 +666,7 @@ func validatePricingBillingMode(pricing []ChannelModelPricing) error {
}
func checkBillingModeRequirements(p ChannelModelPricing) error {
if p.BillingMode == BillingModePerRequest || p.BillingMode == BillingModeImage {
if p.BillingMode == BillingModePerRequest || p.BillingMode == BillingModeImage || p.BillingMode == BillingModeVideo {
if p.PerRequestPrice == nil && len(p.Intervals) == 0 {
return infraerrors.BadRequest(
"BILLING_MODE_MISSING_PRICE",
@@ -962,19 +969,35 @@ func conflictsBetween(a, b modelEntry) bool {
}
}
// toModelEntry 将模型名转换为 modelEntry
// toModelEntry 将模型名转换为 modelEntry(用于模型映射的冲突检测)。
// 归一化必须与 expandMappingToCache 写缓存键的方式一致:映射缓存只做 strings.ToLower。
func toModelEntry(pattern string) modelEntry {
prefix, isWild := splitWildcardSuffix(strings.ToLower(pattern))
return modelEntry{pattern: pattern, prefix: prefix, wildcard: isWild}
}
// toPricingModelEntry 将模型名转换为 modelEntry(用于模型定价的冲突检测)。
//
// 与 toModelEntry 的区别:定价缓存的键走 normalizeChannelPricingModelName
// (额外做 TrimSpace,并把 claude-* 的 "." 换成 "-"),冲突检测必须用同一套归一化,
// 否则两个校验时看着不同、写进缓存后键相同的定价会互相静默覆盖。
func toPricingModelEntry(pattern string) modelEntry {
// 先剥通配符再归一化,与 expandPricingToCache 的处理顺序保持一致
prefix, isWild := splitWildcardSuffix(pattern)
return modelEntry{
pattern: pattern,
prefix: normalizeChannelPricingModelName(prefix),
wildcard: isWild,
}
}
// validateNoConflictingModels 检查定价列表中是否有冲突模型模式(同一平台下)。
// 冲突包括:精确重复、通配符之间的前缀包含、通配符与精确名的前缀匹配。
func validateNoConflictingModels(pricingList []ChannelModelPricing) error {
byPlatform := make(map[string][]modelEntry)
for _, p := range pricingList {
for _, model := range p.Models {
byPlatform[p.Platform] = append(byPlatform[p.Platform], toModelEntry(model))
byPlatform[p.Platform] = append(byPlatform[p.Platform], toPricingModelEntry(model))
}
}
for platform, entries := range byPlatform {
@@ -414,6 +414,44 @@ func TestValidateNoConflictingModels(t *testing.T) {
wantErr: true,
errContains: "conflict",
},
// 以下三例:冲突检测必须与 normalizeChannelPricingModelName 用同一套归一化,
// 否则校验放行、写进缓存后键相同,后写的定价会静默覆盖前一条。
{
name: "claude_dot_and_hyphen_spelling_conflict",
pricingList: []ChannelModelPricing{
{Platform: "anthropic", Models: []string{"claude-sonnet-4.5"}},
{Platform: "anthropic", Models: []string{"claude-sonnet-4-5"}},
},
wantErr: true,
errContains: "conflict",
},
{
name: "claude_dot_and_hyphen_spelling_conflict_wildcard",
pricingList: []ChannelModelPricing{
{Platform: "anthropic", Models: []string{"claude-sonnet-4.5*"}},
{Platform: "anthropic", Models: []string{"claude-sonnet-4-5-x"}},
},
wantErr: true,
errContains: "conflict",
},
{
name: "surrounding_whitespace_conflict",
pricingList: []ChannelModelPricing{
{Platform: "openai", Models: []string{"gpt-5.6"}},
{Platform: "openai", Models: []string{" gpt-5.6 "}},
},
wantErr: true,
errContains: "conflict",
},
{
// 只有 claude-* 前缀才做 "." → "-",别把其它平台也一起归一化了
name: "non_claude_dot_spelling_is_not_normalized",
pricingList: []ChannelModelPricing{
{Platform: "openai", Models: []string{"gpt-5.6"}},
{Platform: "openai", Models: []string{"gpt-5-6"}},
},
wantErr: false,
},
}
for _, tt := range tests {
@@ -469,6 +507,16 @@ func TestValidateNoConflictingMappings(t *testing.T) {
wantErr: true,
errContains: "conflict",
},
{
// 映射缓存(expandMappingToCache)只做 strings.ToLower,不做定价那套
// "." → "-",所以这两个源模式在缓存里是两个不同的键、并不冲突。
// 这条用来卡住:定价侧的归一化修复不能顺手套到映射侧,否则会误报冲突。
name: "mapping keeps dot and hyphen spelling separate",
mapping: map[string]map[string]string{
"anthropic": {"claude-sonnet-4.5": "a", "claude-sonnet-4-5": "b"},
},
wantErr: false,
},
{
name: "wildcard vs exact conflict",
mapping: map[string]map[string]string{
@@ -0,0 +1,279 @@
package service
import (
"context"
"fmt"
"log"
"strings"
"sync"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
)
// cnQuotaProber 抽象额度探测(*CNProviderQuotaService 实现,测试可替换)。
type cnQuotaProber interface {
QueryUsage(ctx context.Context, accountID int64) (*CNProviderQuotaProbeResult, error)
}
// cnQuotaProbeConcurrency 周期任务并发探测额度账号的并发度。
const cnQuotaProbeConcurrency = 4
// CNProviderBalanceCheckService 周期性探测国产供应商账号:
// - payg(按量付费):余额低于阈值则临时停调,恢复则清除(仅清除本服务写入的停调);
// - coding plan:调用 CNProviderQuotaService 探测 5h/weekly 滚动窗口并落 extra 快照,
// 调度阈值评估(cnProviderThresholdCandidates)据此自动停调/恢复。
//
// 克隆自 AccountExpiryService 的 Start/Stop/runOnce + ticker 骨架。
// 余额探测仅覆盖有公开余额端点的 kimi / deepseek;智谱无余额端点,仅靠响应式 429/402。
// 额度探测覆盖 kimi / zhipu 的 coding plan 账号(deepseek 无 coding 套餐)。
type CNProviderBalanceCheckService struct {
accountRepo AccountRepository
balanceService *CNProviderBalanceService
quotaService cnQuotaProber
cfg *config.Config
interval time.Duration
stopCh chan struct{}
stopOnce sync.Once
wg sync.WaitGroup
}
// NewCNProviderBalanceCheckService 构造周期余额/额度检测服务。
// interval <= 0 时 Start() 直接返回(不启动),便于通过配置关闭。
func NewCNProviderBalanceCheckService(
accountRepo AccountRepository,
balanceService *CNProviderBalanceService,
quotaService *CNProviderQuotaService,
cfg *config.Config,
interval time.Duration,
) *CNProviderBalanceCheckService {
return &CNProviderBalanceCheckService{
accountRepo: accountRepo,
balanceService: balanceService,
quotaService: quotaService,
cfg: cfg,
interval: interval,
stopCh: make(chan struct{}),
}
}
func (s *CNProviderBalanceCheckService) Start() {
if s == nil || s.accountRepo == nil || s.balanceService == nil || s.cfg == nil {
return
}
if !s.cfg.Gateway.CNProviders.BalanceCheckEnabled {
return
}
if s.interval <= 0 {
return
}
log.Printf("[CNBalance] started (interval=%s threshold=%.2f)", s.interval, s.cfg.Gateway.CNProviders.BalanceThreshold)
s.wg.Add(1)
go func() {
defer s.wg.Done()
ticker := time.NewTicker(s.interval)
defer ticker.Stop()
// 启动后先等待一个周期再首次探测,避免与进程启动峰重叠。
for {
select {
case <-ticker.C:
s.runOnce()
case <-s.stopCh:
return
}
}
}()
}
func (s *CNProviderBalanceCheckService) Stop() {
if s == nil {
return
}
s.stopOnce.Do(func() {
close(s.stopCh)
})
s.wg.Wait()
}
func (s *CNProviderBalanceCheckService) runOnce() {
// 收集 coding 探测目标(kimi/deepseek + 智谱)与 payg 检查队列。
// coding 探测统一在收集完成后按 4 并发执行:单账号探测 15-20s,串行 ×
// 多账号会耗尽整体预算(120s 上限),排在后面的账号快照会饥饿,
// 连锁影响阈值停调的新鲜度判定。
type quotaTarget struct {
id int64
platform string
}
var quotaTargets []quotaTarget
var paygTargets []*Account
collect := func(platform string, accounts []Account) {
for i := range accounts {
account := &accounts[i]
if !account.IsActive() {
continue
}
// coding 账号:探测滚动窗口并落快照(不要求 Schedulable——已被
// 阈值停调的账号也需要新鲜快照决定是否续停)。
if account.IsCodingPlan() {
quotaTargets = append(quotaTargets, quotaTarget{id: account.ID, platform: account.Platform})
continue
}
// payg 余额探测仅 kimi/deepseek(智谱无公开余额端点,payg 账号
// 依赖响应式 402/429 处理)。
if platform != PlatformZhipu && account.Schedulable {
paygTargets = append(paygTargets, account)
}
}
}
for _, platform := range s.platforms() {
accounts, err := s.accountRepo.ListByPlatform(context.Background(), platform)
if err != nil {
log.Printf("[CNBalance] list %s accounts failed: %v", platform, err)
continue
}
collect(platform, accounts)
}
// 智谱无余额端点,仅进额度探测。
if s.quotaService != nil {
accounts, err := s.accountRepo.ListByPlatform(context.Background(), PlatformZhipu)
if err != nil {
log.Printf("[CNBalance] list %s accounts failed: %v", PlatformZhipu, err)
} else {
collect(PlatformZhipu, accounts)
}
}
// 预算按工作量放大:4 并发 × 15s/批 + payg 每账号 5s,下限 30s 上限 300s。
batches := (len(quotaTargets) + cnQuotaProbeConcurrency - 1) / cnQuotaProbeConcurrency
timeout := 30*time.Second + time.Duration(batches)*15*time.Second + time.Duration(len(paygTargets))*5*time.Second
if timeout > 300*time.Second {
timeout = 300 * time.Second
}
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
threshold := s.cfg.Gateway.CNProviders.BalanceThreshold
paused, cleared := 0, 0
for _, account := range paygTargets {
switch s.checkOne(ctx, account, threshold) {
case cnBalancePaused:
paused++
case cnBalanceCleared:
cleared++
}
}
if len(quotaTargets) > 0 && s.quotaService != nil {
sem := make(chan struct{}, cnQuotaProbeConcurrency)
var wg sync.WaitGroup
for _, target := range quotaTargets {
wg.Add(1)
go func(t quotaTarget) {
defer wg.Done()
sem <- struct{}{}
defer func() { <-sem }()
s.probeQuota(ctx, t.id, t.platform)
}(target)
}
wg.Wait()
}
if paused > 0 || cleared > 0 {
log.Printf("[CNBalance] paused=%d cleared=%d (threshold=%.2f)", paused, cleared, threshold)
}
}
// probeQuota 探测单个 coding plan 账号的滚动窗口用量并落 extra 快照。
// 不在此处做停调/恢复决策:调度阈值评估读取快照统一判定(含暂停账号的续停)。
func (s *CNProviderBalanceCheckService) probeQuota(ctx context.Context, accountID int64, platform string) {
if s.quotaService == nil {
return
}
result, err := s.quotaService.QueryUsage(ctx, accountID)
if err != nil {
log.Printf("[CNBalance] quota probe account %d (%s) failed: %v", accountID, platform, err)
return
}
if result != nil && !result.Success && result.Error != "" {
log.Printf("[CNBalance] quota probe account %d (%s) error: %s", accountID, platform, result.Error)
}
}
type cnBalanceCheckOutcome int
const (
cnBalanceNoChange cnBalanceCheckOutcome = iota
cnBalancePaused
cnBalanceCleared
)
// checkOne 探测单账号余额并决定停调/恢复。探测失败时不动现状(避免瞬时网络抖动
// 误解除或误停调)。
func (s *CNProviderBalanceCheckService) checkOne(ctx context.Context, account *Account, threshold float64) cnBalanceCheckOutcome {
result, err := s.balanceService.QueryBalance(ctx, account.ID)
if err != nil || result == nil || !result.Success {
return cnBalanceNoChange
}
// 双币种(deepseek CNY+USD)任一币种余额达标即可继续调度;仅当全部低于
// 阈值(或不可用)才停调。
low := !result.Available || allCNBalancesBelowThreshold(result, threshold)
if low {
// 已被(任何来源)停调时不覆盖其 reason。
if !account.IsSchedulable() {
return cnBalanceNoChange
}
reason := cnBalanceLowReason(fmt.Sprintf("余额 %.4g %s 低于阈值 %.2f", result.Balance, result.Currency, threshold))
if err := s.accountRepo.SetTempUnschedulable(ctx, account.ID, time.Now().Add(s.cooldown()), reason); err != nil {
log.Printf("[CNBalance] pause account %d failed: %v", account.ID, err)
return cnBalanceNoChange
}
log.Printf("[CNBalance] paused account %d (%s): balance=%.4g %s", account.ID, account.Platform, result.Balance, result.Currency)
return cnBalancePaused
}
// 余额健康:仅清除「本服务写入」的临时停调(reason 前缀匹配),不触碰其他子系统。
if account.TempUnschedulableUntil != nil && strings.HasPrefix(account.TempUnschedulableReason, cnBalanceLowReasonPrefix) {
if err := s.accountRepo.ClearTempUnschedulable(ctx, account.ID); err != nil {
log.Printf("[CNBalance] clear account %d failed: %v", account.ID, err)
return cnBalanceNoChange
}
log.Printf("[CNBalance] reactivated account %d (%s): balance=%.4g %s", account.ID, account.Platform, result.Balance, result.Currency)
return cnBalanceCleared
}
return cnBalanceNoChange
}
func (s *CNProviderBalanceCheckService) platforms() []string {
return []string{PlatformKimi, PlatformDeepseek}
}
// allCNBalancesBelowThreshold 判断全部币种余额是否均低于阈值。
// 无明细时退回主币种判定(与旧行为一致)。
func allCNBalancesBelowThreshold(result *CNProviderBalanceResult, threshold float64) bool {
if len(result.Balances) == 0 {
return result.Balance < threshold
}
for _, entry := range result.Balances {
if entry.Balance >= threshold {
return false
}
}
return true
}
// cooldown 返回临时停调持续时长(= 2× 检测周期),与响应式 402/429 路径一致。
func (s *CNProviderBalanceCheckService) cooldown() time.Duration {
minutes := 10
if s.cfg != nil {
if cfgMin := s.cfg.Gateway.CNProviders.BalanceCheckIntervalMinutes; cfgMin > 0 {
minutes = cfgMin
}
}
cooldown := time.Duration(minutes) * time.Minute * 2
if cooldown < time.Minute {
cooldown = 10 * time.Minute
}
return cooldown
}
@@ -0,0 +1,101 @@
package service
import (
"context"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
// 周期任务对 coding plan 账号的额度探测行为(runOnce 集成路径):
// - kimi coding 账号(含已被阈值停调的)→ 额度探测被调用;
// - 智谱 coding 账号 → 额度探测被调用(智谱不进 kimi/deepseek 余额循环);
// - payg 账号不经过额度探测(走余额路径,本测试不放 payg 账号避免真实网络);
// - 非激活账号完全跳过。
type fakeCNQuotaProber struct {
probed []int64
}
func (f *fakeCNQuotaProber) QueryUsage(ctx context.Context, accountID int64) (*CNProviderQuotaProbeResult, error) {
f.probed = append(f.probed, accountID)
return &CNProviderQuotaProbeResult{Success: true, Persisted: true}, nil
}
type fakeCNCheckRepo struct {
AccountRepository
byPlatform map[string][]Account
}
func (r *fakeCNCheckRepo) ListByPlatform(ctx context.Context, platform string) ([]Account, error) {
return r.byPlatform[platform], nil
}
func TestCNProviderBalanceCheckRunOnceProbesCodingPlanQuota(t *testing.T) {
kimiActive := Account{ID: 1, Platform: PlatformKimi, Type: AccountTypeAPIKey, Status: StatusActive,
Credentials: map[string]any{"account_mode": "coding"}}
// 已被阈值停调的 coding 账号也要刷新快照(决定是否续停)。
kimiPaused := Account{ID: 2, Platform: PlatformKimi, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: false,
Credentials: map[string]any{"account_mode": "coding"}}
// 非激活账号跳过。
kimiInactive := Account{ID: 3, Platform: PlatformKimi, Type: AccountTypeAPIKey, Status: StatusDisabled,
Credentials: map[string]any{"account_mode": "coding"}}
zhipuCoding := Account{ID: 4, Platform: PlatformZhipu, Type: AccountTypeAPIKey, Status: StatusActive,
Credentials: map[string]any{"account_mode": "coding"}}
repo := &fakeCNCheckRepo{byPlatform: map[string][]Account{
PlatformKimi: {kimiActive, kimiPaused, kimiInactive},
PlatformZhipu: {zhipuCoding},
}}
prober := &fakeCNQuotaProber{}
svc := &CNProviderBalanceCheckService{
accountRepo: repo,
quotaService: prober,
cfg: &config.Config{},
}
svc.runOnce()
require.ElementsMatch(t, []int64{1, 2, 4}, prober.probed)
}
// runOnceZhipuQuota 在 quotaService 缺失时安全跳过(Start 门控不启动的老部署路径)。
func TestCNProviderBalanceCheckRunOnceWithoutQuotaService(t *testing.T) {
repo := &fakeCNCheckRepo{byPlatform: map[string][]Account{
PlatformZhipu: {{ID: 4, Platform: PlatformZhipu, Type: AccountTypeAPIKey, Status: StatusActive,
Credentials: map[string]any{"account_mode": "coding"}}},
}}
svc := &CNProviderBalanceCheckService{accountRepo: repo, cfg: &config.Config{}}
require.NotPanics(t, func() { svc.runOnce() })
}
// 双币种(deepseek CNY+USD)停调判定:任一币种达标即不停调,全部低于阈值才停;
// 无明细时退回主币种(兼容旧结果)。
func TestAllCNBalancesBelowThreshold(t *testing.T) {
dualLow := &CNProviderBalanceResult{
Balance: 1.0,
Currency: "CNY",
Balances: []CNProviderBalanceEntry{
{Currency: "CNY", Balance: 1.0},
{Currency: "USD", Balance: 0.5},
},
}
require.True(t, allCNBalancesBelowThreshold(dualLow, 5.0))
dualMixed := &CNProviderBalanceResult{
Balance: 1.0,
Currency: "CNY",
Balances: []CNProviderBalanceEntry{
{Currency: "CNY", Balance: 1.0},
{Currency: "USD", Balance: 20.0},
},
}
require.False(t, allCNBalancesBelowThreshold(dualMixed, 5.0))
// 无明细:按主币种判定(旧行为)。
singleLow := &CNProviderBalanceResult{Balance: 1.0, Currency: "CNY"}
require.True(t, allCNBalancesBelowThreshold(singleLow, 5.0))
singleOK := &CNProviderBalanceResult{Balance: 10.0, Currency: "CNY"}
require.False(t, allCNBalancesBelowThreshold(singleOK, 5.0))
}

Some files were not shown because too many files have changed in this diff Show More