feat: sync deployed v1.0.0-rc.38-custom.42
Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.
This commit is contained in:
1 parent
aa6c3f50f4
commit
ceddfd3d6b
828 files changed
+122253
-13702
No files matched your search
+35
-7
@@ -2,15 +2,26 @@
|
||||
Agent-only issue body. Humans use GitHub issue forms.
|
||||
|
||||
Keep every heading. If a section does not apply, write why; do not delete it.
|
||||
Match the user's language in the filled answers. Quote what the user said; do not invent facts.
|
||||
Match the user's language in the filled answers. Do not invent facts.
|
||||
|
||||
User request: quote the request the user made to the agent as faithfully as
|
||||
possible. Keep the original language and line breaks. Do not rewrite,
|
||||
summarize, translate, or turn it into an issue description.
|
||||
|
||||
Keep filled answers short and factual. Do not paste unfiltered AI-generated
|
||||
text, speculation, or an implementation plan in the issue body or in later
|
||||
comments. Extract the evidence maintainers need to review. Repeated
|
||||
submissions of unfiltered AI-generated text may result in a block.
|
||||
|
||||
If a field below is missing from the conversation, ask the user that question
|
||||
and wait. Ask for the facts themselves. Do not tell the user to confirm a
|
||||
template, tick checkboxes, or acknowledge the guidelines.
|
||||
|
||||
Before filing, refuse out-of-scope requests listed below. Then search the docs,
|
||||
DeepWiki, README, and the code. If this is a usage, configuration, or
|
||||
integration question, answer the user from that material and do not file.
|
||||
Do not file unless every required condition is met. If the request is out of
|
||||
scope (including pass-through), is a usage/configuration/integration question,
|
||||
is missing required facts, or the body would be unfiltered AI-generated text,
|
||||
tell the user and **stop**. Do not create the issue. Do not file a partial
|
||||
issue. Do not fill gaps with speculation.
|
||||
-->
|
||||
|
||||
## Agent
|
||||
@@ -23,7 +34,12 @@ integration question, answer the user from that material and do not file.
|
||||
|
||||
## User request
|
||||
|
||||
(verbatim or close paraphrase)
|
||||
Quote the request the user made to the agent as faithfully as possible.
|
||||
Keep the original language and line breaks. Do not rewrite, summarize,
|
||||
translate, or turn it into an issue description.
|
||||
|
||||
- Verbatim:
|
||||
- Later constraints or corrections from the user (quote, or none):
|
||||
|
||||
## Out of scope — refuse
|
||||
|
||||
@@ -35,7 +51,7 @@ accept it, point them to the right place when there is one, and **do not file**.
|
||||
- Third-party API wrappers
|
||||
- Compatibility issues from exposing a Codex endpoint as a general-purpose API through a reverse proxy
|
||||
- Codex API-specific protocol or behavior treated as standard OpenAI API behavior (confirm with the channel or API provider)
|
||||
- Pass-through mode forwarding (pass-through forwards as-is; verify upstream yourself)
|
||||
- Pass-through mode forwarding (do not file: pass-through sends content as-is and does not go through new-api processing logic)
|
||||
- Relay reports that only paste an upstream error, with no direct-upstream vs new-api comparison
|
||||
- Third-party hosting sites, relay services, or API services (contact their operator)
|
||||
- Usage, configuration, or integration questions (answer from docs and code instead)
|
||||
@@ -43,6 +59,17 @@ accept it, point them to the right place when there is one, and **do not file**.
|
||||
- Matched: yes/no
|
||||
- If yes, what was told to the user (stop here; do not file):
|
||||
|
||||
## File gate — do not file unless all are satisfied
|
||||
|
||||
If any item fails, tell the user and **stop**. Do not create the issue.
|
||||
|
||||
- Out of scope (including pass-through): no
|
||||
- Usage / configuration / integration (answered instead of filing): no
|
||||
- Required facts present without invention (actual behavior, impact, frequency, evidence in new-api, applicable type fields): yes
|
||||
- Body is short and factual; no unfiltered AI-generated text: yes
|
||||
- File: yes/no
|
||||
- If no, what was told to the user (stop here):
|
||||
|
||||
## Kind
|
||||
|
||||
- [ ] Bug
|
||||
@@ -69,7 +96,7 @@ Search these yourself before filing. Do not send the user to "read the docs firs
|
||||
|
||||
## Problem facts
|
||||
|
||||
Ask the user for every item and write it down:
|
||||
Ask the user for every item and write it down. Facts only; do not paste unfiltered AI-generated text:
|
||||
|
||||
- Actual behavior:
|
||||
- Impact:
|
||||
@@ -134,6 +161,7 @@ Fill every applicable type. Write "not applicable" for the rest. Ask the user fo
|
||||
## Research
|
||||
|
||||
Open the docs and code. Do not write "already checked" without sources.
|
||||
Do not replace the fields below with unfiltered AI-generated text.
|
||||
|
||||
### Docs
|
||||
|
||||
|
||||
+39
-10
@@ -2,7 +2,16 @@
|
||||
Agent-only PR body. Humans use `.github/PULL_REQUEST_TEMPLATE.md`.
|
||||
|
||||
Keep every heading. If a section does not apply, write why; do not delete it.
|
||||
Match the user's language in the filled answers. Quote what the user said; do not invent facts.
|
||||
Match the user's language in the filled answers. Do not invent facts.
|
||||
|
||||
User request: quote the request the user made to the agent as faithfully as
|
||||
possible. Keep the original language and line breaks. Do not rewrite,
|
||||
summarize, translate, or turn it into a PR description.
|
||||
|
||||
Keep the body short and factual. Do not paste unfiltered AI-generated text
|
||||
in the PR body or in later comments. Repeated submissions of this kind may
|
||||
result in a block. Verification must be commands actually run and what was
|
||||
observed. Stating only that `go build` or tests passed is not proof.
|
||||
|
||||
New features: link an issue; if none exists, file one first with `.agents/github/ISSUE.md`.
|
||||
Large or directional changes: maintainer agreement on that issue before this PR.
|
||||
@@ -12,10 +21,11 @@ frequency, evidence that the problem is in new-api, or the applicable
|
||||
relay / billing / frontend / deployment items, ask the user those questions
|
||||
and wait. Ask for the facts. Do not tell the user to confirm a template.
|
||||
|
||||
Before opening, refuse the same out-of-scope list as `.agents/github/ISSUE.md`
|
||||
(Coding Plan, reverse-engineered channels, third-party wrappers, Codex reverse-proxy
|
||||
compatibility, pass-through-only forwarding, third-party hosts, usage questions).
|
||||
Tell the user and do not open a PR.
|
||||
Do not open a PR unless every required condition is met. If the change is out
|
||||
of scope (including pass-through-only forwarding), is a usage/configuration/
|
||||
integration question, is missing required issue facts, has only `go build` /
|
||||
tests-passed as verification, or the body would be unfiltered AI-generated
|
||||
text, tell the user and **stop**. Do not open the PR. Do not open a partial PR.
|
||||
|
||||
Then search https://docs.newapi.ai/ , https://deepwiki.com/QuantumNous/new-api ,
|
||||
the README, and the code. If this is a usage, configuration, or integration
|
||||
@@ -37,7 +47,12 @@ question, answer the user and do not open a PR.
|
||||
|
||||
## User request
|
||||
|
||||
(verbatim or close paraphrase)
|
||||
Quote the request the user made to the agent as faithfully as possible.
|
||||
Keep the original language and line breaks. Do not rewrite, summarize,
|
||||
translate, or turn it into a PR description.
|
||||
|
||||
- Verbatim:
|
||||
- Later constraints or corrections from the user (quote, or none):
|
||||
|
||||
## Out of scope — refuse
|
||||
|
||||
@@ -49,13 +64,25 @@ accept it and **do not open a PR**.
|
||||
- Third-party API wrappers
|
||||
- Codex channel-type changes, or compatibility from exposing Codex as a general-purpose API
|
||||
- Codex API-specific protocol or behavior treated as standard OpenAI API behavior
|
||||
- Pass-through-only forwarding
|
||||
- Pass-through-only forwarding (pass-through sends content as-is and does not go through new-api processing logic)
|
||||
- Third-party hosting sites, relay services, or API services
|
||||
- Usage, configuration, or integration (answer from docs and code instead)
|
||||
|
||||
- Matched: yes/no
|
||||
- If yes, what was told to the user (stop here; do not open a PR):
|
||||
|
||||
## Open gate — do not open unless all are satisfied
|
||||
|
||||
If any item fails, tell the user and **stop**. Do not open the PR.
|
||||
|
||||
- Out of scope (including pass-through-only forwarding): no
|
||||
- Usage / configuration / integration (answered instead of opening): no
|
||||
- Required issue facts present without invention: yes
|
||||
- Verification is actual commands or steps and observed results, not only `go build` or tests passed: yes
|
||||
- Body is short and factual; no unfiltered AI-generated text: yes
|
||||
- Open: yes/no
|
||||
- If no, what was told to the user (stop here):
|
||||
|
||||
## Kind
|
||||
|
||||
- [ ] Bug fix
|
||||
@@ -76,7 +103,8 @@ Take these from the linked issue. If a needed item is empty, ask the user that q
|
||||
|
||||
## Change
|
||||
|
||||
(what changed, why it works, grounded in the code actually touched)
|
||||
(what changed, why it works, grounded in the code actually touched.
|
||||
Short and factual; do not paste unfiltered AI-generated text.)
|
||||
|
||||
## Research
|
||||
|
||||
@@ -114,7 +142,8 @@ Open them. Do not write "already checked" without sources.
|
||||
|
||||
## Verification
|
||||
|
||||
Only what was actually run.
|
||||
Only what was actually run. Do not state only that `go build` or tests passed.
|
||||
Each item needs the command or steps and the observed result.
|
||||
|
||||
- Commands and results:
|
||||
- Manual steps and observed result:
|
||||
@@ -133,4 +162,4 @@ Only what was actually run.
|
||||
|
||||
- Single focused change: yes/no (if no, why):
|
||||
- Secrets included: no
|
||||
- Out of scope (Coding Plan / reverse-engineered channel / third-party wrapper / Codex): no
|
||||
- Out of scope (Coding Plan / reverse-engineered channel / third-party wrapper / Codex / pass-through-only forwarding): no
|
||||
@@ -0,0 +1,28 @@
|
||||
# Billing Rules
|
||||
|
||||
Mandatory conventions for quota, pricing, pre-consume, settlement, refund, and usage-quantity code in the Go backend. `AGENTS.md` defines when a task touches billing (see **Billing rules (mandatory read gate)**); any task that matches MUST read this whole file before planning, coding, or reviewing. Tasks that match none of those triggers may skip it. Rules use the same MUST / Do NOT wording as `AGENTS.md` and carry the same weight.
|
||||
|
||||
**Billing expression system:** When working on tiered/dynamic billing (expression-based pricing), MUST read `pkg/billingexpr/expr.md` first. It documents the design philosophy, expression language, full architecture, token normalization rules, quota conversion, and expression versioning. All billing expression changes must follow that document.
|
||||
|
||||
**Built-in model pricing:** New built-in model prices MUST be defined as self-contained billing expressions in `setting/billing_setting/builtin_billing.go`, using real USD per million tokens. Do not add new built-in prices to the legacy model/completion/cache ratio tables. Preserve explicit administrator pricing overrides. Existing legacy prices are migrated only when explicitly requested. Verify published prices and cover applicable context-length thresholds and cache categories.
|
||||
|
||||
**Billing safety invariants:** Quota/billing code MUST never produce a negative charge (a credit) from arithmetic overflow or unvalidated input. Apply defense in depth:
|
||||
|
||||
- Every user-controlled quantity that becomes a billing multiplier (image `n`, video `seconds`/`duration`, resolution/quality ratios, batch counts) MUST be bounded before it reaches quota calculation. Reject out-of-range values at request validation with a 400. Existing bounds: `dto.MaxImageN` for image generation count, `relaycommon.MaxTaskDurationSeconds` for task video duration, `maxTokensLimit` (`relay/helper/valid_request.go`) for `max_tokens`-family fields on every relay format (OpenAI, Claude, Gemini, Responses). Reuse these constants instead of introducing new ad hoc limits for the same concepts. When adding a new relay format or request DTO, bound its max-tokens and count fields in its validator from day one.
|
||||
- Watch for validation bypass paths: passthrough fields (e.g. `Extra["parameters"]`), task `metadata` maps, and multipart form fields can carry the same quantities around the standard DTO validation. Any adaptor that reads a multiplier from such a path must enforce the same bound (or clamp) locally.
|
||||
- Durations parsed from media metadata are user/upstream-controlled too: audio file headers (transcription token counting, TTS response duration) and upstream deduction numbers (e.g. Kling `FinalUnitDeduction`) can claim absurd values. Convert them with saturation before they become token counts.
|
||||
- Never convert a computed quota or token count to `int` with a bare cast like `int(float64(quota) * ratio)`, `int(math.Round(...))` on unbounded input, or `int(decimal.IntPart())`. All quota rounding/conversion is centralized in `common/quota_math.go`; use those helpers: `common.QuotaFromFloat` (truncating) for float products, `common.QuotaRound` (half-away-from-zero) where rounding is intended, and `common.QuotaFromDecimal` for decimal products. `billingexpr.QuotaRound` delegates to `common.QuotaRound`. Do not reintroduce local conversion helpers or bare casts. Single-request saturation stays at the int32 boundary so batch accumulation cannot approach 64-bit wraparound; wallet/top-up conversion uses `common.WalletQuotaFromDecimalStrict` with the JavaScript-safe `common.MaxWalletQuota` boundary. Every clamp/NaN fallback is logged via `common.SysError`.
|
||||
- Saturation events are also audited: each helper has a `*Checked` variant (`common.QuotaFromFloatChecked` / `QuotaRoundChecked` / `QuotaFromDecimalChecked`) that additionally returns a `*common.QuotaClamp` when clamping occurred. Billing paths that compute a charge capture that clamp onto `relayInfo.QuotaClamp` (or thread it into task settlement) and, right before writing the consume/task log, call `attachQuotaSaturation` (in `service/log_info_generate.go`) which nests the marker under the log's `other.admin_info.quota_saturation` and emits a request-correlated `logger.LogWarn`. Nesting under `admin_info` makes it admin-only for free (non-admin log views strip `admin_info`). When adding a new billing path, use the `*Checked` variant and surface the clamp the same way so the anomaly stays auditable in both the admin log UI and backend logs.
|
||||
- Multiplier maps go through `types.PriceData.AddOtherRatio`, which rejects non-positive, NaN, and +Inf ratios. Do not write to `PriceData.OtherRatios` directly, and do not weaken these guards.
|
||||
- Pre-consume (预扣费) and settle (结算/差额) must both be safe: a saturated oversized quota must fail pre-consume with insufficient-quota, never silently wrap. When adding a new billing path (new relay format, new task platform, new adjustment hook), trace the full chain — validation → EstimateBilling/OtherRatios → quota conversion → pre-consume → settle/refund — and confirm each step preserves these invariants.
|
||||
- Fields parsed into unsigned types (`*uint`) accept huge positive JSON numbers (e.g. `18446744073686646784`, a wrapped negative); a `>= 0` check is not sufficient, an upper bound is mandatory.
|
||||
- Regression tests for these invariants belong with the boundary they protect (request validators, converter helpers). See `relay/helper/openai_image_request_test.go`, `relay/common/relay_utils_test.go`, and `common/quota_math_test.go` for the expected style.
|
||||
|
||||
**Upstream image quantity settlement:** The billable image count may be lowered or raised from the upstream response only through `RelayInfo.UpdateImageCount`; a zero or out-of-range count leaves the reserved request quantity in place. Derive the count from image payloads, never from JSON structure:
|
||||
|
||||
- Do NOT count `data` with gjson `data.#`, `len(data)`, or any element count. Providers behind OpenAI-compatible channels return non-standard shapes: `data` as a single object, one image split into a `url` entry and a `b64_json` entry, or entries carrying only `revised_prompt`.
|
||||
- An object-shaped `data` is one image when it has a non-empty `url` or `b64_json`, otherwise zero. For arrays the count is `max(entries with non-empty url, entries with non-empty b64_json)`: a standard response uses one `response_format` so this equals the entry count, a split image bills once, and payload-less entries bill nothing. The reference implementation is `openaiImageResponseCount` in `relay/channel/openai/relay_image.go`; reuse it (or the same rule) for any channel that forwards or rebuilds OpenAI-format image responses. The xAI adaptor already inherits it through `OpenaiImageHandler`.
|
||||
- When an OpenAI-format JSON body is re-emitted as SSE, forward one `image_generation.completed` event per image payload using the same shape handling (object → one event, payload-less entries dropped). The client must never receive only `[DONE]` while being charged for images.
|
||||
- Real SSE streams count `image_generation.completed` / `image_edit.completed` events and only trust a lower count when upstream finished the stream (`done` / `eof`); a client abort must not reduce the charge below the requested quantity.
|
||||
- Channels with their own usage field (Ali `usage.image_count`) use that value, bounded by `dto.MaxImageN`, and fall back to the converted response length. Channels that rebuild responses from typed structs count the typed slice; none of them may read `data.#`.
|
||||
- Regression cases for object, split, and payload-less shapes belong in `relay/channel/openai/image_stream_test.go` (fixed-price, tiered-expression, and JSON-as-SSE forwarding tables). Extend those tables instead of adding new files.
|
||||
@@ -9,14 +9,23 @@ body:
|
||||
<!--
|
||||
If you are an AI coding agent: do not fill this GitHub form. Create the issue with the filled body from `.agents/github/ISSUE.md`.
|
||||
-->
|
||||
## 提交前必读(请勿删除本节)
|
||||
|
||||
- type: textarea
|
||||
id: read-first
|
||||
attributes:
|
||||
label: 提交前必读
|
||||
description: GitHub 的 markdown 说明提交后不会进入 issue。本节会写入正文,请勿删除或清空。
|
||||
value: |
|
||||
- 文档:https://docs.newapi.ai/
|
||||
- 使用问题先看或先问:https://deepwiki.com/QuantumNous/new-api
|
||||
- 转发类问题必须先确认问题发生在 new-api 这一层:请使用等价的脱敏请求分别直连上游和经过 new-api,对比响应状态、错误体与服务端日志。透传模式或未经对比、仅展示上游报错的反馈不接受 issue。
|
||||
- 开启透传请勿反馈问题;透传会直接发送内容,不会经过 new-api 的逻辑处理。
|
||||
- 请以简练语言陈述事实与证据。Issue 正文及后续评论均不得直接粘贴未经人工过滤的大段 AI 文本。
|
||||
- 转发类问题必须先确认问题发生在 new-api 这一层:请使用等价的脱敏请求分别直连上游和经过 new-api,对比响应状态、错误体与服务端日志。未经对比、仅展示上游报错的反馈不接受 issue。
|
||||
- 不接受 Coding Plan、逆向渠道、第三方封装接口,以及将 Codex 接口反代为通用 API 后产生的兼容性问题;Codex API 自身特有的协议或行为也不应被当作标准 OpenAI API 行为,相关问题请先向渠道或接口提供方确认。
|
||||
- 本项目不提供官方托管服务,也不为第三方托管站、中转站或 API 服务提供技术支持;此类问题请联系对应运营方。
|
||||
- 警告:删除本模板、删除小节标题或随意清空内容的 issue,可能会被直接关闭;重复恶意提交者可能会被 block。
|
||||
- 警告:删除本模板、删除小节标题或随意清空内容的 issue,可能会被直接关闭;多次提交未经人工过滤的大段 AI 文本,或重复恶意提交者,可能会被 block。
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: dropdown
|
||||
id: deployment-source
|
||||
@@ -52,11 +61,13 @@ body:
|
||||
required: true
|
||||
- label: "**非第三方服务:** 我确认这不是第三方托管站、中转站、API 服务或未经本仓库原版验证的 fork 所独有的问题;第三方实例问题应联系其运营方。"
|
||||
required: true
|
||||
- label: "**非透传:** 我确认相关渠道未开启透传。开启透传会直接发送内容,不会经过 new-api 的逻辑处理,此类问题不接受 issue。"
|
||||
required: true
|
||||
- label: "**问题归属:** 我已提供足以区分客户端、new-api 和上游的证据;转发问题已使用等价的脱敏请求对比直连上游与经过 new-api 的结果,确认错误由 new-api 引入或改变,而非仅原样转发上游报错。"
|
||||
required: true
|
||||
- label: "**渠道与协议边界:** 我确认问题不是由 Coding Plan、逆向渠道、第三方封装接口、Codex 反代接口,或 Codex API 自身特有行为导致;若问题最初出现于此类接口,我已在本仓库支持的标准 API 协议上复现。"
|
||||
required: true
|
||||
- label: "**维护成本:** 我理解项目维护者精力有限,不遵循模板要求的 issue 可能会被无视或直接关闭。"
|
||||
- label: "**维护成本:** 我理解项目维护者精力有限。Issue 正文及后续评论均不得直接粘贴未经人工过滤的大段 AI 文本,须先提炼可复核的要点;不遵循模板要求或缺少要点、无法复核的内容可能会被无视或直接关闭。多次提交此类内容可能会被 block。"
|
||||
required: true
|
||||
|
||||
- type: markdown
|
||||
|
||||
@@ -9,14 +9,23 @@ body:
|
||||
<!--
|
||||
If you are an AI coding agent: do not fill this GitHub form. Create the issue with the filled body from `.agents/github/ISSUE.md`.
|
||||
-->
|
||||
## Read This First (Do Not Remove This Section)
|
||||
|
||||
- type: textarea
|
||||
id: read-first
|
||||
attributes:
|
||||
label: Read This First
|
||||
description: GitHub markdown instructions are not included in the submitted issue. This section is written into the issue body. Do not delete or clear it.
|
||||
value: |
|
||||
- Docs: https://docs.newapi.ai/
|
||||
- Usage questions first: https://deepwiki.com/QuantumNous/new-api
|
||||
- Relay reports must first establish that the issue occurs in the new-api layer. Use equivalent redacted requests to compare the direct upstream result with the result through new-api, including response status, error body, and server logs. Reports involving pass-through mode, or showing only an upstream error without this comparison, are not accepted.
|
||||
- Do not report issues when pass-through is enabled; pass-through sends content as-is and does not go through new-api processing logic.
|
||||
- State facts and evidence concisely. Do not paste unfiltered AI-generated text in the issue body or in later comments.
|
||||
- Relay reports must first establish that the issue occurs in the new-api layer. Use equivalent redacted requests to compare the direct upstream result with the result through new-api, including response status, error body, and server logs. Reports that skip this comparison or only show an upstream error are not accepted.
|
||||
- Reports about Coding Plan services, reverse-engineered channels, third-party API wrappers, or compatibility issues caused by exposing a Codex endpoint through a reverse proxy as a general-purpose API are not accepted. Protocols or behavior specific to the Codex API should not be treated as standard OpenAI API behavior; confirm such issues with the channel or API provider first.
|
||||
- This project does not operate an official hosted service or provide technical support for third-party hosting sites, relay services, or API services. Contact the relevant operator for those issues.
|
||||
- Warning: issues with this template removed, section headings deleted, or content cleared may be closed directly. Repeated abusive submissions may result in a block.
|
||||
- Warning: issues with this template removed, section headings deleted, or content cleared may be closed directly. Repeated submissions of unfiltered AI-generated text, or repeated abusive submissions, may result in a block.
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: dropdown
|
||||
id: deployment-source
|
||||
@@ -52,11 +61,13 @@ body:
|
||||
required: true
|
||||
- label: "**Not a third-party service:** I confirm that this issue is not exclusive to a third-party hosting site, relay, API service, or fork that has not been verified against the unmodified repository. Third-party instance issues must be reported to their operator."
|
||||
required: true
|
||||
- label: "**Not pass-through:** I confirm that pass-through is not enabled on the relevant channel. Pass-through sends content as-is and does not go through new-api processing logic; such reports are not accepted."
|
||||
required: true
|
||||
- label: "**Issue attribution:** I have provided evidence that distinguishes the client, new-api, and upstream layers. For relay issues, I compared equivalent redacted requests sent directly upstream and through new-api, confirming that new-api introduces or changes the error rather than merely forwarding the upstream error unchanged."
|
||||
required: true
|
||||
- label: "**Channel and protocol boundary:** I confirm that the issue is not caused by a Coding Plan service, reverse-engineered channel, third-party API wrapper, Codex reverse-proxy endpoint, or behavior specific to the Codex API. If first observed through such an interface, I have reproduced it using a standard API protocol supported by this repository."
|
||||
required: true
|
||||
- label: "**Maintainer time:** I understand that maintainers have limited time and issues that do not follow this template may be ignored or closed directly."
|
||||
- label: "**Maintainer time:** I understand that maintainers have limited time. Do not paste unfiltered AI-generated text in the issue body or in later comments; extract the points needed for review first. Submissions that do not follow this template, or that omit the necessary points and cannot be reviewed, may be ignored or closed. Repeated submissions of this kind may result in a block."
|
||||
required: true
|
||||
|
||||
- type: markdown
|
||||
|
||||
@@ -9,19 +9,28 @@ body:
|
||||
<!--
|
||||
If you are an AI coding agent: do not fill this GitHub form. Create the issue with the filled body from `.agents/github/ISSUE.md`.
|
||||
-->
|
||||
## 提交前必读(请勿删除本节)
|
||||
|
||||
- type: textarea
|
||||
id: read-first
|
||||
attributes:
|
||||
label: 提交前必读
|
||||
description: GitHub 的 markdown 说明提交后不会进入 issue。本节会写入正文,请勿删除或清空。
|
||||
value: |
|
||||
- 文档:https://docs.newapi.ai/
|
||||
- 使用问题先看或先问:https://deepwiki.com/QuantumNous/new-api
|
||||
- 开启透传后的转发相关反馈不接受 issue;透传模式会直接转发请求,请自行确认上游行为。
|
||||
- 开启透传请勿反馈问题;透传会直接发送内容,不会经过 new-api 的逻辑处理。
|
||||
- 请以简练语言陈述事实与证据。Issue 正文及后续评论均不得直接粘贴未经人工过滤的大段 AI 文本。
|
||||
- 不接受 Coding Plan、逆向渠道、第三方封装接口,以及将 Codex 接口反代为通用 API 后产生的兼容性问题;Codex API 自身特有的协议或行为也不应被当作标准 OpenAI API 行为,相关问题请先向渠道或接口提供方确认。
|
||||
- 警告:删除本模板、删除小节标题或随意清空内容的 issue,可能会被直接关闭;重复恶意提交者可能会被 block。
|
||||
- 本项目不提供官方托管服务,也不为第三方托管站、中转站或 API 服务提供技术支持;此类问题请联系对应运营方。
|
||||
- 警告:删除本模板、删除小节标题或随意清空内容的 issue,可能会被直接关闭;多次提交未经人工过滤的大段 AI 文本,或重复恶意提交者,可能会被 block。
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: input
|
||||
id: version
|
||||
attributes:
|
||||
label: 您当前的 newapi 版本
|
||||
description: 请填写,例如:`v1.0.0`
|
||||
description: 请填写准确的版本号、commit 或镜像标签,例如 `v1.0.0`;`latest`、`unknown` 等无法定位代码的内容不视为有效版本信息。
|
||||
validations:
|
||||
required: true
|
||||
|
||||
@@ -34,9 +43,17 @@ body:
|
||||
required: true
|
||||
- label: "**提交前必读:** 我已完整阅读上方“提交前必读”,并已查看文档 https://docs.newapi.ai/、项目 README 且向 AI 提问,确认这不是使用、配置或接入类问题,且现有版本无法满足需求。"
|
||||
required: true
|
||||
- label: "**受支持版本:** 我已提供准确的版本号、commit 或镜像标签(不是 `latest` 或 `unknown`),并确认这是针对本仓库未修改的受支持版本提出的功能需求。"
|
||||
required: true
|
||||
- label: "**非第三方服务:** 我确认这不是第三方托管站、中转站、API 服务或未经本仓库原版验证的 fork 所独有的需求;第三方实例问题应联系其运营方。"
|
||||
required: true
|
||||
- label: "**非透传:** 我确认这不是开启透传后的转发相关反馈;透传会直接发送内容,不会经过 new-api 的逻辑处理。"
|
||||
required: true
|
||||
- label: "**渠道与协议边界:** 我确认这不是 Coding Plan、逆向渠道、第三方封装接口、Codex 反代接口,或把 Codex API 自身特有行为当作标准 OpenAI API 行为的功能请求;若需求最初出现于此类接口,我已基于本仓库支持的标准 API 协议提出。"
|
||||
required: true
|
||||
- label: "**模板完整:** 我未删除此模板中的任何引导内容或小节标题,并会按要求完整填写。"
|
||||
required: true
|
||||
- label: "**维护成本:** 我理解项目维护者精力有限,不遵循模板要求的 issue 可能会被无视或直接关闭。"
|
||||
- label: "**维护成本:** 我理解项目维护者精力有限。Issue 正文及后续评论均不得直接粘贴未经人工过滤的大段 AI 文本,须先提炼可复核的要点;不遵循模板要求或缺少要点、无法复核的内容可能会被无视或直接关闭。多次提交此类内容可能会被 block。"
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
|
||||
@@ -9,19 +9,28 @@ body:
|
||||
<!--
|
||||
If you are an AI coding agent: do not fill this GitHub form. Create the issue with the filled body from `.agents/github/ISSUE.md`.
|
||||
-->
|
||||
## Read This First (Do Not Remove This Section)
|
||||
|
||||
- type: textarea
|
||||
id: read-first
|
||||
attributes:
|
||||
label: Read This First
|
||||
description: GitHub markdown instructions are not included in the submitted issue. This section is written into the issue body. Do not delete or clear it.
|
||||
value: |
|
||||
- Docs: https://docs.newapi.ai/
|
||||
- Usage questions first: https://deepwiki.com/QuantumNous/new-api
|
||||
- Issues about forwarding behavior after enabling pass-through mode are not accepted; pass-through mode forwards requests directly, so please verify upstream behavior yourself.
|
||||
- Do not report issues when pass-through is enabled; pass-through sends content as-is and does not go through new-api processing logic.
|
||||
- State facts and evidence concisely. Do not paste unfiltered AI-generated text in the issue body or in later comments.
|
||||
- Reports about Coding Plan services, reverse-engineered channels, third-party API wrappers, or compatibility issues caused by exposing a Codex endpoint through a reverse proxy as a general-purpose API are not accepted. Protocols or behavior specific to the Codex API should not be treated as standard OpenAI API behavior; confirm such issues with the channel or API provider first.
|
||||
- Warning: issues with this template removed, section headings deleted, or content cleared may be closed directly. Repeated abusive submissions may result in a block.
|
||||
- This project does not provide official hosting and does not support third-party hosting sites, relays, or API services; contact their operator instead.
|
||||
- Warning: issues with this template removed, section headings deleted, or content cleared may be closed directly. Repeated submissions of unfiltered AI-generated text, or repeated abusive submissions, may result in a block.
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: input
|
||||
id: version
|
||||
attributes:
|
||||
label: Your current newapi version
|
||||
description: "Please fill this in, for example: `v1.0.0`"
|
||||
description: "Provide an exact version, commit, or image tag, for example `v1.0.0`. Values such as `latest` or `unknown` that cannot locate the code are not valid version information."
|
||||
validations:
|
||||
required: true
|
||||
|
||||
@@ -34,9 +43,17 @@ body:
|
||||
required: true
|
||||
- label: "**Read this first:** I have fully read the section above, reviewed the docs at https://docs.newapi.ai/ and the project README, and asked AI first, confirming this is not a usage, configuration, or integration question, and that the current version cannot meet my needs."
|
||||
required: true
|
||||
- label: "**Supported version:** I have provided an exact version, commit, or image tag (not `latest` or `unknown`) and confirm this feature request is for an unmodified, supported version from this repository."
|
||||
required: true
|
||||
- label: "**Not a third-party service:** I confirm that this request is not exclusive to a third-party hosting site, relay, API service, or fork that has not been verified against the unmodified repository. Third-party instance issues must be reported to their operator."
|
||||
required: true
|
||||
- label: "**Not pass-through:** I confirm this is not a forwarding report after enabling pass-through; pass-through sends content as-is and does not go through new-api processing logic."
|
||||
required: true
|
||||
- label: "**Channel and protocol boundary:** I confirm this is not a feature request for a Coding Plan service, reverse-engineered channel, third-party API wrapper, Codex reverse-proxy endpoint, or treating Codex API-specific behavior as standard OpenAI API behavior. If first observed through such an interface, I have based the request on a standard API protocol supported by this repository."
|
||||
required: true
|
||||
- label: "**Template intact:** I have not removed any guidance or section headings from this template and will complete it as requested."
|
||||
required: true
|
||||
- label: "**Maintainer time:** I understand that maintainers have limited time and issues that do not follow this template may be ignored or closed directly."
|
||||
- label: "**Maintainer time:** I understand that maintainers have limited time. Do not paste unfiltered AI-generated text in the issue body or in later comments; extract the points needed for review first. Submissions that do not follow this template, or that omit the necessary points and cannot be reviewed, may be ignored or closed. Repeated submissions of this kind may result in a block."
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
|
||||
@@ -7,7 +7,8 @@ English template: `.github/PULL_REQUEST_TEMPLATE/en.md`
|
||||
|
||||
> [!IMPORTANT]
|
||||
>
|
||||
> - 描述可用 AI 辅助。提交前请审阅全文,并**声明对其负责**,避免未经核对的直接粘贴。
|
||||
> - 描述可用 AI 辅助。提交前请审阅并精炼全文,仅保留维护者需要复核的要点,并**声明对其负责**。PR 正文及后续评论均不得直接粘贴未经人工过滤的大段 AI 文本;多次提交此类内容可能会被 block。
|
||||
> - 开启透传请勿反馈问题;透传会直接发送内容,不会经过 new-api 的逻辑处理。
|
||||
> - 请按本模板填写后再提交。
|
||||
|
||||
## 🔗 关联任务 / Related Issue
|
||||
@@ -24,17 +25,18 @@ English template: `.github/PULL_REQUEST_TEMPLATE/en.md`
|
||||
- [ ] 📝 文档更新 (Documentation)
|
||||
|
||||
## 📝 变更描述 / Description
|
||||
(简述做了什么、为什么生效。如果难以简述,建议先拆分范围,或在 Issue 中与维护者对齐。)
|
||||
(简述做了什么、为什么生效。不得直接粘贴未经人工过滤的大段 AI 文本。如果难以简述,建议先拆分范围,或在 Issue 中与维护者对齐。)
|
||||
|
||||
## 📸 运行证明 / Proof of Work
|
||||
(请写明如何验证:实际步骤与观察结果。UI 变更请附截图或录屏;Bug 修复请说明复现过程与修复后结果。)
|
||||
(请写明如何验证:实际命令或步骤,以及观察结果。仅声明 `go build` 通过或测试通过,不视为有效证明。UI 变更请附截图或录屏;Bug 修复请说明复现过程与修复后结果。)
|
||||
|
||||
## ✅ 提交前检查项 / Checklist
|
||||
- [ ] **人工确认:** 无论描述是否由 AI 生成,我已审阅全部内容,并声明对其准确性与完整性负责。
|
||||
- [ ] **人工确认:** 无论描述是否由 AI 生成,我已审阅并精炼全文,仅保留可复核的要点,并声明对其准确性与完整性负责。PR 正文及后续评论均未直接粘贴未经人工过滤的大段 AI 文本。
|
||||
- [ ] **非重复提交:** 我已搜索现有的 [Issues](https://github.com/QuantumNous/new-api/issues) 与 [PRs](https://github.com/QuantumNous/new-api/pulls),确认不是重复提交。
|
||||
- [ ] **新功能关联 Issue:** 若此 PR 标记为 `New feature`,我已关联对应 Issue;若尚无 Issue,我已先自行创建。
|
||||
- [ ] **事前沟通:** 若改动较大或涉及方向性变更,已在关联 Issue 中与维护者沟通并达成一致。
|
||||
- [ ] **功能范围:** 本 PR 不是 Coding Plan、逆向渠道、第三方封装接口,也不是对 Codex 渠道类型的改动。
|
||||
- [ ] **非透传:** 本 PR 不是针对开启透传后的转发行为;透传会直接发送内容,不会经过 new-api 的逻辑处理。
|
||||
- [ ] **范围聚焦:** 本 PR 为一项聚焦改动,未包含无关代码。
|
||||
- [ ] **本地验证:** 已在本地运行并通过测试或手动验证,维护者可以据此复核。
|
||||
- [ ] **本地验证:** 已按变更路径实际验证,并写明命令与观察结果。仅声明 `go build` 通过或测试通过,不视为有效证明。
|
||||
- [ ] **安全合规:** 代码中无敏感凭据,且符合项目代码规范。
|
||||
@@ -7,7 +7,8 @@ Chinese template: `.github/PULL_REQUEST_TEMPLATE.md`
|
||||
|
||||
> [!IMPORTANT]
|
||||
>
|
||||
> - AI-assisted descriptions are welcome. Please review the full text before submitting and **take responsibility** for it. Avoid pasting unreviewed content.
|
||||
> - AI-assisted descriptions are welcome. Review and condense the full text before submitting, keep only the points maintainers need to re-check, and **take responsibility** for it. Do not paste unfiltered AI-generated text in the PR body or in later comments. Repeated submissions of this kind may result in a block.
|
||||
> - Do not report issues when pass-through is enabled; pass-through sends content as-is and does not go through new-api processing logic.
|
||||
> - Please complete this template before submitting.
|
||||
|
||||
## Related Issue
|
||||
@@ -24,17 +25,18 @@ Chinese template: `.github/PULL_REQUEST_TEMPLATE.md`
|
||||
- [ ] Documentation
|
||||
|
||||
## Description
|
||||
(Briefly describe what changed and why it works. If that is hard to summarize, consider splitting the scope or aligning with maintainers in an Issue first.)
|
||||
(Briefly describe what changed and why it works. Do not paste unfiltered AI-generated text. If that is hard to summarize, consider splitting the scope or aligning with maintainers in an Issue first.)
|
||||
|
||||
## Proof of Work
|
||||
(Please describe how this was verified: the steps run and what was observed. For UI changes, please include a screenshot or recording. For bug fixes, please describe the reproduction and the result after the fix.)
|
||||
(Record how this was verified: the actual commands or steps and the observed results. Stating only that `go build` or tests passed is not valid proof. For UI changes, include a screenshot or recording. For bug fixes, describe the reproduction and the result after the fix.)
|
||||
|
||||
## Checklist
|
||||
- [ ] **Human review:** Whether or not the description was AI-generated, I have reviewed the full content and take responsibility for its accuracy and completeness.
|
||||
- [ ] **Human review:** Whether or not the description was AI-generated, I have reviewed and condensed the full text, retained only the points needed for review, and take responsibility for its accuracy and completeness. I have not pasted unfiltered AI-generated text in the PR body or in later comments.
|
||||
- [ ] **Not a duplicate:** I have searched existing [Issues](https://github.com/QuantumNous/new-api/issues) and [PRs](https://github.com/QuantumNous/new-api/pulls) and confirmed this is not a duplicate.
|
||||
- [ ] **Feature issue:** If this PR is a New feature, I have linked a corresponding Issue; if none existed, I created one first.
|
||||
- [ ] **Prior discussion:** If this is a large or directional change, I have discussed it with maintainers in the linked Issue and reached agreement.
|
||||
- [ ] **Scope:** This PR is not a Coding Plan, reverse-engineered channel, third-party API wrapper, or a change to the Codex channel type.
|
||||
- [ ] **Not pass-through:** This PR is not about forwarding behavior after enabling pass-through; pass-through sends content as-is and does not go through new-api processing logic.
|
||||
- [ ] **Focused change:** This PR is a single focused change and does not include unrelated code.
|
||||
- [ ] **Local verification:** I have run tests or manually verified locally so maintainers can re-check from this evidence.
|
||||
- [ ] **Local verification:** I verified the changed path and recorded the commands and observed results. Stating only that `go build` or tests passed is not valid proof.
|
||||
- [ ] **Security:** This change does not include secrets and follows the project's coding guidelines.
|
||||
@@ -115,21 +115,22 @@ Inside `relaykit/`, use `kitutil.*` from `relaykit/relayconvert/kitutil/json.go`
|
||||
- Preserve explicit zero values in upstream relay request DTOs: absent client JSON fields must become `nil` and be omitted, while explicit `0`, `0.0`, or `false` values must remain non-`nil` and be sent upstream.
|
||||
- Avoid non-pointer scalars with `omitempty` for optional request parameters, because zero values will be silently dropped during marshal.
|
||||
|
||||
**Billing expression system:** When working on tiered/dynamic billing (expression-based pricing), MUST read `pkg/billingexpr/expr.md` first. It documents the design philosophy, expression language, full architecture, token normalization rules, quota conversion, and expression versioning. All billing expression changes must follow that document.
|
||||
**JavaScript task plugins (mandatory):**
|
||||
|
||||
**Built-in model pricing:** New built-in model prices MUST be defined as self-contained billing expressions in `setting/billing_setting/builtin_billing.go`, using real USD per million tokens. Do not add new built-in prices to the legacy model/completion/cache ratio tables. Preserve explicit administrator pricing overrides. Existing legacy prices are migrated only when explicitly requested. Verify published prices and cover applicable context-length thresholds and cache categories.
|
||||
- Before implementing, modifying, or reviewing JavaScript task plugins or their host API/runtime, MUST read [Task Plugin API v1](docs/plugin-api/v1.md), including its description writing and translation conventions. When changing the plugin contract, also check `docs/plugin-api/v1.schema.json` and `docs/plugin-api/v1.d.ts` for consistency.
|
||||
- For numeric billing fields in `usageSchema` and `usageProfiles[].schema`, `description` MUST name the **billing subject + unit price**, because it labels the price input in the UI. For example, `image_count` uses `Image generation unit price` / `图片生成单价`, not `Generated image count` / `生成图片张数`; `seconds` uses `Video generation unit price` / `视频生成单价`. The field value remains a usage quantity, not a price.
|
||||
- Put units in `unit`. Keep protocol limits, usage sources, defaults, estimation, and settlement details in code comments or technical documentation. Descriptions must be short, equivalent across languages, and free of numeric prices and trailing punctuation. Follow the API document's separate wording rules for actions, booleans, and other enum conditions.
|
||||
- Review metadata wording explicitly before completing plugin work. These are authoring requirements; successful compilation, schema validation, or tests do not verify that descriptions follow them.
|
||||
|
||||
**Billing safety invariants:** Quota/billing code MUST never produce a negative charge (a credit) from arithmetic overflow or unvalidated input. Apply defense in depth:
|
||||
**Billing rules (mandatory read gate):** `.agents/rules/billing.md` holds the billing conventions (expression system, built-in pricing, safety invariants, deriving billable quantities from upstream responses). If a task touches billing as defined below, you MUST read that file in full with your file-read tool — not a grep, a partial skim, memory, or a summary — before planning, coding, or reviewing, and then follow every rule in it. A task touches billing when it meets any of these:
|
||||
|
||||
- Every user-controlled quantity that becomes a billing multiplier (image `n`, video `seconds`/`duration`, resolution/quality ratios, batch counts) MUST be bounded before it reaches quota calculation. Reject out-of-range values at request validation with a 400. Existing bounds: `dto.MaxImageN` for image generation count, `relaycommon.MaxTaskDurationSeconds` for task video duration, `maxTokensLimit` (`relay/helper/valid_request.go`) for `max_tokens`-family fields on every relay format (OpenAI, Claude, Gemini, Responses). Reuse these constants instead of introducing new ad hoc limits for the same concepts. When adding a new relay format or request DTO, bound its max-tokens and count fields in its validator from day one.
|
||||
- Watch for validation bypass paths: passthrough fields (e.g. `Extra["parameters"]`), task `metadata` maps, and multipart form fields can carry the same quantities around the standard DTO validation. Any adaptor that reads a multiplier from such a path must enforce the same bound (or clamp) locally.
|
||||
- Durations parsed from media metadata are user/upstream-controlled too: audio file headers (transcription token counting, TTS response duration) and upstream deduction numbers (e.g. Kling `FinalUnitDeduction`) can claim absurd values. Convert them with saturation before they become token counts.
|
||||
- Never convert a computed quota or token count to `int` with a bare cast like `int(float64(quota) * ratio)`, `int(math.Round(...))` on unbounded input, or `int(decimal.IntPart())`. All quota rounding/conversion is centralized in `common/quota_math.go`; use those helpers: `common.QuotaFromFloat` (truncating) for float products, `common.QuotaRound` (half-away-from-zero) where rounding is intended, and `common.QuotaFromDecimal` for decimal products. `billingexpr.QuotaRound` delegates to `common.QuotaRound`. Do not reintroduce local conversion helpers or bare casts. Single-request saturation stays at the int32 boundary so batch accumulation cannot approach 64-bit wraparound; wallet/top-up conversion uses `common.WalletQuotaFromDecimalStrict` with the JavaScript-safe `common.MaxWalletQuota` boundary. Every clamp/NaN fallback is logged via `common.SysError`.
|
||||
- Saturation events are also audited: each helper has a `*Checked` variant (`common.QuotaFromFloatChecked` / `QuotaRoundChecked` / `QuotaFromDecimalChecked`) that additionally returns a `*common.QuotaClamp` when clamping occurred. Billing paths that compute a charge capture that clamp onto `relayInfo.QuotaClamp` (or thread it into task settlement) and, right before writing the consume/task log, call `attachQuotaSaturation` (in `service/log_info_generate.go`) which nests the marker under the log's `other.admin_info.quota_saturation` and emits a request-correlated `logger.LogWarn`. Nesting under `admin_info` makes it admin-only for free (non-admin log views strip `admin_info`). When adding a new billing path, use the `*Checked` variant and surface the clamp the same way so the anomaly stays auditable in both the admin log UI and backend logs.
|
||||
- Multiplier maps go through `types.PriceData.AddOtherRatio`, which rejects non-positive, NaN, and +Inf ratios. Do not write to `PriceData.OtherRatios` directly, and do not weaken these guards.
|
||||
- Pre-consume (预扣费) and settle (结算/差额) must both be safe: a saturated oversized quota must fail pre-consume with insufficient-quota, never silently wrap. When adding a new billing path (new relay format, new task platform, new adjustment hook), trace the full chain — validation → EstimateBilling/OtherRatios → quota conversion → pre-consume → settle/refund — and confirm each step preserves these invariants.
|
||||
- Fields parsed into unsigned types (`*uint`) accept huge positive JSON numbers (e.g. `18446744073686646784`, a wrapped negative); a `>= 0` check is not sufficient, an upper bound is mandatory.
|
||||
- Regression tests for these invariants belong with the boundary they protect (request validators, converter helpers). See `relay/helper/openai_image_request_test.go`, `relay/common/relay_utils_test.go`, and `common/quota_math_test.go` for the expected style.
|
||||
- Edits `pkg/billingexpr/`, `setting/billing_setting/`, `common/quota_math.go`, `types/price_data.go`, `relay/request_billing.go`, `relay/image_handler.go`, `relay/relay_task.go`, `relay/helper/price.go`, `relay/helper/billing_expr_request.go`, `relay/helper/valid_request.go`, `service/quota.go`, `service/text_quota.go`, `service/image_billing.go`, `service/tiered_settle.go`, `service/task_billing.go`, `service/responses_usage.go`, the billing part of `service/log_info_generate.go`, the `model/pricing*.go` / `model/model_pricing*.go` files, or the billing fields and methods of `relay/common/relay_info.go` (`PriceData`, `TieredBillingSnapshot`, `BillingImageCount`, `UpdateImageCount`).
|
||||
- Reads or writes `PriceData` / `OtherRatios`, quota pre-consume, settlement, refund, or consume-log billing fields anywhere else.
|
||||
- Derives a billable quantity or `Usage` from an upstream response or stream (image counts, seconds, tokens, task deductions) in any channel adaptor, response handler, or task plugin.
|
||||
- Validates, bounds, or forwards a request field that becomes a billing multiplier (`n`, `max_tokens`-family fields, duration, resolution or quality, batch counts), including passthrough and multipart paths.
|
||||
- Adds or changes model prices, or numeric `usageSchema` / `usageProfiles[].schema` fields in a task plugin.
|
||||
|
||||
Tasks that touch none of these (for example unrelated frontend work, authentication, database migrations, or protocol conversion that leaves usage untouched) do not need to read it.
|
||||
|
||||
**Backend test quality:** Backend tests must protect real behavior, API contracts, billing/accounting invariants, data compatibility, or regression paths.
|
||||
|
||||
@@ -145,7 +146,10 @@ Inside `relaykit/`, use `kitutil.*` from `relaykit/relayconvert/kitutil/json.go`
|
||||
- Avoid hand-written assertion helpers unless they encode a reusable project-specific invariant.
|
||||
- When cleaning tests, preserve meaningful regression coverage. If a deleted test covered a real contract indirectly, replace it with a smaller test that asserts that contract directly.
|
||||
|
||||
**Documentation files:** Do NOT add new files under `docs/` or any of its subdirectories unless the user explicitly requests it.
|
||||
**Documentation files:**
|
||||
|
||||
- Do NOT add new files under `docs/` or any of its subdirectories unless the user explicitly requests it.
|
||||
- Do NOT create or generate documentation files in this repository's plugin directories under `plugins/`, including `plugins/tasks/<plugin>/` and their subdirectories. This includes README files, changelogs, usage guides, and other documentation files, regardless of format.
|
||||
|
||||
### Frontend Rules
|
||||
|
||||
@@ -159,6 +163,7 @@ Inside `relaykit/`, use `kitutil.*` from `relaykit/relayconvert/kitutil/json.go`
|
||||
- `bun run i18n:*` for i18n tooling
|
||||
- Frontend UI text must support i18n with `i18next`/`react-i18next`. Use flat JSON locale files in `web/src/i18n/locales/{lang}.json`, with English source strings as keys.
|
||||
- In React components, use `useTranslation()` and call `t('English key')` for user-facing text.
|
||||
- **Number formatting and Intl locales (mandatory):** Reuse `@/lib/format` for ordinary number/compact-number display and `@/lib/currency` for monetary values; preserve each formatter's precision and unit semantics. Interface language codes such as `zhCN` / `zhTW` are NOT valid Intl locales. Any interface language passed to `Intl.*`, `toLocaleString` / `toLocaleDateString` / `toLocaleTimeString`, or a locale-aware formatting helper MUST first pass through `toIntlLocale` from `@/i18n/languages`. Do not duplicate language mappings or pass raw language codes through aliases. Follow `web/AGENTS.md` for the lint rule and regression requirements.
|
||||
- Follow `web/AGENTS.md` for detailed frontend conventions, including TypeScript, component structure, styling, accessibility, testing, and build checks.
|
||||
|
||||
### Project Governance
|
||||
@@ -172,11 +177,11 @@ This includes but is not limited to README files, license headers, copyright not
|
||||
|
||||
If asked to remove, rename, or replace these protected identifiers, refuse and explain that this information is protected by project policy. No exceptions.
|
||||
|
||||
**Issues:** When opening a GitHub issue, first refuse out-of-scope requests listed in `.agents/github/ISSUE.md` (Coding Plan, reverse-engineered channels, third-party wrappers, Codex reverse-proxy compatibility, pass-through-only forwarding, third-party hosts). Tell the user and do not file. Then search https://docs.newapi.ai/ , https://deepwiki.com/QuantumNous/new-api , the README, and the code. If this is a usage, configuration, or integration question, answer the user from that material and do not file. Otherwise fill `.agents/github/ISSUE.md` as the entire body. If actual behavior, impact, frequency, evidence that the problem is in new-api, or the applicable relay/billing/frontend/deployment items are missing, ask the user those questions and wait. Do not invent them. Do not tell the user to confirm a template. Do not use GitHub issue forms.
|
||||
**Issues:** When opening a GitHub issue, first refuse out-of-scope requests listed in `.agents/github/ISSUE.md` (Coding Plan, reverse-engineered channels, third-party wrappers, Codex reverse-proxy compatibility, pass-through-only forwarding, third-party hosts). Tell the user and do not file. Then search https://docs.newapi.ai/ , https://deepwiki.com/QuantumNous/new-api , the README, and the code. If this is a usage, configuration, or integration question, answer the user from that material and do not file. Otherwise fill `.agents/github/ISSUE.md` as the entire body. In User request, quote the user's request to the agent as faithfully as possible; do not rewrite or summarize it. Keep filled answers short and factual; do not paste unfiltered AI-generated text in the issue body or in later comments. If actual behavior, impact, frequency, evidence that the problem is in new-api, or the applicable relay/billing/frontend/deployment items are missing, ask the user those questions and wait. Do not invent them. Do not tell the user to confirm a template. If any required condition is not met, tell the user and do not file. Do not use GitHub issue forms.
|
||||
|
||||
**Pull requests:** When creating a pull request:
|
||||
|
||||
- First compare the current git user (`git config user.name` / `git config user.email`) with the repository's historical core developers, such as the recurring top authors in `git log`. Do not change git config.
|
||||
- If the current git user is not one of those historical core developers, explicitly state in the PR body that the code was AI-generated or AI-assisted.
|
||||
- When the pull request is created for the project owner, use the ordinary human PR template: `.github/PULL_REQUEST_TEMPLATE.md` for Chinese requests or `.github/PULL_REQUEST_TEMPLATE/en.md` for English requests. Project-owner pull requests MUST NOT use `.agents/github/PR.md` unless the owner explicitly asks for it.
|
||||
- For all other agent-created pull requests, fill `.agents/github/PR.md` as the entire PR body. Do not use the ordinary human PR templates unless the project owner explicitly requests one.
|
||||
- For all other agent-created pull requests, fill `.agents/github/PR.md` as the entire PR body. Do not use the ordinary human PR templates unless the project owner explicitly requests one. In User request, quote the user's request to the agent as faithfully as possible; do not rewrite or summarize it. Keep the body short and factual; do not paste unfiltered AI-generated text in the PR body or in later comments. Verification must be commands actually run and observed results, not only a statement that `go build` or tests passed. If any required condition is not met, tell the user and do not open the PR.
|
||||
@@ -13,3 +13,4 @@ Rules:
|
||||
- Do not replace the Read with a grep, glob, or partial skim.
|
||||
- After reading, follow every rule in `AGENTS.md` for the rest of the work.
|
||||
- If the task touches `web/`, also Read `web/AGENTS.md` before editing frontend files.
|
||||
- If the task touches billing as defined under **Billing rules (mandatory read gate)** in `AGENTS.md`, also Read `.agents/rules/billing.md` in full before planning or editing. Tasks outside that definition may skip it.
|
||||
+25
-13
@@ -1,20 +1,32 @@
|
||||
# 自定义版本说明
|
||||
|
||||
当前生产版本基于官方 `v1.0.0-rc.36`(`ea7cb0ba4e0f82e2bfa5e55752eb68bdf902f71b`),版本号为 `v1.0.0-rc.36-custom.2`。
|
||||
当前版本为 v1.0.0-rc.38-custom.42-20260926,基于官方 v1.0.0-rc.38,对应 2026-09-26 已部署的自定义源码。
|
||||
|
||||
包含以下生产定制:
|
||||
## 主要定制
|
||||
|
||||
- 仅展示存在可用渠道的分组;禁用、无渠道和不可用分组不进入用户筛选。
|
||||
- 固定令牌分组可配置一个备用分组;主分组无路由或发生可重试上游错误时按顺序切换,实际命中的分组用于计费和日志。
|
||||
- 固定令牌配置了备用分组时,备用分组切换优先于会话亲和规则的“失败后不重试”;上游返回可重试错误(如 503)时不再停在主分组渠道。不可重试错误、显式指定渠道和 auto 分组行为保持不变。
|
||||
- 分组状态页展示请求成功率、首字延迟、平均延迟、TPS、请求缓存和令牌缓存指标。
|
||||
- 支持按分组配置动态计费表达式,并在价格页按所选分组展示。
|
||||
- 用户侧隐藏上游错误细节,完整错误仅保留在 root 日志范围;客户端断开不计为上游失败。
|
||||
- 固定令牌支持有序分组池、备用分组、可重试错误切换及分组可用性判断;按实际命中的分组计费和记录日志。
|
||||
- 支持分组动态计费表达式、请求明细、路由过程及逐次请求结算。
|
||||
- 用户侧隐藏上游错误细节;客户端断开与真实上游故障分别处理。
|
||||
- Astra 模型指纹抽检、响应模型守卫、鹈鹕检测及检测费用结算。
|
||||
- 独立公开 uptime 页,使用 shadcn/ui,支持分组别名、隐藏分组和管理员重置当前分组检测记录。
|
||||
- 主站 /group-status 使用真实分组名并展示全部分组,检测请求使用真实分组标识;独立 uptime 页应用公开展示配置。
|
||||
- 渠道每 15 分钟探活,展示最近 18 次记录;正常用户请求可恢复分组正常状态。
|
||||
- 支持 Astra 的分组排在前面,其他分组使用紧凑布局;鹈鹕样例按时间倒序横向滚动,最新在左,点击缩略图预览。
|
||||
- 鹈鹕检测通过 manxue.ai 访客接口请求公开作品,保存可跳转的作品链接;连接失败时本地生成并标记为不确定。
|
||||
- 作品脚本在隔离 iframe 中运行,以保留腿部、踏板等动画绘制;缩略图按容器宽度等比缩放。
|
||||
|
||||
生产镜像通过以下命令构建:
|
||||
## 从源码构建
|
||||
|
||||
```sh
|
||||
docker build --pull=false -f Dockerfile.request-status -t local/new-api:v1.0.0-rc.36-custom.2-20260911 .
|
||||
```
|
||||
docker build -f Dockerfile.local-release -t local/new-api:v1.0.0-rc.38-custom.42-20260926 .
|
||||
|
||||
`Dockerfile.request-status` 会执行前端类型检查、相关 Vitest 与 oxlint 检查、Go 回归测试和生产构建。
|
||||
此入口从 Bun、Go 和 Debian 基础镜像构建前端与服务,不要求预先存在旧版 local/new-api 镜像。VERSION 用于编译版本号;Compose 如果设置了 VERSION 环境变量,也应同步更新。
|
||||
|
||||
其他 Dockerfile.* 和 build-fingerprint-backend.sh 保留为历史构建参考。部分增量构建文件依赖服务器上的旧版镜像或已编译二进制,新环境请使用 Dockerfile.local-release。
|
||||
|
||||
## 运行与升级
|
||||
|
||||
使用现有外部 PostgreSQL 和 Valkey,连接信息通过部署环境提供。仓库不包含生产凭据、数据库内容或线上 Compose 配置。
|
||||
|
||||
升级前备份数据库及 Compose 配置,部署后检查容器健康、版本号、主站分组状态和公开 uptime 页面。Uptime Kuma 的停用与域名反向代理属于部署配置,不会仅因替换应用镜像自动完成。
|
||||
|
||||
公开作品链接受 manxue.ai 的保留期限约束;没有公开链接的历史作品继续使用本站预览。
|
||||
@@ -0,0 +1,3 @@
|
||||
FROM local/new-api:v1.0.0-rc.38-custom.41-20260926
|
||||
COPY new-api /new-api
|
||||
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.42-20260926"
|
||||
@@ -0,0 +1,3 @@
|
||||
FROM local/new-api:v1.0.0-rc.38-custom.24-20260923
|
||||
COPY new-api /new-api
|
||||
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.25-20260923"
|
||||
@@ -0,0 +1,4 @@
|
||||
FROM local/new-api:v1.0.0-rc.38-custom.26-20260923
|
||||
COPY new-api /new-api
|
||||
COPY LICENSE NOTICE THIRD-PARTY-LICENSES.md /licenses/
|
||||
COPY pkg/model_fingerprint/LICENSE-*.txt /licenses/model-fingerprint/
|
||||
@@ -0,0 +1,31 @@
|
||||
FROM oven/bun:1.4.0 AS web
|
||||
WORKDIR /build/web
|
||||
COPY web/package.json web/bun.lock ./
|
||||
RUN bun install --frozen-lockfile
|
||||
COPY web/ ./
|
||||
RUN bun run build
|
||||
|
||||
FROM golang:1.25.1 AS build
|
||||
ENV GOMAXPROCS=2 GOFLAGS=-p=2 GO111MODULE=on CGO_ENABLED=0 GOWORK=off
|
||||
ARG TARGETOS
|
||||
ARG TARGETARCH
|
||||
ENV GOOS=${TARGETOS:-linux} GOARCH=${TARGETARCH:-amd64}
|
||||
WORKDIR /build
|
||||
COPY . .
|
||||
COPY --from=web /build/web/dist ./web/dist
|
||||
RUN --mount=type=cache,target=/go/pkg/mod \
|
||||
--mount=type=cache,target=/root/.cache/go-build \
|
||||
go build -buildvcs=false -ldflags "-s -w -X 'github.com/QuantumNous/new-api/common.Version=$(cat VERSION)'" -o new-api
|
||||
|
||||
FROM debian:bookworm-slim
|
||||
RUN apt-get update \
|
||||
&& apt-get install -y --no-install-recommends ca-certificates tzdata libasan8 wget \
|
||||
&& rm -rf /var/lib/apt/lists/* \
|
||||
&& update-ca-certificates
|
||||
COPY --from=build /build/new-api /
|
||||
COPY LICENSE NOTICE THIRD-PARTY-LICENSES.md /licenses/
|
||||
COPY pkg/model_fingerprint/LICENSE-*.txt /licenses/model-fingerprint/
|
||||
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.42-20260926"
|
||||
EXPOSE 3000
|
||||
WORKDIR /data
|
||||
ENTRYPOINT ["/new-api"]
|
||||
@@ -0,0 +1,3 @@
|
||||
FROM local/new-api:v1.0.0-rc.38-custom.38-20260926
|
||||
COPY new-api /new-api
|
||||
LABEL org.opencontainers.image.version="v1.0.0-rc.38-custom.39-20260926"
|
||||
@@ -8,9 +8,10 @@ WORKDIR /build/web
|
||||
COPY --from=webdeps /build/web/node_modules ./node_modules
|
||||
COPY web ./
|
||||
COPY VERSION /build/VERSION
|
||||
COPY pkg/billingexpr/testdata /build/pkg/billingexpr/testdata
|
||||
RUN npm run typecheck
|
||||
RUN node node_modules/vitest/vitest.mjs run src/features/group-status/__tests__ src/features/pricing/__tests__/group-pricing.test.ts src/features/keys/lib/__tests__/auto-group-form.test.ts src/features/keys/components/__tests__/api-keys-mutate-drawer.test.tsx
|
||||
RUN npx --no-install oxlint -c .oxlintrc.json src/features/group-status src/features/dashboard/types.ts src/features/keys/components/api-keys-mutate-drawer.tsx src/features/keys/lib/api-key-form.ts src/features/keys/types.ts
|
||||
RUN node node_modules/vitest/vitest.mjs run src/features/usage-logs/components/__tests__ src/features/system-settings/request-policies src/features/group-status/__tests__ src/features/pricing/__tests__/group-pricing.test.ts src/features/keys/lib/__tests__/auto-group-form.test.ts src/features/keys/components/__tests__/api-keys-mutate-drawer.test.tsx src/features/keys/components/__tests__/astra-cost-notice.test.tsx src/features/channels/components/__tests__/responses-websocket-setting.test.tsx src/features/profile/__tests__/settings.test.tsx
|
||||
RUN npx --no-install oxlint -c .oxlintrc.json src/features/group-status src/features/dashboard/types.ts src/features/keys/components/astra-cost-notice.tsx src/features/keys/components/__tests__/astra-cost-notice.test.tsx src/features/keys/components/api-keys-mutate-drawer.tsx src/features/keys/lib/api-key-form.ts src/features/keys/types.ts src/features/channels/components/astra-model-guard-setting.tsx src/features/channels/components/drawers/channel-mutate-drawer.tsx src/features/channels/components/__tests__/responses-websocket-setting.test.tsx src/features/channels/lib/channel-form.ts src/features/channels/lib/channel-configuration.ts src/features/profile/lib/user-settings.ts src/features/profile/types.ts src/features/usage-logs/components/routing-attempts.tsx src/features/usage-logs/components/model-badge.tsx src/features/usage-logs/components/common-logs-filter-bar.tsx
|
||||
RUN DISABLE_ESLINT_PLUGIN='true' VITE_REACT_APP_VERSION=$(cat /build/VERSION) npm run build
|
||||
|
||||
FROM golang:1.26.1-alpine@sha256:2389ebfa5b7f43eeafbd6be0c3700cc46690ef842ad962f6c5bd6be49ed82039 AS build
|
||||
@@ -25,7 +26,12 @@ RUN go mod download
|
||||
COPY . .
|
||||
COPY --from=webbuild /build/web/dist ./web/dist
|
||||
RUN go test ./pkg/perf_metrics ./model ./controller ./service ./relay/helper -run 'TestGroupBuckets|TestRequestStatus|TestGroupModelPricing|TestGetEnabledGroupSetRequiresEnabledAbilityAndChannel|TestSelectableGroupsRequireEnabledChannel|TestAddTokenPersistsFixedBackupGroup|TestUpdateTokenAutoGroupsTriStateAndFixedBackup|TestAddTokenRejectsInvalidFixedBackupGroup|TestCacheGetRandomSatisfiedChannelUsesTokenAutoGroupsWhenGlobalAutoIsEmpty|TestCacheGetRandomSatisfiedChannelUsesFixedTokenBackupGroup|TestCacheGetRandomSatisfiedChannelSkipsUnavailableFixedPrimary|TestAppendStreamStatusTreatsClientGoneAsDisplayNormal|TestClientFacingOpenAIErrorHidesUpstreamDetails|TestClientFacingClaudeErrorHidesUpstreamDetails|TestClientFacingErrorKeepsLocalValidationMessage|TestPublicUpstreamErrorContent' -count=1
|
||||
RUN go test ./controller ./service -run 'TestFixedBackup|TestShouldSkipRetryAfterChannelAffinityFailure|TestShouldRetryHonorsPinRetryMode' -count=1
|
||||
RUN go test ./model ./controller ./service -run 'TestChannelGetSettingDefaultsAstraGuardOn|TestFixedBackup|TestShouldSkipRetryAfterChannelAffinityFailure|TestShouldRetryHonorsPinRetryMode|TestGroupPool|TestBuildGroupPoolTiers|TestNextTierIndex|TestGroupAvailability|TestGetGroupPool|TestSetGroupPool' -count=1
|
||||
RUN go test ./model ./service ./controller ./relay/helper -run 'TestRequestDetails|TestSelectionTrace|TestResponsesInterruptedStreamHealth|TestResponsesWebSocketReusesConnectionAndSettlesEachRequest|TestStreamScanner' -count=1
|
||||
RUN cd relaykit && GOWORK=off go build ./...
|
||||
RUN go test ./relay -run 'TestResponseGuard|TestResponseModel' -count=1
|
||||
RUN go test ./controller ./middleware ./service -run 'TestTokenAstraMigration|TestRelayAttemptGuardResult|TestRelayAttemptSwitchSafety|TestTokenAuth|TestAstra' -count=1
|
||||
RUN go test ./pkg/model_fingerprint ./controller ./service ./router -run 'TestFingerprint|TestAstraSampling|TestSystemTask|TestChannelPermission' -count=1
|
||||
RUN go build -ldflags "-s -w -X 'github.com/QuantumNous/new-api/common.Version=$(cat VERSION)'" -o new-api
|
||||
|
||||
FROM debian:bookworm-slim@sha256:f06537653ac770703bc45b4b113475bd402f451e85223f0f2837acbf89ab020a
|
||||
@@ -35,6 +41,7 @@ RUN apt-get update \
|
||||
&& update-ca-certificates
|
||||
COPY --from=build /build/new-api /
|
||||
COPY LICENSE NOTICE THIRD-PARTY-LICENSES.md /licenses/
|
||||
COPY pkg/model_fingerprint/LICENSE-*.txt /licenses/model-fingerprint/
|
||||
EXPOSE 3000
|
||||
WORKDIR /data
|
||||
ENTRYPOINT ["/new-api"]
|
||||
ENTRYPOINT ["/new-api"]
|
||||
@@ -322,3 +322,7 @@ this software, either in source code form or as a compiled binary, for any
|
||||
purpose, commercial or non-commercial, and by any means.
|
||||
|
||||
For more information, please refer to https://unlicense.org/
|
||||
|
||||
## Numerical model fingerprint assets
|
||||
|
||||
The embedded bank and scorer in `pkg/model_fingerprint` are from [gpt-model-check](https://github.com/Katao123/gpt-model-check), commit `3d6716a0903cf061daea001d5d44ee410846634e`, using the ModelTrace bank dated 2026-09-05. The upstream gpt-model-check, ModelTrace, and is-gpt-nerfed MIT notices are included alongside the sources and under `/licenses/model-fingerprint/` in the image. Relative candidate weights are not calibrated identity probabilities.
|
||||
@@ -0,0 +1,9 @@
|
||||
set -eu
|
||||
go test ./pkg/perf_metrics ./model ./controller ./service ./relay/helper -run 'TestGroupBuckets|TestRequestStatus|TestGroupModelPricing|TestGetEnabledGroupSetRequiresEnabledAbilityAndChannel|TestSelectableGroupsRequireEnabledChannel|TestAddTokenPersistsFixedBackupGroup|TestUpdateTokenAutoGroupsTriStateAndFixedBackup|TestAddTokenRejectsInvalidFixedBackupGroup|TestCacheGetRandomSatisfiedChannelUsesTokenAutoGroupsWhenGlobalAutoIsEmpty|TestCacheGetRandomSatisfiedChannelUsesFixedTokenBackupGroup|TestCacheGetRandomSatisfiedChannelSkipsUnavailableFixedPrimary|TestAppendStreamStatusTreatsClientGoneAsDisplayNormal|TestClientFacingOpenAIErrorHidesUpstreamDetails|TestClientFacingClaudeErrorHidesUpstreamDetails|TestClientFacingErrorKeepsLocalValidationMessage|TestPublicUpstreamErrorContent' -count=1
|
||||
go test ./model ./controller ./service -run 'TestChannelGetSettingDefaultsAstraGuardOn|TestFixedBackup|TestShouldSkipRetryAfterChannelAffinityFailure|TestShouldRetryHonorsPinRetryMode|TestGroupPool|TestBuildGroupPoolTiers|TestNextTierIndex|TestGroupAvailability|TestGetGroupPool|TestSetGroupPool' -count=1
|
||||
go test ./model ./service ./controller ./relay/helper -run 'TestRequestDetails|TestSelectionTrace|TestResponsesInterruptedStreamHealth|TestResponsesWebSocketReusesConnectionAndSettlesEachRequest|TestStreamScanner' -count=1
|
||||
(cd relaykit && GOWORK=off go build ./...)
|
||||
go test ./relay -run 'TestResponseGuard|TestResponseModel' -count=1
|
||||
go test ./controller ./middleware ./service -run 'TestTokenAstraMigration|TestRelayAttemptGuardResult|TestRelayAttemptSwitchSafety|TestTokenAuth|TestAstra' -count=1
|
||||
go test ./pkg/model_fingerprint ./controller ./service ./router -run 'TestFingerprint|TestAstraSampling|TestSystemTask|TestChannelPermission' -count=1
|
||||
go build -ldflags "-s -w -X 'github.com/QuantumNous/new-api/common.Version=$(cat VERSION)'" -o new-api
|
||||
@@ -0,0 +1,28 @@
|
||||
package common
|
||||
|
||||
import (
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
)
|
||||
|
||||
// GetAdvancedCustomPreset returns fresh defaults.
|
||||
func GetAdvancedCustomPreset(channelType int) *dto.AdvancedCustomConfig {
|
||||
if channelType != constant.ChannelTypeVLLM && channelType != constant.ChannelTypeSGLang {
|
||||
return nil
|
||||
}
|
||||
config := &dto.AdvancedCustomConfig{}
|
||||
for _, path := range []string{"/v1/chat/completions", "/v1/completions", "/v1/responses", "/v1/embeddings", "/v1/messages", dto.AdvancedCustomModelListPath} {
|
||||
config.Routes = append(config.Routes, dto.AdvancedCustomRoute{
|
||||
IncomingPath: path, UpstreamPath: path, Converter: "none",
|
||||
Auth: &dto.AdvancedCustomRouteAuth{Type: dto.AdvancedCustomAuthTypeHeader, Name: "Authorization", Value: "Bearer {api_key}"},
|
||||
})
|
||||
}
|
||||
if channelType == constant.ChannelTypeSGLang {
|
||||
for _, path := range []string{"/v1/rerank", "/rerank"} {
|
||||
config.Routes = append(config.Routes, dto.AdvancedCustomRoute{
|
||||
IncomingPath: path, UpstreamPath: "/v1/rerank", Converter: dto.AdvancedCustomConverterSGLangRerank,
|
||||
})
|
||||
}
|
||||
}
|
||||
return config
|
||||
}
|
||||
+1
-1
@@ -75,7 +75,7 @@ func ChannelType2APIType(channelType int) (int, bool) {
|
||||
apiType = constant.APITypeReplicate
|
||||
case constant.ChannelTypeCodex:
|
||||
apiType = constant.APITypeCodex
|
||||
case constant.ChannelTypeAdvancedCustom:
|
||||
case constant.ChannelTypeAdvancedCustom, constant.ChannelTypeVLLM, constant.ChannelTypeSGLang:
|
||||
apiType = constant.APITypeAdvancedCustom
|
||||
case constant.ChannelTypeSub2API:
|
||||
apiType = constant.APITypeSub2API
|
||||
|
||||
@@ -5,6 +5,7 @@ import (
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
|
||||
"github.com/abema/go-mp4"
|
||||
"github.com/go-audio/aiff"
|
||||
@@ -20,6 +21,7 @@ import (
|
||||
// 它不再依赖外部的 ffmpeg 或 ffprobe 程序。
|
||||
func GetAudioDuration(ctx context.Context, f io.ReadSeeker, ext string) (duration float64, err error) {
|
||||
SysLog(fmt.Sprintf("GetAudioDuration: ext=%s", ext))
|
||||
ext = strings.ToLower(ext)
|
||||
// 根据文件扩展名选择解析器
|
||||
switch ext {
|
||||
case ".mp3":
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
package common
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// minimalWAV builds a valid 16-bit PCM mono WAV with the given number of sample frames.
|
||||
func minimalWAV(sampleRate uint32, frames int) []byte {
|
||||
const numChannels = 1
|
||||
const bitsPerSample = 16
|
||||
dataSize := frames * numChannels * (bitsPerSample / 8)
|
||||
buf := &bytes.Buffer{}
|
||||
buf.WriteString("RIFF")
|
||||
binary.Write(buf, binary.LittleEndian, uint32(36+dataSize))
|
||||
buf.WriteString("WAVE")
|
||||
buf.WriteString("fmt ")
|
||||
binary.Write(buf, binary.LittleEndian, uint32(16))
|
||||
binary.Write(buf, binary.LittleEndian, uint16(1)) // PCM
|
||||
binary.Write(buf, binary.LittleEndian, uint16(numChannels))
|
||||
binary.Write(buf, binary.LittleEndian, sampleRate)
|
||||
byteRate := sampleRate * numChannels * (bitsPerSample / 8)
|
||||
binary.Write(buf, binary.LittleEndian, byteRate)
|
||||
binary.Write(buf, binary.LittleEndian, uint16(numChannels*(bitsPerSample/8)))
|
||||
binary.Write(buf, binary.LittleEndian, uint16(bitsPerSample))
|
||||
buf.WriteString("data")
|
||||
binary.Write(buf, binary.LittleEndian, uint32(dataSize))
|
||||
buf.Write(make([]byte, dataSize))
|
||||
return buf.Bytes()
|
||||
}
|
||||
|
||||
func TestGetAudioDurationCaseInsensitiveExt(t *testing.T) {
|
||||
wav := minimalWAV(8000, 8000) // 1 second
|
||||
|
||||
for _, ext := range []string{".wav", ".WAV", ".Wav"} {
|
||||
d, err := GetAudioDuration(context.Background(), bytes.NewReader(wav), ext)
|
||||
if err != nil {
|
||||
t.Fatalf("ext %q: unexpected error: %v", ext, err)
|
||||
}
|
||||
if d < 0.99 || d > 1.01 {
|
||||
t.Fatalf("ext %q: expected ~1s duration, got %v", ext, d)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetAudioDurationUppercaseNotUnsupported(t *testing.T) {
|
||||
_, err := GetAudioDuration(context.Background(), bytes.NewReader([]byte("not audio")), ".MP3")
|
||||
if err != nil && strings.Contains(err.Error(), "unsupported audio format") {
|
||||
t.Fatalf("uppercase extension rejected as unsupported: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -28,6 +28,8 @@ func GetEndpointTypesByChannelType(channelType int, modelName string) []constant
|
||||
endpointTypes = []constant.EndpointType{constant.EndpointTypeOpenAI}
|
||||
case constant.ChannelTypeXai:
|
||||
endpointTypes = []constant.EndpointType{constant.EndpointTypeOpenAI, constant.EndpointTypeOpenAIResponse}
|
||||
case constant.ChannelTypeVLLM, constant.ChannelTypeSGLang:
|
||||
endpointTypes = GetAdvancedCustomPreset(channelType).SupportedEndpointTypesForModel(modelName)
|
||||
case constant.ChannelTypeSora:
|
||||
endpointTypes = []constant.EndpointType{constant.EndpointTypeOpenAIVideo}
|
||||
case constant.ChannelTypeSub2API, constant.ChannelTypeNewAPI:
|
||||
|
||||
+4
-2
@@ -9,6 +9,8 @@ import (
|
||||
"github.com/gin-gonic/gin/binding"
|
||||
)
|
||||
|
||||
type RawMessage = json.RawMessage
|
||||
|
||||
// hostJSONCodec is the single place where the host chooses its JSON engine.
|
||||
// Swap the implementation here (for example to sonic.ConfigStd) and every
|
||||
// common.* and kitutil.* JSON helper, including relaykit DTO (un)marshalling,
|
||||
@@ -73,11 +75,11 @@ func IndentJson(data []byte) ([]byte, error) {
|
||||
return buffer.Bytes(), nil
|
||||
}
|
||||
|
||||
func GetJsonType(data json.RawMessage) string {
|
||||
func GetJsonType(data RawMessage) string {
|
||||
return kitutil.GetJsonType(data)
|
||||
}
|
||||
|
||||
// JsonRawMessageToString returns JSON strings as their decoded value and other JSON values as raw text.
|
||||
func JsonRawMessageToString(data json.RawMessage) string {
|
||||
func JsonRawMessageToString(data RawMessage) string {
|
||||
return kitutil.JsonRawMessageToString(data)
|
||||
}
|
||||
+23
-4
@@ -2,6 +2,9 @@ package common
|
||||
|
||||
import "strings"
|
||||
|
||||
// ZImagePromptExtendMultiplier is the legacy Ali image request surcharge.
|
||||
const ZImagePromptExtendMultiplier = 2
|
||||
|
||||
var (
|
||||
// OpenAIResponseOnlyModels is a list of models that are only available for OpenAI responses.
|
||||
OpenAIResponseOnlyModels = []string{
|
||||
@@ -12,7 +15,20 @@ var (
|
||||
ImageGenerationModels = []string{
|
||||
"dall-e-3",
|
||||
"dall-e-2",
|
||||
"gpt-image-1",
|
||||
"prefix:dall-e", // Deprecated upstream models; retained for compatible routes.
|
||||
"gpt-image-",
|
||||
"qwen-image",
|
||||
"z-image",
|
||||
"wan2.7-image-pro",
|
||||
"wan2.7-image",
|
||||
"wan2.6-image",
|
||||
"wan2.6-t2i",
|
||||
"wan2.5-t2i-preview",
|
||||
"wan2.2-t2i-flash",
|
||||
"wan2.2-t2i-plus",
|
||||
"wanx2.1-t2i-turbo",
|
||||
"wanx2.1-t2i-plus",
|
||||
"wanx2.0-t2i-turbo",
|
||||
"prefix:imagen-",
|
||||
"flux-",
|
||||
"flux.1-",
|
||||
@@ -38,10 +54,13 @@ func IsOpenAIResponseOnlyModel(modelName string) bool {
|
||||
func IsImageGenerationModel(modelName string) bool {
|
||||
modelName = strings.ToLower(modelName)
|
||||
for _, m := range ImageGenerationModels {
|
||||
if strings.Contains(modelName, m) {
|
||||
return true
|
||||
if prefix, ok := strings.CutPrefix(m, "prefix:"); ok {
|
||||
if strings.HasPrefix(modelName, prefix) {
|
||||
return true
|
||||
}
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(m, "prefix:") && strings.HasPrefix(modelName, strings.TrimPrefix(m, "prefix:")) {
|
||||
if strings.Contains(modelName, m) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
package common_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestWanEndpointsDistinguishImagesFromVideos(t *testing.T) {
|
||||
for _, name := range []string{
|
||||
"wan2.7-image-pro", "wan2.7-image", "wan2.6-image", "wan2.6-t2i",
|
||||
"wan2.5-t2i-preview", "wan2.2-t2i-flash", "wan2.2-t2i-plus",
|
||||
"wanx2.1-t2i-turbo", "wanx2.1-t2i-plus", "wanx2.0-t2i-turbo",
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
assert.Contains(t, common.GetEndpointTypesByChannelType(constant.ChannelTypeAli, name), constant.EndpointTypeImageGeneration)
|
||||
})
|
||||
}
|
||||
for _, name := range []string{
|
||||
"wanx2.1-t2v-plus", "wanx2.1-t2v-turbo", "wanx2.1-i2v-plus", "wanx2.1-i2v-turbo",
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
assert.NotContains(t, common.GetEndpointTypesByChannelType(constant.ChannelTypeAli, name), constant.EndpointTypeImageGeneration)
|
||||
})
|
||||
}
|
||||
}
|
||||
+186
-45
@@ -1,70 +1,211 @@
|
||||
package common
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// rateLimitRequest stores one accepted request in a key's sliding window.
|
||||
type rateLimitRequest struct {
|
||||
next *rateLimitRequest
|
||||
timestamp int64
|
||||
}
|
||||
|
||||
// rateLimitQueue keeps accepted requests ordered from oldest to newest.
|
||||
type rateLimitQueue struct {
|
||||
head *rateLimitRequest
|
||||
tail *rateLimitRequest
|
||||
length int
|
||||
}
|
||||
|
||||
// append adds an accepted request without preallocating for the configured limit.
|
||||
func (q *rateLimitQueue) append(timestamp int64) {
|
||||
request := &rateLimitRequest{timestamp: timestamp}
|
||||
if q.tail == nil {
|
||||
q.head = request
|
||||
q.tail = request
|
||||
} else {
|
||||
q.tail.next = request
|
||||
q.tail = request
|
||||
}
|
||||
q.length++
|
||||
}
|
||||
|
||||
// removeExpired releases every expired request at the front of the queue.
|
||||
func (q *rateLimitQueue) removeExpired(now int64, duration int64) {
|
||||
if q.head == nil || now-q.head.timestamp < duration {
|
||||
return
|
||||
}
|
||||
|
||||
// Requests are time ordered, so an expired tail means the whole queue expired.
|
||||
if now-q.tail.timestamp >= duration {
|
||||
q.clear()
|
||||
return
|
||||
}
|
||||
|
||||
for now-q.head.timestamp >= duration {
|
||||
expired := q.head
|
||||
q.head = expired.next
|
||||
expired.next = nil
|
||||
q.length--
|
||||
}
|
||||
}
|
||||
|
||||
// clear releases the whole request chain when its key expires.
|
||||
func (q *rateLimitQueue) clear() {
|
||||
q.head = nil
|
||||
q.tail = nil
|
||||
q.length = 0
|
||||
}
|
||||
|
||||
// rateLimitEntry is both a rate-limit bucket and a node in the key-level LRU.
|
||||
type rateLimitEntry struct {
|
||||
lastActive time.Time
|
||||
element *list.Element
|
||||
requests rateLimitQueue
|
||||
key string
|
||||
}
|
||||
|
||||
// InMemoryRateLimiter implements a sliding-window limiter with idle-key eviction.
|
||||
// Reservations are tracked apart from the LRU so an in-flight request survives
|
||||
// its key being evicted while it is still running.
|
||||
type InMemoryRateLimiter struct {
|
||||
store map[string]*[]int64
|
||||
store map[string]*rateLimitEntry
|
||||
lru *list.List
|
||||
reservations map[string]int
|
||||
mutex sync.Mutex
|
||||
expirationDuration time.Duration
|
||||
}
|
||||
|
||||
// Init initializes the limiter once. Repeated calls leave the first configuration unchanged.
|
||||
func (l *InMemoryRateLimiter) Init(expirationDuration time.Duration) {
|
||||
if l.store == nil {
|
||||
l.mutex.Lock()
|
||||
if l.store == nil {
|
||||
l.store = make(map[string]*[]int64)
|
||||
l.expirationDuration = expirationDuration
|
||||
if expirationDuration > 0 {
|
||||
go l.clearExpiredItems()
|
||||
}
|
||||
}
|
||||
l.mutex.Unlock()
|
||||
l.mutex.Lock()
|
||||
defer l.mutex.Unlock()
|
||||
|
||||
if l.store != nil {
|
||||
return
|
||||
}
|
||||
|
||||
l.store = make(map[string]*rateLimitEntry)
|
||||
l.lru = list.New()
|
||||
l.reservations = make(map[string]int)
|
||||
l.expirationDuration = expirationDuration
|
||||
if expirationDuration > 0 {
|
||||
go l.clearExpiredItems(time.NewTicker(expirationDuration).C)
|
||||
}
|
||||
}
|
||||
|
||||
func (l *InMemoryRateLimiter) clearExpiredItems() {
|
||||
for {
|
||||
time.Sleep(l.expirationDuration)
|
||||
l.mutex.Lock()
|
||||
now := time.Now().Unix()
|
||||
for key := range l.store {
|
||||
queue := l.store[key]
|
||||
size := len(*queue)
|
||||
if size == 0 || now-(*queue)[size-1] > int64(l.expirationDuration.Seconds()) {
|
||||
delete(l.store, key)
|
||||
}
|
||||
}
|
||||
l.mutex.Unlock()
|
||||
// clearExpiredItems periodically removes expired entries from the LRU tail.
|
||||
func (l *InMemoryRateLimiter) clearExpiredItems(ticks <-chan time.Time) {
|
||||
for now := range ticks {
|
||||
l.deleteExpiredEntries(now)
|
||||
}
|
||||
}
|
||||
|
||||
// deleteExpiredEntries walks only the oldest LRU entries and stops at the first active key.
|
||||
func (l *InMemoryRateLimiter) deleteExpiredEntries(now time.Time) {
|
||||
if l.expirationDuration <= 0 {
|
||||
return
|
||||
}
|
||||
|
||||
l.mutex.Lock()
|
||||
defer l.mutex.Unlock()
|
||||
|
||||
for {
|
||||
oldest := l.lru.Back()
|
||||
if oldest == nil {
|
||||
return
|
||||
}
|
||||
|
||||
entry := oldest.Value.(*rateLimitEntry)
|
||||
if now.Sub(entry.lastActive) < l.expirationDuration {
|
||||
return
|
||||
}
|
||||
|
||||
delete(l.store, entry.key)
|
||||
l.lru.Remove(oldest)
|
||||
entry.element = nil
|
||||
entry.requests.clear()
|
||||
}
|
||||
}
|
||||
|
||||
// touchEntry returns the key's bucket, creating it on demand, and marks it as
|
||||
// the most recently active key. The caller must hold the mutex.
|
||||
func (l *InMemoryRateLimiter) touchEntry(key string, now time.Time) *rateLimitEntry {
|
||||
entry, ok := l.store[key]
|
||||
if !ok {
|
||||
entry = &rateLimitEntry{
|
||||
key: key,
|
||||
lastActive: now,
|
||||
}
|
||||
entry.element = l.lru.PushFront(entry)
|
||||
l.store[key] = entry
|
||||
return entry
|
||||
}
|
||||
entry.lastActive = now
|
||||
l.lru.MoveToFront(entry.element)
|
||||
return entry
|
||||
}
|
||||
|
||||
// Request parameter duration's unit is seconds
|
||||
func (l *InMemoryRateLimiter) Request(key string, maxRequestNum int, duration int64) bool {
|
||||
l.mutex.Lock()
|
||||
defer l.mutex.Unlock()
|
||||
// [old <-- new]
|
||||
queue, ok := l.store[key]
|
||||
now := time.Now().Unix()
|
||||
if ok {
|
||||
if len(*queue) < maxRequestNum {
|
||||
*queue = append(*queue, now)
|
||||
return true
|
||||
} else {
|
||||
if now-(*queue)[0] >= duration {
|
||||
*queue = (*queue)[1:]
|
||||
*queue = append(*queue, now)
|
||||
return true
|
||||
} else {
|
||||
return false
|
||||
}
|
||||
}
|
||||
} else {
|
||||
s := make([]int64, 0, maxRequestNum)
|
||||
l.store[key] = &s
|
||||
*(l.store[key]) = append(*(l.store[key]), now)
|
||||
|
||||
now := time.Now()
|
||||
entry := l.touchEntry(key, now)
|
||||
entry.requests.removeExpired(now.Unix(), duration)
|
||||
|
||||
allowed := entry.requests.length < maxRequestNum
|
||||
if allowed {
|
||||
entry.requests.append(now.Unix())
|
||||
}
|
||||
return true
|
||||
|
||||
return allowed
|
||||
}
|
||||
|
||||
// RateLimitReservation counts an active request towards admission until its
|
||||
// outcome is known. Failed requests release the slot without consuming quota.
|
||||
type RateLimitReservation struct {
|
||||
limiter *InMemoryRateLimiter
|
||||
key string
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
// Reserve admits a request when the key's accepted requests plus its in-flight
|
||||
// reservations stay below maxRequests. It returns nil when the key is saturated.
|
||||
func (l *InMemoryRateLimiter) Reserve(key string, maxRequests int, duration int64) *RateLimitReservation {
|
||||
l.mutex.Lock()
|
||||
defer l.mutex.Unlock()
|
||||
|
||||
now := time.Now()
|
||||
entry := l.touchEntry(key, now)
|
||||
entry.requests.removeExpired(now.Unix(), duration)
|
||||
if entry.requests.length+l.reservations[key] >= maxRequests {
|
||||
return nil
|
||||
}
|
||||
l.reservations[key]++
|
||||
return &RateLimitReservation{limiter: l, key: key}
|
||||
}
|
||||
|
||||
// Complete releases the reservation. A successful outcome records the request
|
||||
// in the key's sliding window; a failed one only frees the slot.
|
||||
func (r *RateLimitReservation) Complete(success bool) {
|
||||
if r == nil {
|
||||
return
|
||||
}
|
||||
r.once.Do(func() {
|
||||
l := r.limiter
|
||||
l.mutex.Lock()
|
||||
defer l.mutex.Unlock()
|
||||
l.reservations[r.key]--
|
||||
if l.reservations[r.key] == 0 {
|
||||
delete(l.reservations, r.key)
|
||||
}
|
||||
if success {
|
||||
now := time.Now()
|
||||
l.touchEntry(r.key, now).requests.append(now.Unix())
|
||||
}
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,150 @@
|
||||
package common
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestInMemoryRateLimiterAllocatesTimestampsOnDemand(t *testing.T) {
|
||||
var limiter InMemoryRateLimiter
|
||||
limiter.Init(0)
|
||||
|
||||
require.True(t, limiter.Request("client", 1_000_000, 60))
|
||||
|
||||
entry := limiter.store["client"]
|
||||
require.NotNil(t, entry)
|
||||
assert.Equal(t, 1, entry.requests.length)
|
||||
assert.Same(t, entry.requests.head, entry.requests.tail)
|
||||
assert.Nil(t, entry.requests.head.next)
|
||||
|
||||
limiter.deleteExpiredEntries(time.Now().Add(time.Hour))
|
||||
assert.Contains(t, limiter.store, "client")
|
||||
}
|
||||
|
||||
func TestInMemoryRateLimiterEnforcesWindowBoundary(t *testing.T) {
|
||||
var queue rateLimitQueue
|
||||
queue.append(100)
|
||||
queue.append(100)
|
||||
|
||||
queue.removeExpired(159, 60)
|
||||
assert.Equal(t, 2, queue.length)
|
||||
|
||||
queue.removeExpired(160, 60)
|
||||
assert.Zero(t, queue.length)
|
||||
assert.Nil(t, queue.head)
|
||||
assert.Nil(t, queue.tail)
|
||||
}
|
||||
|
||||
func TestInMemoryRateLimiterRemovesAllExpiredRequestsPerCall(t *testing.T) {
|
||||
var queue rateLimitQueue
|
||||
for range 5 {
|
||||
queue.append(100)
|
||||
}
|
||||
queue.append(110)
|
||||
|
||||
queue.removeExpired(110, 10)
|
||||
assert.Equal(t, 1, queue.length)
|
||||
require.NotNil(t, queue.head)
|
||||
assert.EqualValues(t, 110, queue.head.timestamp)
|
||||
assert.Same(t, queue.head, queue.tail)
|
||||
}
|
||||
|
||||
func TestInMemoryRateLimiterKeepsNonExpiredRequests(t *testing.T) {
|
||||
var queue rateLimitQueue
|
||||
for range 3 {
|
||||
queue.append(100)
|
||||
}
|
||||
|
||||
queue.removeExpired(109, 10)
|
||||
assert.Equal(t, 3, queue.length)
|
||||
assert.NotNil(t, queue.head)
|
||||
assert.NotNil(t, queue.tail)
|
||||
}
|
||||
|
||||
func TestInMemoryRateLimiterEvictsOnlyExpiredLRUTail(t *testing.T) {
|
||||
var limiter InMemoryRateLimiter
|
||||
limiter.Init(0)
|
||||
limiter.expirationDuration = 10 * time.Second
|
||||
|
||||
require.True(t, limiter.Request("active", 1, 100))
|
||||
require.True(t, limiter.Request("idle", 1, 100))
|
||||
assert.False(t, limiter.Request("active", 1, 100))
|
||||
|
||||
assert.Equal(t, "active", limiter.lru.Front().Value.(*rateLimitEntry).key)
|
||||
assert.Equal(t, "idle", limiter.lru.Back().Value.(*rateLimitEntry).key)
|
||||
|
||||
idleEntry := limiter.store["idle"]
|
||||
now := time.Now()
|
||||
idleEntry.lastActive = now.Add(-10 * time.Second)
|
||||
limiter.store["active"].lastActive = now
|
||||
limiter.deleteExpiredEntries(now)
|
||||
|
||||
assert.NotContains(t, limiter.store, "idle")
|
||||
assert.Contains(t, limiter.store, "active")
|
||||
assert.Equal(t, 1, limiter.lru.Len())
|
||||
assert.Nil(t, idleEntry.element)
|
||||
assert.Nil(t, idleEntry.requests.head)
|
||||
assert.Nil(t, idleEntry.requests.tail)
|
||||
assert.Zero(t, idleEntry.requests.length)
|
||||
}
|
||||
|
||||
func TestInMemoryRateLimiterCleanupTicksEvictExpiredKeys(t *testing.T) {
|
||||
var limiter InMemoryRateLimiter
|
||||
limiter.Init(0)
|
||||
limiter.expirationDuration = 10 * time.Second
|
||||
require.True(t, limiter.Request("idle", 1, 100))
|
||||
|
||||
now := time.Now()
|
||||
limiter.store["idle"].lastActive = now.Add(-10 * time.Second)
|
||||
ticks := make(chan time.Time, 1)
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
limiter.clearExpiredItems(ticks)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
ticks <- now
|
||||
close(ticks)
|
||||
<-done
|
||||
|
||||
assert.Empty(t, limiter.store)
|
||||
assert.Zero(t, limiter.lru.Len())
|
||||
}
|
||||
|
||||
func TestInMemoryRateLimiterInitKeepsFirstConfiguration(t *testing.T) {
|
||||
var limiter InMemoryRateLimiter
|
||||
limiter.Init(time.Hour)
|
||||
limiter.Init(0)
|
||||
|
||||
assert.Equal(t, time.Hour, limiter.expirationDuration)
|
||||
require.True(t, limiter.Request("client", 1, 60))
|
||||
limiter.deleteExpiredEntries(time.Now().Add(time.Hour))
|
||||
assert.NotContains(t, limiter.store, "client")
|
||||
}
|
||||
|
||||
func TestInMemoryRateLimiterConcurrentInitializationAndRequests(t *testing.T) {
|
||||
var limiter InMemoryRateLimiter
|
||||
var allowed atomic.Int64
|
||||
var waitGroup sync.WaitGroup
|
||||
|
||||
for range 100 {
|
||||
waitGroup.Add(1)
|
||||
go func() {
|
||||
defer waitGroup.Done()
|
||||
limiter.Init(0)
|
||||
if limiter.Request("client", 10, 60) {
|
||||
allowed.Add(1)
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
waitGroup.Wait()
|
||||
assert.EqualValues(t, 10, allowed.Load())
|
||||
assert.Len(t, limiter.store, 1)
|
||||
assert.Equal(t, 10, limiter.store["client"].requests.length)
|
||||
}
|
||||
@@ -59,10 +59,14 @@ const (
|
||||
ChannelTypeSub2API = 59
|
||||
ChannelTypeNewAPI = 60
|
||||
ChannelTypeTaskPlugin = 61
|
||||
ChannelTypeVLLM = 62
|
||||
ChannelTypeSGLang = 63
|
||||
ChannelTypeDummy // this one is only for count, do not add any channel after this
|
||||
|
||||
)
|
||||
|
||||
// ChannelBaseURLs 保存各渠道类型的内置默认 Base URL。
|
||||
// 非空值会通过 /api/channel/default_base_urls 下发到前端,作为渠道表单的 API 地址占位提示。
|
||||
var ChannelBaseURLs = []string{
|
||||
"", // 0
|
||||
"https://api.openai.com", // 1
|
||||
@@ -126,6 +130,8 @@ var ChannelBaseURLs = []string{
|
||||
"", //59
|
||||
"", //60
|
||||
"", //61
|
||||
"", //62
|
||||
"", //63
|
||||
}
|
||||
|
||||
func GetChannelBaseURL(channelType int) string {
|
||||
@@ -194,6 +200,8 @@ var ChannelTypeNames = map[int]string{
|
||||
ChannelTypeSub2API: "Sub2API",
|
||||
ChannelTypeNewAPI: "New API",
|
||||
ChannelTypeTaskPlugin: "Task Plugin",
|
||||
ChannelTypeVLLM: "vLLM",
|
||||
ChannelTypeSGLang: "SGLang",
|
||||
}
|
||||
|
||||
func GetChannelTypeName(channelType int) string {
|
||||
@@ -226,3 +234,13 @@ var ChannelSpecialBases = map[string]ChannelSpecialBase{
|
||||
OpenAIBaseURL: "https://ark.cn-beijing.volces.com/api/coding/v3",
|
||||
},
|
||||
}
|
||||
|
||||
// IsAdvancedCustomChannel includes named channels backed by route presets.
|
||||
func IsAdvancedCustomChannel(channelType int) bool {
|
||||
switch channelType {
|
||||
case ChannelTypeAdvancedCustom, ChannelTypeVLLM, ChannelTypeSGLang:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
+17
-12
@@ -11,16 +11,20 @@ const (
|
||||
ContextKeyRequestStartTime ContextKey = "request_start_time"
|
||||
|
||||
/* token related keys */
|
||||
ContextKeyTokenUnlimited ContextKey = "token_unlimited_quota"
|
||||
ContextKeyTokenKey ContextKey = "token_key"
|
||||
ContextKeyTokenId ContextKey = "token_id"
|
||||
ContextKeyTokenGroup ContextKey = "token_group"
|
||||
ContextKeyOriginTasks ContextKey = "origin_tasks"
|
||||
ContextKeyChannelConstraints ContextKey = "channel_constraints"
|
||||
ContextKeyTokenModelLimitEnabled ContextKey = "token_model_limit_enabled"
|
||||
ContextKeyTokenModelLimit ContextKey = "token_model_limit"
|
||||
ContextKeyTokenCrossGroupRetry ContextKey = "token_cross_group_retry"
|
||||
ContextKeyTokenAutoGroups ContextKey = "token_auto_groups"
|
||||
ContextKeyTokenUnlimited ContextKey = "token_unlimited_quota"
|
||||
ContextKeyTokenKey ContextKey = "token_key"
|
||||
ContextKeyTokenId ContextKey = "token_id"
|
||||
ContextKeyTokenGroup ContextKey = "token_group"
|
||||
ContextKeyOriginTasks ContextKey = "origin_tasks"
|
||||
ContextKeyChannelConstraints ContextKey = "channel_constraints"
|
||||
ContextKeyTokenModelLimitEnabled ContextKey = "token_model_limit_enabled"
|
||||
ContextKeyTokenModelLimit ContextKey = "token_model_limit"
|
||||
ContextKeyTokenCrossGroupRetry ContextKey = "token_cross_group_retry"
|
||||
ContextKeyTokenAstraDetectionDisabled ContextKey = "token_astra_detection_disabled"
|
||||
ContextKeyTokenAstraAutoSwitch ContextKey = "token_astra_auto_switch"
|
||||
ContextKeyTokenAutoGroups ContextKey = "token_auto_groups"
|
||||
ContextKeyTokenGroupPool ContextKey = "token_group_pool"
|
||||
ContextKeyTokenFailedGroups ContextKey = "token_failed_groups"
|
||||
|
||||
/* channel related keys */
|
||||
ContextKeyChannelId ContextKey = "channel_id"
|
||||
@@ -66,8 +70,9 @@ const (
|
||||
ContextKeyAdminRejectReason ContextKey = "admin_reject_reason"
|
||||
|
||||
// ContextKeyLanguage stores the user's language preference for i18n
|
||||
ContextKeyLanguage ContextKey = "language"
|
||||
ContextKeyIsStream ContextKey = "is_stream"
|
||||
ContextKeyLanguage ContextKey = "language"
|
||||
ContextKeyIsStream ContextKey = "is_stream"
|
||||
ContextKeyResponseStreamStatus ContextKey = "response_stream_status"
|
||||
|
||||
// ContextKeyAuditLogged marks that the current request has already recorded
|
||||
// a manage/operation audit log inside the handler. When set, the admin-audit
|
||||
|
||||
@@ -0,0 +1,256 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math/big"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
|
||||
"github.com/QuantumNous/new-api/relay"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type astraSamplingHandler struct{}
|
||||
|
||||
type astraSamplingTaskPayload struct {
|
||||
Force bool `json:"force,omitempty"`
|
||||
}
|
||||
|
||||
func (astraSamplingHandler) Type() string { return service.AstraSamplingTaskType }
|
||||
func (astraSamplingHandler) Enabled() bool { return service.AstraSamplingEnabled() }
|
||||
func (astraSamplingHandler) Interval() time.Duration { return service.AstraSamplingSchedulerInterval() }
|
||||
func (astraSamplingHandler) ScheduleFromCreation() bool { return true }
|
||||
func (astraSamplingHandler) NewPayload() any { return nil }
|
||||
|
||||
func (astraSamplingHandler) ScheduleDue(ctx context.Context, now time.Time) (bool, error) {
|
||||
channels, err := service.AstraSamplingChannelsDue(ctx, now)
|
||||
return len(channels) > 0, err
|
||||
}
|
||||
|
||||
func (astraSamplingHandler) Run(parent context.Context, task *model.SystemTask, runnerID string) {
|
||||
ctx, cancel := context.WithTimeout(parent, 25*time.Minute)
|
||||
defer cancel()
|
||||
payload := astraSamplingTaskPayload{}
|
||||
if err := task.DecodePayload(&payload); err != nil {
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err)
|
||||
return
|
||||
}
|
||||
channels, err := service.AstraSamplingChannelsDue(ctx, time.Now())
|
||||
if payload.Force {
|
||||
channels, err = service.AstraSamplingChannels(ctx)
|
||||
}
|
||||
if err != nil {
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, errors.New("channel query failed"))
|
||||
return
|
||||
}
|
||||
state := service.AstraSamplingState{Rounds: []service.AstraSamplingRound{}}
|
||||
var mu sync.Mutex
|
||||
var persistErr error
|
||||
// Persist before sending, and after every probe. A failed lease or restart
|
||||
// never resumes a paid probe; unfinished records remain visible as errors.
|
||||
persist := func() bool {
|
||||
if persistErr != nil {
|
||||
return false
|
||||
}
|
||||
persistErr = model.UpdateSystemTaskState(task.TaskID, runnerID, state)
|
||||
if persistErr != nil {
|
||||
cancel()
|
||||
}
|
||||
return persistErr == nil
|
||||
}
|
||||
jobs := make(chan int)
|
||||
var workers sync.WaitGroup
|
||||
for worker := 0; worker < 2; worker++ {
|
||||
workers.Add(1)
|
||||
go func() {
|
||||
defer workers.Done()
|
||||
for id := range jobs {
|
||||
if ctx.Err() != nil {
|
||||
continue
|
||||
}
|
||||
mu.Lock()
|
||||
index := len(state.Rounds)
|
||||
state.Rounds = append(state.Rounds, service.AstraSamplingRound{ID: common.GetRandomString(16), ChannelID: id, StartedAt: time.Now().Unix(), Version: fingerprint.Version, Result: fingerprint.Result{Status: "running"}})
|
||||
ok := persist()
|
||||
mu.Unlock()
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
channel, loadErr := model.GetChannelById(id, true)
|
||||
samples := []fingerprint.Sample{}
|
||||
reason := ""
|
||||
usage := []fingerprint.ProbeResponse{}
|
||||
if loadErr != nil || channel.Status != common.ChannelStatusEnabled {
|
||||
reason = "channel_unavailable"
|
||||
} else {
|
||||
for probe := 0; probe < 3 && ctx.Err() == nil; probe++ {
|
||||
mu.Lock()
|
||||
state.Rounds[index].Sent++
|
||||
ok = persist()
|
||||
mu.Unlock()
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
answer, count, probeErr := runAstraFingerprintProbe(ctx, channel)
|
||||
usage = append(usage, answer)
|
||||
if probeErr != nil {
|
||||
if errors.Is(probeErr, fingerprint.ErrInvalidSample) {
|
||||
reason = "invalid_sample"
|
||||
} else {
|
||||
reason = "transport_error"
|
||||
break
|
||||
}
|
||||
} else {
|
||||
samples = append(samples, fingerprint.Sample{Text: answer.Text, Count: count})
|
||||
}
|
||||
mu.Lock()
|
||||
state.Rounds[index].Usage = append([]fingerprint.ProbeResponse(nil), usage...)
|
||||
ok = persist()
|
||||
mu.Unlock()
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
result, scoreErr := fingerprint.Score(samples)
|
||||
if scoreErr != nil {
|
||||
reason = "scoring_error"
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
reason = "interrupted"
|
||||
}
|
||||
if reason != "" && reason != "invalid_sample" {
|
||||
result.Status = "error"
|
||||
}
|
||||
mu.Lock()
|
||||
state.Rounds[index].Result = result
|
||||
state.Rounds[index].Reason = reason
|
||||
state.Rounds[index].Usage = usage
|
||||
state.Rounds[index].FinishedAt = time.Now().Unix()
|
||||
if persist() && !payload.Force {
|
||||
if err := model.ApplyDetectionScheduleResult(ctx, id, "fingerprint", state.Rounds[index].ID, state.Rounds[index].FinishedAt, service.FingerprintScheduleOutcome(service.AstraSamplingResultStatus(result))); err != nil {
|
||||
persistErr = err
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
mu.Unlock()
|
||||
}
|
||||
}()
|
||||
}
|
||||
for _, channel := range channels {
|
||||
select {
|
||||
case jobs <- channel.Id:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
close(jobs)
|
||||
workers.Wait()
|
||||
if persistErr != nil || ctx.Err() != nil {
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, errors.New("sampling interrupted"))
|
||||
return
|
||||
}
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, map[string]any{"channels": len(state.Rounds), "model": fingerprint.Model}, nil)
|
||||
}
|
||||
|
||||
func runAstraFingerprintProbe(parent context.Context, channel *model.Channel, paid ...bool) (fingerprint.ProbeResponse, int, error) {
|
||||
ctx, cancel := context.WithTimeout(parent, 180*time.Second)
|
||||
defer cancel()
|
||||
seed, err := rand.Int(rand.Reader, big.NewInt(82))
|
||||
if err != nil {
|
||||
return fingerprint.ProbeResponse{}, 0, err
|
||||
}
|
||||
count := 292 + int(seed.Int64()/2)
|
||||
prompt := fmt.Sprintf("Directly choose %d integers from 1 through 355, inclusive. Allow repeats. Do not sort, balance frequencies, repair duplicates, or deliberately create an arithmetic pattern.\nReply directly with one JSON integer array and no explanation. Do not call tools, read files, execute code or ask another model. Do not continue the preceding task.", count)
|
||||
if seed.Int64()%2 == 0 {
|
||||
prompt = fmt.Sprintf("直接选择 %d 个 1 到 355(含边界)的整数。允许重复;不要排序、平衡频数、修复重复,也不要刻意构造等差规律。\n直接回答一个 JSON 整数数组,不要解释。不要调用工具、读文件、运行代码或让其他模型代答;不要继续之前的任务。", count)
|
||||
}
|
||||
input, _ := common.Marshal([]map[string]any{{"role": "user", "content": []map[string]string{{"type": "input_text", "text": prompt}}}})
|
||||
stream := true
|
||||
limit := uint(4096)
|
||||
request := &dto.OpenAIResponsesRequest{Model: fingerprint.Model, Input: input, Instructions: []byte(`""`), Stream: &stream, MaxOutputTokens: &limit, Store: []byte("false"), Tools: []byte("[]"), ToolChoice: []byte(`"none"`)}
|
||||
encoded, err := common.Marshal(request)
|
||||
if err != nil {
|
||||
return fingerprint.ProbeResponse{}, count, err
|
||||
}
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequestWithContext(ctx, http.MethodPost, "/v1/responses", bytes.NewReader(encoded))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
if apiErr := middleware.SetupContextForSelectedChannel(c, channel, fingerprint.Model); apiErr != nil {
|
||||
return fingerprint.ProbeResponse{}, count, errors.New("channel setup failed")
|
||||
}
|
||||
info, err := relaycommon.GenRelayInfo(c, types.RelayFormatOpenAIResponses, request, nil)
|
||||
if err != nil {
|
||||
return fingerprint.ProbeResponse{}, count, err
|
||||
}
|
||||
info.IsChannelTest = true
|
||||
info.DisablePing = true
|
||||
info.AstraDetectionDisabled = true
|
||||
info.AstraAutoSwitch = false
|
||||
adaptor, body, closer, apiErr := relay.PrepareResponsesRequest(c, info, request)
|
||||
if closer != nil {
|
||||
defer closer.Close()
|
||||
}
|
||||
if apiErr != nil {
|
||||
return fingerprint.ProbeResponse{}, count, errors.New("probe preparation failed")
|
||||
}
|
||||
if len(paid) > 0 && paid[0] {
|
||||
bounded, boundedCloser, boundErr := boundDetectionBody(body, request)
|
||||
if boundErr != nil {
|
||||
return fingerprint.ProbeResponse{}, count, boundErr
|
||||
}
|
||||
defer boundedCloser.Close()
|
||||
body = bounded
|
||||
}
|
||||
response, err := adaptor.DoRequest(c, info, body)
|
||||
if err != nil {
|
||||
return fingerprint.ProbeResponse{}, count, errors.New("probe transport failed")
|
||||
}
|
||||
resp, ok := response.(*http.Response)
|
||||
if !ok || resp == nil {
|
||||
return fingerprint.ProbeResponse{}, count, errors.New("unsupported probe transport")
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return fingerprint.ProbeResponse{}, count, errors.New("probe HTTP failure")
|
||||
}
|
||||
// Deliberately bypass DoResponse: this is a platform-funded probe, with no
|
||||
// customer token, wallet charge, session binding or generated response relay.
|
||||
answer, err := fingerprint.ReadResponse(resp.Body, resp.Header.Get("Content-Type"))
|
||||
return answer, count, err
|
||||
}
|
||||
|
||||
func TestAstraFingerprint(c *gin.Context) {
|
||||
if !service.AstraSamplingEnabled() {
|
||||
c.JSON(http.StatusConflict, gin.H{"success": false, "message": "Astra sampling is disabled"})
|
||||
return
|
||||
}
|
||||
latest, err := model.GetLatestSystemTasks([]string{service.AstraSamplingTaskType})
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if last := latest[service.AstraSamplingTaskType]; last != nil && time.Now().Unix()-last.CreatedAt < 300 {
|
||||
c.JSON(http.StatusTooManyRequests, gin.H{"success": false, "message": "Please wait five minutes between sampling rounds"})
|
||||
return
|
||||
}
|
||||
task, created, err := service.EnqueueSystemTask(service.AstraSamplingTaskType, astraSamplingTaskPayload{Force: true})
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "data": gin.H{"task_id": task.TaskID, "created": created}})
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestAstraSamplingUsesIsolatedResponsesWithoutGuard(t *testing.T) {
|
||||
requests := 0
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
requests++
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
if r.URL.Path != "/v1/responses" || gjson.GetBytes(body, "model").String() != fingerprint.Model || !gjson.GetBytes(body, "stream").Bool() {
|
||||
t.Errorf("wrong probe request %s %s", r.URL.Path, body)
|
||||
}
|
||||
if gjson.GetBytes(body, "previous_response_id").Exists() || gjson.GetBytes(body, "conversation").Exists() {
|
||||
t.Error("probe carried customer session")
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = io.WriteString(w, "data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"model\":\"gpt-5.6-luna\",\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"[1,2,3]\"}]}]}}\n\n")
|
||||
}))
|
||||
defer server.Close()
|
||||
url := server.URL
|
||||
channel := &model.Channel{Id: 1, Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: fingerprint.Model}
|
||||
r, count, err := runAstraFingerprintProbe(context.Background(), channel)
|
||||
if err != nil || r.DeclaredModel != "gpt-5.6-luna" || count < 292 || count > 332 || requests != 1 {
|
||||
t.Fatalf("probe %+v count=%d requests=%d err=%v", r, count, requests, err)
|
||||
}
|
||||
}
|
||||
@@ -1,8 +1,10 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
@@ -43,6 +45,11 @@ var auditContentTemplates = map[string]string{
|
||||
"user.reset_passkey": "Reset the user passkey",
|
||||
"option.update": "Updated system setting ${key}",
|
||||
|
||||
"option.passkey_domains": "Updated Passkey domains: removed ${domains}; affected ${known}; unknown ${unknown}",
|
||||
"option.passkey_domains_confirmed": "Confirmed removal of Passkey domains: ${domains}; affected ${known}; unknown ${unknown}",
|
||||
"option.passkey_domains_blocked": "Passkey domain change blocked: ${domains}; affected ${known}; unknown ${unknown}",
|
||||
"option.passkey_domains_failed": "Passkey domain update failed",
|
||||
|
||||
"channel.create": "Created channel ${name} (type ${type}, count ${count})",
|
||||
"channel.update": "Updated channel ${name} (ID: ${id})",
|
||||
"channel.delete": "Deleted channel ${name} (ID: ${id})",
|
||||
@@ -65,6 +72,33 @@ var auditContentTemplates = map[string]string{
|
||||
"subscription.user_plan_reset": "Reset active plan ${plan_id} subscriptions for user ${target_user_id}",
|
||||
}
|
||||
|
||||
func recordPasskeyDomainAudit(c *gin.Context, change *model.PasskeyDomainChange, confirmed bool, err error) {
|
||||
confirmed = confirmed && err == nil && change != nil && len(change.RemovedRPIDs) > 0
|
||||
params := map[string]any{"success": err == nil, "confirmed": confirmed}
|
||||
if change != nil {
|
||||
params["domains"] = strings.Join(change.RemovedRPIDs, ", ")
|
||||
params["removed_rp_ids"] = change.RemovedRPIDs
|
||||
params["known"] = change.AffectedCredentials
|
||||
params["unknown"] = change.UnknownCredentials
|
||||
params["previous_rp_id"] = change.PreviousRPID
|
||||
params["effective_rp_id"] = change.EffectiveRPID
|
||||
}
|
||||
action := "option.passkey_domains"
|
||||
if errors.Is(err, model.ErrPasskeyDomainRemovalConfirmation) {
|
||||
action = "option.passkey_domains_blocked"
|
||||
} else if err != nil {
|
||||
action = "option.passkey_domains_failed"
|
||||
} else if confirmed && change != nil && len(change.RemovedRPIDs) > 0 {
|
||||
action = "option.passkey_domains_confirmed"
|
||||
}
|
||||
auditInfo := &model.AuditRequestInfo{
|
||||
Method: c.Request.Method, Route: c.FullPath(), Path: c.FullPath(),
|
||||
Status: c.Writer.Status(), Success: err == nil,
|
||||
}
|
||||
model.RecordOperationAuditLog(c.GetInt("id"), c.GetInt("role"), auditContentEN(action, params), c.ClientIP(), action, params, auditOperatorInfo(c), auditInfo, c)
|
||||
markAuditLogged(c)
|
||||
}
|
||||
|
||||
// auditContentEN 按 action 模板渲染英文兜底文本;未登记的 action 退回 action 本身。
|
||||
func auditContentEN(action string, params map[string]any) string {
|
||||
tmpl, ok := auditContentTemplates[action]
|
||||
|
||||
@@ -76,7 +76,7 @@ func TestSecurityLoginCodeCompletesOnce(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
factor := &model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}
|
||||
require.NoError(t, model.DB.Create(factor).Error)
|
||||
challenge, err := service.StartLoginVerification(user, "password")
|
||||
challenge, err := service.StartLoginVerification(user, "password", nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, challenge)
|
||||
code, err := totp.GenerateCode(factor.Secret, time.Now())
|
||||
@@ -121,7 +121,7 @@ func TestSecurityLoginRejectsChangedOrExpiredAuthorization(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
factor := &model.TwoFA{UserId: user.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}
|
||||
require.NoError(t, model.DB.Create(factor).Error)
|
||||
challenge, err := service.StartLoginVerification(user, "password")
|
||||
challenge, err := service.StartLoginVerification(user, "password", nil)
|
||||
require.NoError(t, err)
|
||||
method := "2fa"
|
||||
switch change {
|
||||
@@ -183,7 +183,7 @@ func TestSecurityLoginPasskeyDoesNotRequireAdditionalTwoFA(t *testing.T) {
|
||||
require.Equal(t, "required", result.Data.Options.PublicKey.UserVerification)
|
||||
flowToken, challenge = result.Data.FlowToken, result.Data.Options.PublicKey.Challenge
|
||||
} else {
|
||||
pending, err := service.StartLoginVerification(user, "oauth:github")
|
||||
pending, err := service.StartLoginVerification(user, "oauth:github", nil)
|
||||
require.NoError(t, err)
|
||||
parentToken = pending.FlowToken
|
||||
flowToken, challenge = beginSecurityLoginPasskey(t, parentToken)
|
||||
@@ -225,7 +225,7 @@ func TestSecurityLoginPasskeyDoesNotRequireAdditionalTwoFA(t *testing.T) {
|
||||
func TestSecurityLoginPasskeyConcurrentCompletionCreatesOneSession(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
key := newSecurityLoginPasskey(t, user.Id)
|
||||
pending, err := service.StartLoginVerification(user, "password")
|
||||
pending, err := service.StartLoginVerification(user, "password", nil)
|
||||
require.NoError(t, err)
|
||||
requests := make([]string, 2)
|
||||
for index := range requests {
|
||||
@@ -266,7 +266,7 @@ func TestSecurityLoginPasskeyConcurrentCompletionCreatesOneSession(t *testing.T)
|
||||
func TestSecurityLoginSessionFailureRollsBackChallengeConsumption(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
key := newSecurityLoginPasskey(t, user.Id)
|
||||
pending, err := service.StartLoginVerification(user, "password")
|
||||
pending, err := service.StartLoginVerification(user, "password", nil)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, model.DB.Callback().Create().Before("gorm:create").Register("login_session_failure", func(tx *gorm.DB) {
|
||||
if tx.Statement.Table == "user_sessions" {
|
||||
@@ -410,7 +410,7 @@ func TestSecurityLoginPasskeyCannotCompleteAnotherChallenge(t *testing.T) {
|
||||
t.Run(fmt.Sprintf("other-user=%t", otherUser), func(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
key := newSecurityLoginPasskey(t, user.Id)
|
||||
first, err := service.StartLoginVerification(user, "password")
|
||||
first, err := service.StartLoginVerification(user, "password", nil)
|
||||
require.NoError(t, err)
|
||||
passkeyToken, challenge := beginSecurityLoginPasskey(t, first.FlowToken)
|
||||
if otherUser {
|
||||
@@ -418,7 +418,7 @@ func TestSecurityLoginPasskeyCannotCompleteAnotherChallenge(t *testing.T) {
|
||||
require.NoError(t, model.DB.Create(user).Error)
|
||||
newSecurityLoginPasskey(t, user.Id)
|
||||
}
|
||||
second, err := service.StartLoginVerification(user, "password")
|
||||
second, err := service.StartLoginVerification(user, "password", nil)
|
||||
require.NoError(t, err)
|
||||
body, err := common.Marshal(map[string]any{"flow_token": second.FlowToken, "passkey_flow_token": passkeyToken, "credential": securityPasskeyResponse(t, key, challenge, false, 0)})
|
||||
require.NoError(t, err)
|
||||
@@ -816,3 +816,310 @@ func TestOAuthBindProviderErrorConsumesSessionBoundFlow(t *testing.T) {
|
||||
assert.Zero(t, provider.exchangeCalls)
|
||||
assert.Zero(t, provider.userInfoCalls)
|
||||
}
|
||||
|
||||
// legacyGitHubOAuthProvider follows the GitHub provider contract: the numeric
|
||||
// account ID identifies the user, Extra["legacy_id"] carries the username and
|
||||
// the verified email list is served on demand. Lookups go through the real
|
||||
// github_id model queries.
|
||||
type legacyGitHubOAuthProvider struct {
|
||||
authFlowTestOAuthProvider
|
||||
providerUserID string
|
||||
legacyID string
|
||||
verifiedEmails []string
|
||||
verifiedEmailsErr error
|
||||
verifiedEmailCalls int
|
||||
}
|
||||
|
||||
func (*legacyGitHubOAuthProvider) GetName() string { return "GitHub" }
|
||||
func (provider *legacyGitHubOAuthProvider) GetUserInfo(context.Context, *oauth.OAuthToken) (*oauth.OAuthUser, error) {
|
||||
return &oauth.OAuthUser{ProviderUserID: provider.providerUserID, Username: provider.legacyID, Extra: map[string]any{"legacy_id": provider.legacyID}}, nil
|
||||
}
|
||||
func (provider *legacyGitHubOAuthProvider) GetVerifiedEmails(context.Context, *oauth.OAuthToken) ([]string, error) {
|
||||
provider.verifiedEmailCalls++
|
||||
return provider.verifiedEmails, provider.verifiedEmailsErr
|
||||
}
|
||||
func (*legacyGitHubOAuthProvider) IsUserIDTaken(providerUserID string) bool {
|
||||
return model.IsGitHubIdAlreadyTaken(providerUserID)
|
||||
}
|
||||
func (*legacyGitHubOAuthProvider) FillUserByProviderID(user *model.User, providerUserID string) error {
|
||||
user.GitHubId = providerUserID
|
||||
return user.FillUserByGitHubId()
|
||||
}
|
||||
func (*legacyGitHubOAuthProvider) SetProviderUserID(user *model.User, providerUserID string) {
|
||||
user.GitHubId = providerUserID
|
||||
}
|
||||
func (*legacyGitHubOAuthProvider) GetProviderPrefix() string { return "github_" }
|
||||
func (*legacyGitHubOAuthProvider) ProviderUserIDColumn() string { return "github_id" }
|
||||
|
||||
// legacyGitHubOAuthLogin registers provider for one test and completes an OAuth
|
||||
// login callback with it.
|
||||
func legacyGitHubOAuthLogin(t *testing.T, provider *legacyGitHubOAuthProvider) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
const slug = "github-legacy-login-test"
|
||||
oauth.Register(slug, provider)
|
||||
t.Cleanup(func() { oauth.Unregister(slug) })
|
||||
token, _, err := model.CreateAuthFlow(model.AuthFlowCreate{Purpose: model.AuthFlowPurposeOAuth, Provider: slug, Intent: model.AuthFlowIntentLogin, Payload: `{}`, ExpiresAt: time.Now().Add(time.Minute)})
|
||||
require.NoError(t, err)
|
||||
router := gin.New()
|
||||
router.GET("/api/oauth/:provider", HandleOAuth)
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, httptest.NewRequest(http.MethodGet, "/api/oauth/"+slug+"?state="+token+"&code=provider-code", nil))
|
||||
return response
|
||||
}
|
||||
|
||||
// legacyGitHubBindingAudits returns the account binding audit rows in creation
|
||||
// order together with their encoded parameters.
|
||||
func legacyGitHubBindingAudits(t *testing.T) ([]model.AuditLog, []string) {
|
||||
t.Helper()
|
||||
var audits []model.AuditLog
|
||||
require.NoError(t, model.LOG_DB.Where("action = ?", "user.binding_bind").Order("id").Find(&audits).Error)
|
||||
params := make([]string, 0, len(audits))
|
||||
for _, audit := range audits {
|
||||
require.NotNil(t, audit.Other.Op)
|
||||
encoded, err := common.Marshal(audit.Other.Op.Params)
|
||||
require.NoError(t, err)
|
||||
params = append(params, string(encoded))
|
||||
}
|
||||
return audits, params
|
||||
}
|
||||
|
||||
func TestOAuthLoginLegacyGitHubBindingRequiresAccountEvidence(t *testing.T) {
|
||||
const declined = "This GitHub account cannot be linked to an existing account automatically, please sign in or register another way and then link GitHub in account settings"
|
||||
tests := []struct {
|
||||
name string
|
||||
existingGitHubID string
|
||||
existingDeleted bool
|
||||
withoutEmail bool
|
||||
legacyID string
|
||||
verifiedEmails []string
|
||||
verifiedEmailsErr error
|
||||
registerEnabled bool
|
||||
expectLogin bool
|
||||
expectMigration bool
|
||||
expectNewAccount bool
|
||||
expectEmailCalls int
|
||||
expectAuditParams string
|
||||
}{
|
||||
{name: "numeric legacy value registers a new account", existingGitHubID: "424242", legacyID: "424242", verifiedEmails: []string{"legacy-github@example.com"}, registerEnabled: true, expectLogin: true, expectNewAccount: true},
|
||||
{name: "numeric legacy value with registration disabled", existingGitHubID: "424242", legacyID: "424242", verifiedEmails: []string{"legacy-github@example.com"}},
|
||||
{name: "soft-deleted legacy row registers a new account", existingGitHubID: "octocat-legacy", existingDeleted: true, legacyID: "octocat-legacy", verifiedEmails: []string{"legacy-github@example.com"}, registerEnabled: true, expectLogin: true, expectNewAccount: true},
|
||||
{
|
||||
name: "verified email match migrates", existingGitHubID: "octocat-legacy", legacyID: "octocat-legacy",
|
||||
verifiedEmails: []string{"other@example.com", " Legacy-GitHub@Example.com "}, expectLogin: true, expectMigration: true, expectEmailCalls: 1,
|
||||
// No SMTP server is configured in tests, so the notification attempt is recorded as failed.
|
||||
expectAuditParams: `{"provider":"github","legacy_migration":true,"legacy_id":"octocat-legacy","provider_user_id":"900001","verified_email_matched":true,"success":true,"notification_failed":true}`,
|
||||
},
|
||||
{
|
||||
name: "no matching verified email declines", existingGitHubID: "octocat-legacy", legacyID: "octocat-legacy",
|
||||
verifiedEmails: []string{"other@example.com"}, registerEnabled: true, expectEmailCalls: 1,
|
||||
expectAuditParams: `{"provider":"github","legacy_migration":true,"legacy_id":"octocat-legacy","provider_user_id":"900001","success":false,"reason":"no_matching_evidence"}`,
|
||||
},
|
||||
{
|
||||
name: "verified emails failure declines", existingGitHubID: "octocat-legacy", legacyID: "octocat-legacy",
|
||||
verifiedEmailsErr: errors.New("emails unavailable"), registerEnabled: true, expectEmailCalls: 1,
|
||||
expectAuditParams: `{"provider":"github","legacy_migration":true,"legacy_id":"octocat-legacy","provider_user_id":"900001","success":false,"reason":"verified_emails_unavailable"}`,
|
||||
},
|
||||
{
|
||||
name: "account without email declines before asking the provider", existingGitHubID: "octocat-legacy", withoutEmail: true, legacyID: "octocat-legacy",
|
||||
verifiedEmails: []string{"legacy-github@example.com"}, registerEnabled: true,
|
||||
expectAuditParams: `{"provider":"github","legacy_migration":true,"legacy_id":"octocat-legacy","provider_user_id":"900001","success":false,"reason":"no_matching_evidence"}`,
|
||||
},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
setupSecurityEnrollmentTest(t)
|
||||
previousRegister := common.RegisterEnabled
|
||||
common.RegisterEnabled = test.registerEnabled
|
||||
t.Cleanup(func() { common.RegisterEnabled = previousRegister })
|
||||
existing := &model.User{Username: "legacy-github", Email: "legacy-github@example.com", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", AffCode: "legacy-github", AuthVersion: 1, GitHubId: test.existingGitHubID}
|
||||
if test.withoutEmail {
|
||||
existing.Email = ""
|
||||
}
|
||||
require.NoError(t, model.DB.Create(existing).Error)
|
||||
if test.existingDeleted {
|
||||
require.NoError(t, model.DB.Delete(existing).Error)
|
||||
}
|
||||
provider := &legacyGitHubOAuthProvider{providerUserID: "900001", legacyID: test.legacyID, verifiedEmails: test.verifiedEmails, verifiedEmailsErr: test.verifiedEmailsErr}
|
||||
response := legacyGitHubOAuthLogin(t, provider)
|
||||
var result struct {
|
||||
Success bool `json:"success"`
|
||||
Message string `json:"message"`
|
||||
Data struct {
|
||||
User struct {
|
||||
Id int `json:"id"`
|
||||
} `json:"user"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.Equal(t, test.expectLogin, result.Success, response.Body.String())
|
||||
assert.Equal(t, test.expectEmailCalls, provider.verifiedEmailCalls)
|
||||
|
||||
var reloaded model.User
|
||||
require.NoError(t, model.DB.Unscoped().First(&reloaded, existing.Id).Error)
|
||||
audits, params := legacyGitHubBindingAudits(t)
|
||||
if test.expectAuditParams == "" {
|
||||
assert.Empty(t, audits)
|
||||
} else {
|
||||
require.Len(t, audits, 1)
|
||||
assert.Equal(t, existing.Id, audits[0].UserId)
|
||||
assert.Equal(t, common.RoleCommonUser, audits[0].ActorRole)
|
||||
assert.Equal(t, test.expectMigration, audits[0].Success)
|
||||
assert.JSONEq(t, test.expectAuditParams, params[0])
|
||||
}
|
||||
if test.expectMigration {
|
||||
assert.Equal(t, existing.Id, result.Data.User.Id)
|
||||
assert.Equal(t, "900001", reloaded.GitHubId)
|
||||
assert.NotEmpty(t, response.Header().Values("Set-Cookie"))
|
||||
return
|
||||
}
|
||||
assert.Equal(t, test.existingGitHubID, reloaded.GitHubId, "the existing binding must stay untouched")
|
||||
if test.expectNewAccount {
|
||||
var created model.User
|
||||
require.NoError(t, model.DB.Where("github_id = ?", "900001").First(&created).Error)
|
||||
assert.NotEqual(t, existing.Id, created.Id)
|
||||
assert.Equal(t, created.Id, result.Data.User.Id)
|
||||
return
|
||||
}
|
||||
assert.Empty(t, response.Header().Values("Set-Cookie"))
|
||||
if test.expectAuditParams != "" {
|
||||
assert.Equal(t, declined, result.Message)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOAuthLoginLegacyGitHubBindingMigratesAfterLoginVerification(t *testing.T) {
|
||||
for _, interference := range []string{"none", "pending id claimed by another account", "binding relinked meanwhile"} {
|
||||
t.Run(interference, func(t *testing.T) {
|
||||
setupSecurityEnrollmentTest(t)
|
||||
existing := &model.User{Username: "legacy-github", Email: "legacy-github@example.com", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", AffCode: "legacy-github", AuthVersion: 1, GitHubId: "octocat-legacy"}
|
||||
require.NoError(t, model.DB.Create(existing).Error)
|
||||
factor := &model.TwoFA{UserId: existing.Id, Secret: "JBSWY3DPEHPK3PXP", IsEnabled: true}
|
||||
require.NoError(t, model.DB.Create(factor).Error)
|
||||
provider := &legacyGitHubOAuthProvider{providerUserID: "900001", legacyID: "octocat-legacy", verifiedEmails: []string{"legacy-github@example.com"}}
|
||||
response := legacyGitHubOAuthLogin(t, provider)
|
||||
var challenge struct {
|
||||
Success bool `json:"success"`
|
||||
Data service.LoginChallenge `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &challenge))
|
||||
require.True(t, challenge.Success, response.Body.String())
|
||||
require.True(t, challenge.Data.RequireVerification)
|
||||
assert.Equal(t, []service.VerificationMethodOption{{Method: "2fa", Available: true}}, challenge.Data.Methods)
|
||||
assert.Empty(t, response.Header().Values("Set-Cookie"))
|
||||
assert.Zero(t, provider.verifiedEmailCalls, "a second factor is checked before the provider emails")
|
||||
var reloaded model.User
|
||||
require.NoError(t, model.DB.First(&reloaded, existing.Id).Error)
|
||||
assert.Equal(t, "octocat-legacy", reloaded.GitHubId, "the binding waits for the verification")
|
||||
audits, _ := legacyGitHubBindingAudits(t)
|
||||
assert.Empty(t, audits)
|
||||
|
||||
expectedGitHubID, expectLogin := "900001", true
|
||||
switch interference {
|
||||
case "pending id claimed by another account":
|
||||
other := &model.User{Username: "other-github", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", AffCode: "other-github", AuthVersion: 1, GitHubId: "900001"}
|
||||
require.NoError(t, model.DB.Create(other).Error)
|
||||
expectedGitHubID, expectLogin = "octocat-legacy", false
|
||||
case "binding relinked meanwhile":
|
||||
require.NoError(t, model.DB.Model(existing).Update("github_id", "relinked-legacy").Error)
|
||||
expectedGitHubID = "relinked-legacy"
|
||||
}
|
||||
code, err := totp.GenerateCode(factor.Secret, time.Now())
|
||||
require.NoError(t, err)
|
||||
body, err := common.Marshal(map[string]string{"flow_token": challenge.Data.FlowToken, "code": code})
|
||||
require.NoError(t, err)
|
||||
for attempt := range 2 {
|
||||
response = securityEnrollmentRequest(http.MethodPost, "/api/user/login/verify", string(body), "", service.AuthIdentity{}, VerifyLogin)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.Equal(t, attempt == 0 && expectLogin, result.Success, response.Body.String())
|
||||
if attempt == 0 && !expectLogin {
|
||||
assert.Equal(t, "ACCOUNT_ALREADY_BOUND", result.Code)
|
||||
}
|
||||
if result.Success {
|
||||
assert.NotEmpty(t, response.Header().Values("Set-Cookie"))
|
||||
} else {
|
||||
assert.Empty(t, response.Header().Values("Set-Cookie"))
|
||||
}
|
||||
}
|
||||
require.NoError(t, model.DB.First(&reloaded, existing.Id).Error)
|
||||
assert.Equal(t, expectedGitHubID, reloaded.GitHubId)
|
||||
sessions, err := model.CountActiveUserSessions(existing.Id, time.Now().Unix())
|
||||
require.NoError(t, err)
|
||||
expectedSessions := 0
|
||||
if expectLogin {
|
||||
expectedSessions = 1
|
||||
}
|
||||
assert.EqualValues(t, expectedSessions, sessions, "the consumed challenge cannot be replayed")
|
||||
audits, params := legacyGitHubBindingAudits(t)
|
||||
if interference != "none" {
|
||||
assert.Empty(t, audits)
|
||||
return
|
||||
}
|
||||
require.Len(t, audits, 1)
|
||||
assert.Equal(t, existing.Id, audits[0].UserId)
|
||||
assert.Equal(t, common.RoleCommonUser, audits[0].ActorRole)
|
||||
assert.True(t, audits[0].Success)
|
||||
// No SMTP server is configured in tests, so the notification attempt is recorded as failed.
|
||||
assert.JSONEq(t, `{"provider":"github","legacy_migration":true,"legacy_id":"octocat-legacy","provider_user_id":"900001","verification_method":"2fa","success":true,"notification_failed":true}`, params[0])
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOAuthBindIgnoresLegacyGitHubUsernames(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
ownGitHubID string
|
||||
otherGitHubID string
|
||||
legacyID string
|
||||
expectBound bool
|
||||
}{
|
||||
{name: "another account's legacy username does not block binding", otherGitHubID: "octocat-legacy", legacyID: "octocat-legacy", expectBound: true},
|
||||
{name: "own legacy username is replaced", ownGitHubID: "octocat-legacy", otherGitHubID: "unrelated-legacy", legacyID: "octocat-legacy", expectBound: true},
|
||||
{name: "numeric legacy value does not collide", otherGitHubID: "424242", legacyID: "424242", expectBound: true},
|
||||
{name: "numeric ID held by another account is rejected", otherGitHubID: "900001", legacyID: "octocat-legacy"},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
if test.ownGitHubID != "" {
|
||||
require.NoError(t, model.DB.Model(user).Update("github_id", test.ownGitHubID).Error)
|
||||
}
|
||||
other := &model.User{Username: "other-github", Role: common.RoleCommonUser, Status: common.UserStatusEnabled, Group: "default", AffCode: "other-github", AuthVersion: 1, GitHubId: test.otherGitHubID}
|
||||
require.NoError(t, model.DB.Create(other).Error)
|
||||
const slug = "github-legacy-bind-test"
|
||||
oauth.Register(slug, &legacyGitHubOAuthProvider{providerUserID: "900001", legacyID: test.legacyID})
|
||||
t.Cleanup(func() { oauth.Unregister(slug) })
|
||||
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: service.VerificationScopeAccountBind, Context: []byte(`{"provider":"` + slug + `"}`)}, service.VerificationMethodPassword)
|
||||
started := securityEnrollmentRequest(http.MethodPost, "/api/oauth/state", `{"provider":"`+slug+`","intent":"bind"}`, proof, identity, GenerateOAuthCode)
|
||||
var flow struct {
|
||||
Data struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(started.Body.Bytes(), &flow))
|
||||
require.NotEmpty(t, flow.Data.FlowToken, started.Body.String())
|
||||
response := securityEnrollmentRequest(http.MethodGet, "/api/oauth/"+slug+"?state="+flow.Data.FlowToken+"&code=provider-code", "", "", identity, func(c *gin.Context) {
|
||||
c.Params = gin.Params{{Key: "provider", Value: slug}}
|
||||
HandleOAuth(c)
|
||||
})
|
||||
var result struct {
|
||||
Success bool `json:"success"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.Equal(t, test.expectBound, result.Success, response.Body.String())
|
||||
|
||||
var bound, untouched model.User
|
||||
require.NoError(t, model.DB.First(&bound, user.Id).Error)
|
||||
require.NoError(t, model.DB.First(&untouched, other.Id).Error)
|
||||
assert.Equal(t, test.otherGitHubID, untouched.GitHubId)
|
||||
if test.expectBound {
|
||||
assert.Equal(t, "900001", bound.GitHubId)
|
||||
return
|
||||
}
|
||||
assert.Equal(t, test.ownGitHubID, bound.GitHubId)
|
||||
assert.Equal(t, "This GitHub account has already been bound", result.Message)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -146,7 +146,7 @@ func TestSessionLimitDoesNotRecordRejectedLoginAsSuccessful(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/user/login", nil)
|
||||
setupLogin(user, c)
|
||||
setupLogin(user, nil, c)
|
||||
|
||||
assert.Equal(t, http.StatusConflict, recorder.Code)
|
||||
var stored model.User
|
||||
|
||||
@@ -11,10 +11,8 @@ import (
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/setting/config"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestUpdateOptionRejectsInvalidTaskBillingExpressions(t *testing.T) {
|
||||
@@ -103,42 +101,16 @@ func TestUpdateOptionRejectsUsageExpressionWithoutTaskPlugin(t *testing.T) {
|
||||
assert.Contains(t, recorder.Body.String(), "no task plugin usage schema")
|
||||
}
|
||||
|
||||
func setupBillingAliasOptionDB(t *testing.T) {
|
||||
t.Helper()
|
||||
previousDB := model.DB
|
||||
previousLogDB := model.LOG_DB
|
||||
previousType := common.MainDatabaseType()
|
||||
previousCache := common.MemoryCacheEnabled
|
||||
previousMap := common.OptionMap
|
||||
previousRedis := common.RedisEnabled
|
||||
database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Option{}, &model.Log{}, &model.AuditLog{}, &model.User{}))
|
||||
model.DB = database
|
||||
model.LOG_DB = database
|
||||
common.SetMainDatabaseType(common.DatabaseTypeSQLite)
|
||||
common.MemoryCacheEnabled = false
|
||||
common.RedisEnabled = false
|
||||
common.OptionMap = map[string]string{}
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
model.LOG_DB = previousLogDB
|
||||
common.SetMainDatabaseType(previousType)
|
||||
common.MemoryCacheEnabled = previousCache
|
||||
common.OptionMap = previousMap
|
||||
common.RedisEnabled = previousRedis
|
||||
model.InitChannelCache()
|
||||
})
|
||||
}
|
||||
|
||||
func TestUpdateOptionAliasBillingExprUsesPluginSchema(t *testing.T) {
|
||||
setupBillingAliasOptionDB(t)
|
||||
database := modelManagementDB(t, "sqlite", "")
|
||||
require.NoError(t, database.AutoMigrate(&model.Log{}))
|
||||
const pluginKey = "billing-alias-probe"
|
||||
source := `
|
||||
export const meta = {
|
||||
apiVersion: 1, key: "billing-alias-probe", name: "Billing Alias Probe", version: "1.0.0", author: {name: "Test"},
|
||||
models: ["declared-model"], fetchMode: "per_task",
|
||||
usageSchema: {seconds: {type: "number", unit: "second"}}
|
||||
usageSchema: {seconds: {type: "number", unit: "second"}, image_count: {type: "number", unit: "count"}},
|
||||
usageProfiles: [{models: ["declared-model"], schema: {seconds: {type: "number", unit: "second"}}}]
|
||||
};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {}; }
|
||||
@@ -191,10 +163,14 @@ export function parseTaskResult() { return {}; }
|
||||
assert.Equal(t, http.StatusOK, accepted.Code)
|
||||
assert.Contains(t, accepted.Body.String(), `"success":true`)
|
||||
|
||||
rejectedKey := putExpr("alias-model", `u("clips")`)
|
||||
rejectedKey := putExpr("alias-model", `u("image_count")`)
|
||||
assert.Equal(t, http.StatusOK, rejectedKey.Code)
|
||||
assert.Contains(t, rejectedKey.Body.String(), `"success":false`)
|
||||
assert.Contains(t, rejectedKey.Body.String(), `usage key \"clips\" is not declared`)
|
||||
assert.Contains(t, rejectedKey.Body.String(), `usage key \"image_count\" is not declared`)
|
||||
|
||||
rejectedDeclared := putExpr("declared-model", `u("image_count")`)
|
||||
assert.Contains(t, rejectedDeclared.Body.String(), `"success":false`)
|
||||
assert.Contains(t, rejectedDeclared.Body.String(), `usage key \"image_count\" is not declared`)
|
||||
|
||||
unresolvable := putExpr("unknown-alias-model", `u("seconds")`)
|
||||
assert.Equal(t, http.StatusOK, unresolvable.Code)
|
||||
|
||||
@@ -275,6 +275,67 @@ func updateChannelSiliconFlowBalance(channel *model.Channel) (float64, error) {
|
||||
return balance, nil
|
||||
}
|
||||
|
||||
func getDeepSeekBalanceUSD(response DeepSeekUsageResponse, usdExchangeRate float64) (float64, error) {
|
||||
var usdBalance *string
|
||||
var cnyBalance *string
|
||||
|
||||
for i := range response.BalanceInfos {
|
||||
balanceInfo := &response.BalanceInfos[i]
|
||||
switch balanceInfo.Currency {
|
||||
case "USD":
|
||||
if usdBalance == nil {
|
||||
usdBalance = &balanceInfo.TotalBalance
|
||||
}
|
||||
case "CNY":
|
||||
if cnyBalance == nil {
|
||||
cnyBalance = &balanceInfo.TotalBalance
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if usdBalance != nil {
|
||||
balance, err := strconv.ParseFloat(*usdBalance, 64)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if math.IsNaN(balance) || math.IsInf(balance, 0) {
|
||||
return 0, errors.New("USD balance must be finite")
|
||||
}
|
||||
if balance < 0 {
|
||||
return 0, errors.New("USD balance must be non-negative")
|
||||
}
|
||||
return balance, nil
|
||||
}
|
||||
if cnyBalance == nil {
|
||||
return 0, errors.New("currency USD or CNY not found")
|
||||
}
|
||||
if math.IsNaN(usdExchangeRate) || math.IsInf(usdExchangeRate, 0) {
|
||||
return 0, errors.New("USD exchange rate must be finite")
|
||||
}
|
||||
if usdExchangeRate <= 0 {
|
||||
return 0, errors.New("USD exchange rate must be greater than zero")
|
||||
}
|
||||
|
||||
balanceCNY, err := strconv.ParseFloat(*cnyBalance, 64)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if math.IsNaN(balanceCNY) || math.IsInf(balanceCNY, 0) {
|
||||
return 0, errors.New("CNY balance must be finite")
|
||||
}
|
||||
if balanceCNY < 0 {
|
||||
return 0, errors.New("CNY balance must be non-negative")
|
||||
}
|
||||
balanceUSD := decimal.NewFromFloat(balanceCNY).Div(decimal.NewFromFloat(usdExchangeRate)).InexactFloat64()
|
||||
if math.IsNaN(balanceUSD) || math.IsInf(balanceUSD, 0) {
|
||||
return 0, errors.New("converted USD balance must be finite")
|
||||
}
|
||||
if balanceUSD < 0 {
|
||||
return 0, errors.New("converted USD balance must be non-negative")
|
||||
}
|
||||
return balanceUSD, nil
|
||||
}
|
||||
|
||||
func updateChannelDeepSeekBalance(channel *model.Channel) (float64, error) {
|
||||
url := "https://api.deepseek.com/user/balance"
|
||||
body, err := GetResponseBody("GET", url, channel, GetAuthHeader(channel.Key))
|
||||
@@ -286,17 +347,7 @@ func updateChannelDeepSeekBalance(channel *model.Channel) (float64, error) {
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
index := -1
|
||||
for i, balanceInfo := range response.BalanceInfos {
|
||||
if balanceInfo.Currency == "CNY" {
|
||||
index = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if index == -1 {
|
||||
return 0, errors.New("currency CNY not found")
|
||||
}
|
||||
balance, err := strconv.ParseFloat(response.BalanceInfos[index].TotalBalance, 64)
|
||||
balance, err := getDeepSeekBalanceUSD(response, operation_setting.USDExchangeRate)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
+141
-5
@@ -73,6 +73,68 @@ func clearChannelInfo(channel *model.Channel) {
|
||||
}
|
||||
}
|
||||
|
||||
func channelIDsFromChannels(channels []*model.Channel) []int {
|
||||
ids := make([]int, 0, len(channels))
|
||||
for _, channel := range channels {
|
||||
if channel != nil && channel.Id > 0 {
|
||||
ids = append(ids, channel.Id)
|
||||
}
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
func closeActiveChannelWebSockets(channelIDs []int) {
|
||||
service.CloseActiveWebSocketsForChannels(channelIDs, service.ChannelDisabledCloseReason)
|
||||
}
|
||||
|
||||
func hasEnabledMultiKey(channel *model.Channel) bool {
|
||||
if channel == nil || !channel.ChannelInfo.IsMultiKey {
|
||||
return true
|
||||
}
|
||||
keys := channel.GetKeys()
|
||||
if len(keys) == 0 {
|
||||
return false
|
||||
}
|
||||
for i := range keys {
|
||||
if channel.ChannelInfo.MultiKeyStatusList == nil {
|
||||
return true
|
||||
}
|
||||
if status, ok := channel.ChannelInfo.MultiKeyStatusList[i]; !ok || status == common.ChannelStatusEnabled {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func disableMultiKeyChannelIfUnavailable(channel *model.Channel) bool {
|
||||
if channel == nil || !channel.ChannelInfo.IsMultiKey || hasEnabledMultiKey(channel) {
|
||||
return false
|
||||
}
|
||||
if channel.Status != common.ChannelStatusEnabled {
|
||||
return true
|
||||
}
|
||||
channel.Status = common.ChannelStatusManuallyDisabled
|
||||
info := channel.GetOtherInfo()
|
||||
info["status_reason"] = model.ChannelStatusReasonAllKeysDisabled
|
||||
info["status_time"] = common.GetTimestamp()
|
||||
channel.SetOtherInfo(info)
|
||||
return true
|
||||
}
|
||||
|
||||
func restoreMultiKeyChannelIfAvailable(channel *model.Channel) {
|
||||
if channel.Status != common.ChannelStatusManuallyDisabled || !hasEnabledMultiKey(channel) {
|
||||
return
|
||||
}
|
||||
info := channel.GetOtherInfo()
|
||||
if info["status_reason"] != model.ChannelStatusReasonAllKeysDisabled {
|
||||
return
|
||||
}
|
||||
channel.Status = common.ChannelStatusEnabled
|
||||
info["status_reason"] = ""
|
||||
info["status_time"] = common.GetTimestamp()
|
||||
channel.SetOtherInfo(info)
|
||||
}
|
||||
|
||||
func applyChannelStatusFilter(query *gorm.DB, statusFilter int) *gorm.DB {
|
||||
if statusFilter == common.ChannelStatusEnabled {
|
||||
return query.Where("status = ?", common.ChannelStatusEnabled)
|
||||
@@ -94,11 +156,38 @@ func buildChannelListQuery(group string, statusFilter int, typeFilter int) *gorm
|
||||
}
|
||||
|
||||
func GetChannelOps(c *gin.Context) {
|
||||
snapshot := model.CurrentRequestPolicy()
|
||||
automaticDisable, source := snapshot.AutoDisable, "global"
|
||||
if value, present := c.GetQuery("auto_ban"); present {
|
||||
autoBan, err := strconv.ParseBool(value)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
automaticDisable = snapshot.AutoDisable && autoBan
|
||||
if snapshot.AutoDisable {
|
||||
source = "global_and_channel"
|
||||
if !autoBan {
|
||||
source = "channel"
|
||||
}
|
||||
}
|
||||
}
|
||||
common.ApiSuccess(c, gin.H{
|
||||
"retry_times": common.RetryTimes,
|
||||
"retry_times": snapshot.RetryTimes,
|
||||
"request_policy": gin.H{"automatic_disable": automaticDisable, "source": source},
|
||||
})
|
||||
}
|
||||
|
||||
func GetChannelDefaultBaseURLs(c *gin.Context) {
|
||||
baseURLs := make(map[int]string)
|
||||
for channelType, baseURL := range constant.ChannelBaseURLs {
|
||||
if baseURL != "" {
|
||||
baseURLs[channelType] = baseURL
|
||||
}
|
||||
}
|
||||
common.ApiSuccess(c, baseURLs)
|
||||
}
|
||||
|
||||
func GetAllChannels(c *gin.Context) {
|
||||
pageInfo := common.GetPageQuery(c)
|
||||
channelData := make([]*model.Channel, 0)
|
||||
@@ -485,6 +574,12 @@ func validateChannel(channel *model.Channel, isAdd bool) error {
|
||||
if channel.Type == constant.ChannelTypeNewAPI && strings.TrimSpace(channel.GetBaseURL()) == "" {
|
||||
return fmt.Errorf("New API channel base URL cannot be empty")
|
||||
}
|
||||
if channel.Type == constant.ChannelTypeVLLM && strings.TrimSpace(channel.GetBaseURL()) == "" {
|
||||
return fmt.Errorf("vLLM channel base URL cannot be empty")
|
||||
}
|
||||
if channel.Type == constant.ChannelTypeSGLang && strings.TrimSpace(channel.GetBaseURL()) == "" {
|
||||
return fmt.Errorf("SGLang channel base URL cannot be empty")
|
||||
}
|
||||
|
||||
// 如果是添加操作,检查 channel 和 key 是否为空
|
||||
if isAdd {
|
||||
@@ -757,6 +852,7 @@ func DeleteChannel(c *gin.Context) {
|
||||
"id": id,
|
||||
"name": channelName,
|
||||
})
|
||||
closeActiveChannelWebSockets([]int{id})
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
@@ -765,6 +861,13 @@ func DeleteChannel(c *gin.Context) {
|
||||
}
|
||||
|
||||
func DeleteDisabledChannel(c *gin.Context) {
|
||||
var ids []int
|
||||
if err := model.DB.Model(&model.Channel{}).
|
||||
Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled).
|
||||
Pluck("id", &ids).Error; err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
rows, err := model.DeleteDisabledChannel()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
@@ -777,6 +880,7 @@ func DeleteDisabledChannel(c *gin.Context) {
|
||||
recordManageAudit(c, "channel.delete_disabled", map[string]any{
|
||||
"count": rows,
|
||||
})
|
||||
closeActiveChannelWebSockets(ids)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
@@ -807,6 +911,12 @@ func DisableTagChannels(c *gin.Context) {
|
||||
})
|
||||
return
|
||||
}
|
||||
channels, err := model.GetChannelsByTag(channelTag.Tag, false, false)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
ids := channelIDsFromChannels(channels)
|
||||
err = model.DisableChannelByTag(channelTag.Tag)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
@@ -816,6 +926,7 @@ func DisableTagChannels(c *gin.Context) {
|
||||
recordManageAudit(c, "channel.tag_disable", map[string]any{
|
||||
"tag": channelTag.Tag,
|
||||
})
|
||||
closeActiveChannelWebSockets(ids)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
@@ -936,6 +1047,7 @@ func DeleteChannelBatch(c *gin.Context) {
|
||||
recordManageAudit(c, "channel.delete_batch", map[string]any{
|
||||
"count": deletedCount,
|
||||
})
|
||||
closeActiveChannelWebSockets(channelBatch.Ids)
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
@@ -1170,6 +1282,9 @@ func UpdateChannelStatus(c *gin.Context) {
|
||||
changed := model.UpdateChannelStatus(id, "", req.Status, "manual operation")
|
||||
if changed {
|
||||
model.InitChannelCache()
|
||||
if req.Status != common.ChannelStatusEnabled {
|
||||
closeActiveChannelWebSockets([]int{id})
|
||||
}
|
||||
}
|
||||
recordManageAudit(c, "channel.status_update", map[string]any{
|
||||
"id": id,
|
||||
@@ -1190,14 +1305,21 @@ func BatchUpdateChannelStatus(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
changedCount := 0
|
||||
var disabledIDs []int
|
||||
for _, id := range req.Ids {
|
||||
if model.UpdateChannelStatus(id, "", req.Status, "manual batch operation") {
|
||||
changedCount++
|
||||
if req.Status != common.ChannelStatusEnabled {
|
||||
disabledIDs = append(disabledIDs, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
if changedCount > 0 {
|
||||
model.InitChannelCache()
|
||||
}
|
||||
if len(disabledIDs) > 0 {
|
||||
closeActiveChannelWebSockets(disabledIDs)
|
||||
}
|
||||
recordManageAudit(c, "channel.status_update_batch", map[string]any{
|
||||
"count": changedCount,
|
||||
"total": len(req.Ids),
|
||||
@@ -1714,13 +1836,16 @@ func ManageMultiKeys(c *gin.Context) {
|
||||
|
||||
channel.ChannelInfo.MultiKeyStatusList[keyIndex] = 2 // disabled
|
||||
|
||||
shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel)
|
||||
err = channel.Update()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
model.InitChannelCache()
|
||||
if shouldCloseWebSocket {
|
||||
closeActiveChannelWebSockets([]int{channel.Id})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "密钥已禁用",
|
||||
@@ -1755,6 +1880,7 @@ func ManageMultiKeys(c *gin.Context) {
|
||||
if channel.ChannelInfo.MultiKeyDisabledReason != nil {
|
||||
delete(channel.ChannelInfo.MultiKeyDisabledReason, keyIndex)
|
||||
}
|
||||
restoreMultiKeyChannelIfAvailable(channel)
|
||||
|
||||
err = channel.Update()
|
||||
if err != nil {
|
||||
@@ -1779,6 +1905,7 @@ func ManageMultiKeys(c *gin.Context) {
|
||||
channel.ChannelInfo.MultiKeyStatusList = make(map[int]int)
|
||||
channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64)
|
||||
channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string)
|
||||
restoreMultiKeyChannelIfAvailable(channel)
|
||||
|
||||
err = channel.Update()
|
||||
if err != nil {
|
||||
@@ -1827,13 +1954,16 @@ func ManageMultiKeys(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel)
|
||||
err = channel.Update()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
model.InitChannelCache()
|
||||
if shouldCloseWebSocket {
|
||||
closeActiveChannelWebSockets([]int{channel.Id})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": fmt.Sprintf("已禁用 %d 个密钥", disabledCount),
|
||||
@@ -1907,13 +2037,16 @@ func ManageMultiKeys(c *gin.Context) {
|
||||
channel.ChannelInfo.MultiKeyDisabledTime = newDisabledTime
|
||||
channel.ChannelInfo.MultiKeyDisabledReason = newDisabledReason
|
||||
|
||||
shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel)
|
||||
err = channel.Update()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
model.InitChannelCache()
|
||||
if shouldCloseWebSocket {
|
||||
closeActiveChannelWebSockets([]int{channel.Id})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "密钥已删除",
|
||||
@@ -1975,13 +2108,16 @@ func ManageMultiKeys(c *gin.Context) {
|
||||
channel.ChannelInfo.MultiKeyDisabledTime = newDisabledTime
|
||||
channel.ChannelInfo.MultiKeyDisabledReason = newDisabledReason
|
||||
|
||||
shouldCloseWebSocket := disableMultiKeyChannelIfUnavailable(channel)
|
||||
err = channel.Update()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
model.InitChannelCache()
|
||||
if shouldCloseWebSocket {
|
||||
closeActiveChannelWebSockets([]int{channel.Id})
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": fmt.Sprintf("已删除 %d 个自动禁用的密钥", deletedCount),
|
||||
|
||||
@@ -0,0 +1,116 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGetDeepSeekBalanceUSD(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
responseJSON string
|
||||
usdExchangeRate float64
|
||||
want float64
|
||||
wantErrContains string
|
||||
}{
|
||||
{
|
||||
name: "prefers USD when USD precedes CNY",
|
||||
responseJSON: `{"balance_infos":[{"currency":"USD","total_balance":"12.50"},{"currency":"CNY","total_balance":"73.00"}]}`,
|
||||
usdExchangeRate: 7.3,
|
||||
want: 12.5,
|
||||
},
|
||||
{
|
||||
name: "prefers USD when CNY precedes USD",
|
||||
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"73.00"},{"currency":"USD","total_balance":"12.50"}]}`,
|
||||
usdExchangeRate: 7.3,
|
||||
want: 12.5,
|
||||
},
|
||||
{
|
||||
name: "converts CNY when USD is absent",
|
||||
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"73.00"}]}`,
|
||||
usdExchangeRate: 7.3,
|
||||
want: 10,
|
||||
},
|
||||
{
|
||||
name: "returns error when USD and CNY are absent",
|
||||
responseJSON: `{"balance_infos":[{"currency":"EUR","total_balance":"10.00"}]}`,
|
||||
usdExchangeRate: 7.3,
|
||||
wantErrContains: "currency USD or CNY not found",
|
||||
},
|
||||
{
|
||||
name: "returns USD parse error instead of falling back to CNY",
|
||||
responseJSON: `{"balance_infos":[{"currency":"USD","total_balance":"invalid"},{"currency":"CNY","total_balance":"73.00"}]}`,
|
||||
usdExchangeRate: 7.3,
|
||||
wantErrContains: "invalid syntax",
|
||||
},
|
||||
{
|
||||
name: "rejects NaN USD balance",
|
||||
responseJSON: `{"balance_infos":[{"currency":"USD","total_balance":"NaN"}]}`,
|
||||
usdExchangeRate: 7.3,
|
||||
wantErrContains: "USD balance must be finite",
|
||||
},
|
||||
{
|
||||
name: "rejects negative USD balance",
|
||||
responseJSON: `{"balance_infos":[{"currency":"USD","total_balance":"-1.00"}]}`,
|
||||
usdExchangeRate: 7.3,
|
||||
wantErrContains: "USD balance must be non-negative",
|
||||
},
|
||||
{
|
||||
name: "rejects positive infinity CNY balance",
|
||||
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"+Inf"}]}`,
|
||||
usdExchangeRate: 7.3,
|
||||
wantErrContains: "CNY balance must be finite",
|
||||
},
|
||||
{
|
||||
name: "rejects negative CNY balance",
|
||||
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"-7.30"}]}`,
|
||||
usdExchangeRate: 7.3,
|
||||
wantErrContains: "CNY balance must be non-negative",
|
||||
},
|
||||
{
|
||||
name: "returns error for non-positive CNY exchange rate",
|
||||
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"73.00"}]}`,
|
||||
usdExchangeRate: 0,
|
||||
wantErrContains: "USD exchange rate must be greater than zero",
|
||||
},
|
||||
{
|
||||
name: "rejects NaN CNY exchange rate",
|
||||
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"73.00"}]}`,
|
||||
usdExchangeRate: math.NaN(),
|
||||
wantErrContains: "USD exchange rate must be finite",
|
||||
},
|
||||
{
|
||||
name: "rejects positive infinity CNY exchange rate",
|
||||
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"73.00"}]}`,
|
||||
usdExchangeRate: math.Inf(1),
|
||||
wantErrContains: "USD exchange rate must be finite",
|
||||
},
|
||||
{
|
||||
name: "rejects CNY conversion overflow",
|
||||
responseJSON: `{"balance_infos":[{"currency":"CNY","total_balance":"1.7976931348623157e+308"}]}`,
|
||||
usdExchangeRate: math.SmallestNonzeroFloat64,
|
||||
wantErrContains: "converted USD balance must be finite",
|
||||
},
|
||||
}
|
||||
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
var response DeepSeekUsageResponse
|
||||
require.NoError(t, common.Unmarshal([]byte(test.responseJSON), &response))
|
||||
|
||||
balance, err := getDeepSeekBalanceUSD(response, test.usdExchangeRate)
|
||||
if test.wantErrContains != "" {
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), test.wantErrContains)
|
||||
return
|
||||
}
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.InDelta(t, test.want, balance, 1e-12)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,242 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/prometheus/common/expfmt"
|
||||
)
|
||||
|
||||
type inferenceEndpointStatus struct {
|
||||
Status int `json:"status"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type inferenceModel struct {
|
||||
ID string `json:"id"`
|
||||
Root string `json:"root"`
|
||||
MaxModelLen *int64 `json:"max_model_len"`
|
||||
}
|
||||
|
||||
type inferenceMetric struct {
|
||||
Name string `json:"name"`
|
||||
Labels map[string]string `json:"labels"`
|
||||
Value float64 `json:"value"`
|
||||
}
|
||||
|
||||
type inferenceStatus struct {
|
||||
SampledAt int64 `json:"sampled_at"`
|
||||
Endpoints map[string]inferenceEndpointStatus `json:"endpoints"`
|
||||
Version string `json:"version"`
|
||||
Models []inferenceModel `json:"models"`
|
||||
Metrics []inferenceMetric `json:"metrics"`
|
||||
RawMetrics string `json:"raw_metrics"`
|
||||
}
|
||||
|
||||
func GetVLLMChannelStatus(c *gin.Context) {
|
||||
getInferenceChannelStatus(c, constant.ChannelTypeVLLM)
|
||||
}
|
||||
|
||||
func GetSGLangChannelStatus(c *gin.Context) {
|
||||
getInferenceChannelStatus(c, constant.ChannelTypeSGLang)
|
||||
}
|
||||
|
||||
func getInferenceChannelStatus(c *gin.Context, channelType int) {
|
||||
c.Header("Cache-Control", "no-store")
|
||||
id, err := strconv.Atoi(c.Param("id"))
|
||||
if err != nil || id <= 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "Invalid channel id"})
|
||||
return
|
||||
}
|
||||
channel, err := model.GetChannelById(id, true)
|
||||
if err != nil || channel == nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"success": false, "message": "Channel not found"})
|
||||
return
|
||||
}
|
||||
if channel.Type != channelType {
|
||||
common.ApiError(c, errors.New("Channel type does not match the status endpoint"))
|
||||
return
|
||||
}
|
||||
status, err := fetchInferenceStatus(c.Request.Context(), channel)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, status)
|
||||
}
|
||||
|
||||
// fetchInferenceStatus reads only fixed paths on an operator-configured channel.
|
||||
// Each endpoint can fail independently; upstream errors never echo credentials.
|
||||
func fetchInferenceStatus(ctx context.Context, channel *model.Channel) (*inferenceStatus, error) {
|
||||
if channel.Type != constant.ChannelTypeVLLM && channel.Type != constant.ChannelTypeSGLang {
|
||||
return nil, errors.New("This operation is only supported for vLLM or SGLang channels")
|
||||
}
|
||||
baseURL := strings.TrimRight(strings.TrimSpace(channel.GetBaseURL()), "/")
|
||||
parsedURL, err := url.Parse(baseURL)
|
||||
if err != nil || parsedURL.Host == "" || parsedURL.User != nil || parsedURL.RawQuery != "" || parsedURL.Fragment != "" || (parsedURL.Scheme != "http" && parsedURL.Scheme != "https") {
|
||||
return nil, errors.New("Invalid inference server address")
|
||||
}
|
||||
key, _, keyErr := channel.GetNextEnabledKey()
|
||||
if keyErr != nil {
|
||||
return nil, errors.New("No enabled channel key")
|
||||
}
|
||||
key = strings.TrimSpace(key)
|
||||
headers, err := buildFetchModelsHeaders(channel, key)
|
||||
if err != nil {
|
||||
return nil, errors.New("Invalid channel header override")
|
||||
}
|
||||
if (key == "" || key == "EMPTY") && headers.Get("Authorization") == "Bearer "+key {
|
||||
headers.Del("Authorization")
|
||||
}
|
||||
settings := channel.GetSetting()
|
||||
client, err := service.GetHttpClientWithProxySettings(settings.Proxy, settings)
|
||||
if err != nil {
|
||||
return nil, errors.New("Invalid channel proxy")
|
||||
}
|
||||
// Do not mutate the shared client or forward channel credentials on redirects.
|
||||
statusClient := *client
|
||||
statusClient.CheckRedirect = func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }
|
||||
ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
versionPath, metricPrefix := "/version", "vllm:"
|
||||
if channel.Type == constant.ChannelTypeSGLang {
|
||||
versionPath, metricPrefix = "/server_info", "sglang:"
|
||||
}
|
||||
paths := []string{"/health", versionPath, "/v1/models", "/metrics"}
|
||||
results := make([]struct {
|
||||
inferenceEndpointStatus
|
||||
body []byte
|
||||
at int64
|
||||
}, len(paths))
|
||||
var wg sync.WaitGroup
|
||||
for i, path := range paths {
|
||||
wg.Go(func() {
|
||||
result := &results[i]
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, baseURL+path, nil)
|
||||
if err != nil {
|
||||
result.Error = "request_failed"
|
||||
return
|
||||
}
|
||||
req.Header = headers.Clone()
|
||||
req.Host = headers.Get("Host")
|
||||
resp, err := statusClient.Do(req)
|
||||
if err != nil {
|
||||
result.Error = "request_failed"
|
||||
return
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
result.Status = resp.StatusCode
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
result.Error = "http_error"
|
||||
return
|
||||
}
|
||||
const maxResponseBytes = 4 << 20
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, maxResponseBytes+1))
|
||||
if err != nil {
|
||||
result.Error = "request_failed"
|
||||
return
|
||||
}
|
||||
if len(body) > maxResponseBytes {
|
||||
result.Error = "response_too_large"
|
||||
return
|
||||
}
|
||||
result.body = body
|
||||
result.at = time.Now().UnixMilli()
|
||||
})
|
||||
}
|
||||
wg.Wait()
|
||||
if errors.Is(ctx.Err(), context.Canceled) {
|
||||
return nil, context.Canceled
|
||||
}
|
||||
status := &inferenceStatus{
|
||||
SampledAt: time.Now().UnixMilli(),
|
||||
Endpoints: make(map[string]inferenceEndpointStatus, len(paths)),
|
||||
Models: []inferenceModel{},
|
||||
Metrics: []inferenceMetric{},
|
||||
}
|
||||
for i, path := range paths {
|
||||
result := &results[i]
|
||||
if result.Error == "" {
|
||||
switch path {
|
||||
case versionPath:
|
||||
// /server_info includes the launch configuration. Decode only the
|
||||
// version so API keys and other server settings never reach the browser.
|
||||
var version struct {
|
||||
Version string `json:"version"`
|
||||
}
|
||||
if err := common.Unmarshal(result.body, &version); err != nil || version.Version == "" {
|
||||
result.Error = "invalid_response"
|
||||
} else {
|
||||
status.Version = version.Version
|
||||
}
|
||||
case "/v1/models":
|
||||
var models struct {
|
||||
Data []inferenceModel `json:"data"`
|
||||
}
|
||||
if err := common.Unmarshal(result.body, &models); err != nil || models.Data == nil {
|
||||
result.Error = "invalid_response"
|
||||
} else {
|
||||
status.Models = models.Data
|
||||
}
|
||||
case "/metrics":
|
||||
metrics, err := parseInferenceMetrics(string(result.body), metricPrefix)
|
||||
if err != nil {
|
||||
result.Error = "invalid_response"
|
||||
} else {
|
||||
status.Metrics = metrics
|
||||
status.RawMetrics = string(result.body)
|
||||
status.SampledAt = result.at
|
||||
}
|
||||
}
|
||||
}
|
||||
status.Endpoints[path] = result.inferenceEndpointStatus
|
||||
}
|
||||
return status, nil
|
||||
}
|
||||
|
||||
func parseInferenceMetrics(raw string, prefix string) ([]inferenceMetric, error) {
|
||||
var parser expfmt.TextParser
|
||||
families, err := parser.TextToMetricFamilies(strings.NewReader(raw))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
metrics := []inferenceMetric{}
|
||||
for name, family := range families {
|
||||
if !strings.HasPrefix(name, prefix) && name != "process_start_time_seconds" {
|
||||
continue
|
||||
}
|
||||
samples, err := expfmt.ExtractSamples(&expfmt.DecodeOptions{}, family)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for _, sample := range samples {
|
||||
name := string(sample.Metric["__name__"])
|
||||
value := float64(sample.Value)
|
||||
if strings.HasSuffix(name, "_bucket") || math.IsNaN(value) || math.IsInf(value, 0) || value < 0 {
|
||||
continue
|
||||
}
|
||||
labels := make(map[string]string, len(sample.Metric)-1)
|
||||
for label, value := range sample.Metric {
|
||||
if label != "__name__" {
|
||||
labels[string(label)] = string(value)
|
||||
}
|
||||
}
|
||||
metrics = append(metrics, inferenceMetric{Name: name, Labels: labels, Value: value})
|
||||
}
|
||||
}
|
||||
return metrics, nil
|
||||
}
|
||||
@@ -0,0 +1,196 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestVLLMStatusPartialFailureAndCredentials(t *testing.T) {
|
||||
service.InitHttpClient()
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Equal(t, http.MethodGet, r.Method)
|
||||
assert.Equal(t, "Bearer test-vllm-key", r.Header.Get("Authorization"))
|
||||
assert.Equal(t, "status-panel", r.Header.Get("X-Source"))
|
||||
switch r.URL.Path {
|
||||
case "/prefix/health":
|
||||
w.WriteHeader(http.StatusOK)
|
||||
case "/prefix/version":
|
||||
_, _ = w.Write([]byte(`{"version":"0.25.2"}`))
|
||||
case "/prefix/v1/models":
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"served-model","root":"source/model","max_model_len":1048576}]}`))
|
||||
case "/prefix/metrics":
|
||||
http.Error(w, "test-vllm-key should not reach the browser", http.StatusForbidden)
|
||||
default:
|
||||
t.Errorf("unexpected path: %s", r.URL.Path)
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}
|
||||
}))
|
||||
defer upstream.Close()
|
||||
base := upstream.URL + "/prefix/"
|
||||
override := `{"X-Source":"status-panel"}`
|
||||
channel := &model.Channel{Type: constant.ChannelTypeVLLM, BaseURL: &base, Key: "test-vllm-key", HeaderOverride: &override}
|
||||
status, err := fetchInferenceStatus(context.Background(), channel)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 200, status.Endpoints["/health"].Status)
|
||||
assert.Equal(t, "0.25.2", status.Version)
|
||||
require.Len(t, status.Models, 1)
|
||||
assert.Equal(t, "source/model", status.Models[0].Root)
|
||||
require.NotNil(t, status.Models[0].MaxModelLen)
|
||||
assert.EqualValues(t, 1048576, *status.Models[0].MaxModelLen)
|
||||
assert.Equal(t, inferenceEndpointStatus{Status: 403, Error: "http_error"}, status.Endpoints["/metrics"])
|
||||
assert.Empty(t, status.RawMetrics)
|
||||
assert.Empty(t, status.Metrics)
|
||||
body, err := common.Marshal(status)
|
||||
require.NoError(t, err)
|
||||
assert.NotContains(t, string(body), channel.Key)
|
||||
}
|
||||
|
||||
func TestVLLMStatusRejectsInvalidTargetsAndRedirects(t *testing.T) {
|
||||
service.InitHttpClient()
|
||||
for _, base := range []string{"", "file:///etc/passwd", "http://user:secret@localhost", "http://localhost?key=secret"} {
|
||||
_, err := fetchInferenceStatus(context.Background(), &model.Channel{Type: constant.ChannelTypeVLLM, BaseURL: &base})
|
||||
require.Error(t, err)
|
||||
assert.NotContains(t, err.Error(), "secret")
|
||||
}
|
||||
_, err := fetchInferenceStatus(context.Background(), &model.Channel{Type: constant.ChannelTypeOllama})
|
||||
require.Error(t, err)
|
||||
target := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
t.Error("redirect must not forward channel credentials")
|
||||
}))
|
||||
defer target.Close()
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Empty(t, r.Header.Get("Authorization"))
|
||||
http.Redirect(w, r, target.URL, http.StatusFound)
|
||||
}))
|
||||
defer upstream.Close()
|
||||
status, err := fetchInferenceStatus(context.Background(), &model.Channel{Type: constant.ChannelTypeVLLM, BaseURL: &upstream.URL, Key: "EMPTY"})
|
||||
require.NoError(t, err)
|
||||
for _, endpoint := range status.Endpoints {
|
||||
assert.Equal(t, inferenceEndpointStatus{Status: 302, Error: "http_error"}, endpoint)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVLLMStatusInvalidAndOversizedResponses(t *testing.T) {
|
||||
service.InitHttpClient()
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/health":
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
case "/version":
|
||||
_, _ = w.Write([]byte(`{"version":false}`))
|
||||
case "/v1/models":
|
||||
_, _ = w.Write([]byte(`{"data":[]}`))
|
||||
case "/metrics":
|
||||
_, _ = w.Write([]byte(strings.Repeat(" ", (4<<20)+1)))
|
||||
}
|
||||
}))
|
||||
defer upstream.Close()
|
||||
channel := &model.Channel{Type: constant.ChannelTypeVLLM, BaseURL: &upstream.URL, Key: "EMPTY"}
|
||||
status, err := fetchInferenceStatus(context.Background(), channel)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 503, status.Endpoints["/health"].Status)
|
||||
assert.Equal(t, "invalid_response", status.Endpoints["/version"].Error)
|
||||
assert.Empty(t, status.Endpoints["/v1/models"].Error)
|
||||
assert.Equal(t, "response_too_large", status.Endpoints["/metrics"].Error)
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
_, err = fetchInferenceStatus(ctx, channel)
|
||||
assert.ErrorIs(t, err, context.Canceled)
|
||||
}
|
||||
|
||||
func TestParseVLLMMetricsPreservesSeriesAndHistogramTotals(t *testing.T) {
|
||||
metrics, err := parseInferenceMetrics(`# TYPE vllm:num_requests_running gauge
|
||||
vllm:num_requests_running{engine="0",model_name="model,a"} 0
|
||||
vllm:num_requests_running{model_name="model,a",engine="1"} 2
|
||||
# TYPE vllm:time_to_first_token_seconds histogram
|
||||
vllm:time_to_first_token_seconds_bucket{le="+Inf"} 4
|
||||
vllm:time_to_first_token_seconds_sum 10
|
||||
vllm:time_to_first_token_seconds_count 4
|
||||
# TYPE vllm:kv_cache_usage_perc gauge
|
||||
vllm:kv_cache_usage_perc NaN
|
||||
vllm:bad_value +Inf
|
||||
# TYPE process_start_time_seconds gauge
|
||||
process_start_time_seconds 100
|
||||
unrelated_metric 123
|
||||
`, "vllm:")
|
||||
require.NoError(t, err)
|
||||
assert.ElementsMatch(t, []inferenceMetric{
|
||||
{Name: "vllm:num_requests_running", Labels: map[string]string{"engine": "0", "model_name": "model,a"}, Value: 0},
|
||||
{Name: "vllm:num_requests_running", Labels: map[string]string{"engine": "1", "model_name": "model,a"}, Value: 2},
|
||||
{Name: "vllm:time_to_first_token_seconds_sum", Labels: map[string]string{}, Value: 10},
|
||||
{Name: "vllm:time_to_first_token_seconds_count", Labels: map[string]string{}, Value: 4},
|
||||
{Name: "process_start_time_seconds", Labels: map[string]string{}, Value: 100},
|
||||
}, metrics)
|
||||
_, err = parseInferenceMetrics("vllm:broken{engine=oops} 1\n", "vllm:")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
// Opt-in verification against an operator-provided instance; credentials stay in the environment.
|
||||
func TestVLLMStatusLive(t *testing.T) {
|
||||
base := os.Getenv("NEW_API_VLLM_TEST_URL")
|
||||
if base == "" {
|
||||
t.Skip("NEW_API_VLLM_TEST_URL is not configured")
|
||||
}
|
||||
service.InitHttpClient()
|
||||
status, err := fetchInferenceStatus(context.Background(), &model.Channel{
|
||||
Type: constant.ChannelTypeVLLM, BaseURL: &base, Key: os.Getenv("NEW_API_VLLM_TEST_KEY"),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
for path, endpoint := range status.Endpoints {
|
||||
assert.Equal(t, 200, endpoint.Status, path)
|
||||
assert.Empty(t, endpoint.Error, path)
|
||||
}
|
||||
assert.NotEmpty(t, status.Version)
|
||||
assert.NotEmpty(t, status.Models)
|
||||
assert.NotEmpty(t, status.Metrics)
|
||||
assert.NotEmpty(t, status.RawMetrics)
|
||||
}
|
||||
|
||||
func TestSGLangStatusUsesDocumentedEndpointsAndRedactsServerConfiguration(t *testing.T) {
|
||||
service.InitHttpClient()
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
assert.Equal(t, http.MethodGet, r.Method)
|
||||
assert.Equal(t, "Bearer test-key", r.Header.Get("Authorization"))
|
||||
switch r.URL.Path {
|
||||
case "/health":
|
||||
w.WriteHeader(http.StatusOK)
|
||||
case "/server_info":
|
||||
_, _ = w.Write([]byte(`{"version":"0.5.19","api_key":"server-secret","admin_api_key":"admin-secret","internal_states":[{"private":"private-state"}]}`))
|
||||
case "/v1/models":
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"served-model","root":"served-model","max_model_len":32768}]}`))
|
||||
case "/metrics":
|
||||
_, _ = w.Write([]byte("# TYPE sglang:num_running_reqs gauge\nsglang:num_running_reqs{dp_rank=\"0\"} 0\nsglang:num_running_reqs{dp_rank=\"1\"} 2\n# TYPE sglang:time_to_first_token_seconds histogram\nsglang:time_to_first_token_seconds_bucket{le=\"+Inf\"} 2\nsglang:time_to_first_token_seconds_sum 1.5\nsglang:time_to_first_token_seconds_count 2\nvllm:num_requests_running 9\n"))
|
||||
default:
|
||||
t.Errorf("unexpected SGLang path: %s", r.URL.Path)
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}
|
||||
}))
|
||||
defer upstream.Close()
|
||||
status, err := fetchInferenceStatus(context.Background(), &model.Channel{Type: constant.ChannelTypeSGLang, BaseURL: &upstream.URL, Key: "test-key"})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "0.5.19", status.Version)
|
||||
require.Len(t, status.Models, 1)
|
||||
assert.Equal(t, "served-model", status.Models[0].ID)
|
||||
assert.ElementsMatch(t, []inferenceMetric{
|
||||
{Name: "sglang:num_running_reqs", Labels: map[string]string{"dp_rank": "0"}, Value: 0},
|
||||
{Name: "sglang:num_running_reqs", Labels: map[string]string{"dp_rank": "1"}, Value: 2},
|
||||
{Name: "sglang:time_to_first_token_seconds_sum", Labels: map[string]string{}, Value: 1.5},
|
||||
{Name: "sglang:time_to_first_token_seconds_count", Labels: map[string]string{}, Value: 2},
|
||||
}, status.Metrics)
|
||||
encoded, err := common.Marshal(status)
|
||||
require.NoError(t, err)
|
||||
for _, secret := range []string{"test-key", "server-secret", "admin-secret", "private-state", "api_key", "internal_states"} {
|
||||
assert.NotContains(t, string(encoded), secret)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,111 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestMultiKeyEnableRestoresOnlyExhaustedChannels(t *testing.T) {
|
||||
previousDB, previousLogDB := model.DB, model.LOG_DB
|
||||
previousType, previousLogType := common.MainDatabaseType(), common.LogDatabaseType()
|
||||
previousMaster, previousCache, previousRedis, previousSQLite := common.IsMasterNode, common.MemoryCacheEnabled, common.RedisEnabled, common.SQLitePath
|
||||
t.Cleanup(func() {
|
||||
model.DB, model.LOG_DB = previousDB, previousLogDB
|
||||
common.SetDatabaseTypes(previousType, previousLogType)
|
||||
common.IsMasterNode, common.MemoryCacheEnabled, common.RedisEnabled, common.SQLitePath = previousMaster, previousCache, previousRedis, previousSQLite
|
||||
})
|
||||
t.Setenv("SQL_DSN", os.Getenv("TEST_CHANNEL_SQL_DSN"))
|
||||
t.Setenv("LOG_SQL_DSN", "")
|
||||
common.IsMasterNode, common.MemoryCacheEnabled, common.RedisEnabled = false, false, false
|
||||
common.SQLitePath = filepath.Join(t.TempDir(), "channel.db")
|
||||
require.NoError(t, model.InitDB())
|
||||
database := model.DB
|
||||
sqlDB, err := database.DB()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
|
||||
model.LOG_DB = database
|
||||
common.SetLogDatabaseType(common.MainDatabaseType())
|
||||
require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.User{}, &model.Log{}, &model.AuditLog{}))
|
||||
root := &model.User{Username: "multi-key-review-root", Role: common.RoleRootUser, Status: common.UserStatusEnabled}
|
||||
require.NoError(t, database.Create(root).Error)
|
||||
t.Cleanup(func() { require.NoError(t, database.Unscoped().Delete(root).Error) })
|
||||
versionQuery := "SELECT VERSION()"
|
||||
if common.UsingMainDatabase(common.DatabaseTypeSQLite) {
|
||||
versionQuery = "SELECT sqlite_version()"
|
||||
}
|
||||
var version string
|
||||
require.NoError(t, database.Raw(versionQuery).Scan(&version).Error)
|
||||
t.Logf("database=%s version=%s", common.MainDatabaseType(), version)
|
||||
|
||||
for _, cacheEnabled := range []bool{false, true} {
|
||||
for _, action := range []string{"enable_key", "enable_all_keys"} {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
initialStatus int
|
||||
manualOverride string
|
||||
wantStatus int
|
||||
}{
|
||||
{name: "key exhaustion restores", initialStatus: common.ChannelStatusEnabled, wantStatus: common.ChannelStatusEnabled},
|
||||
{name: "manual disable is preserved", initialStatus: common.ChannelStatusManuallyDisabled, wantStatus: common.ChannelStatusManuallyDisabled},
|
||||
{name: "manual disable after exhaustion is preserved", initialStatus: common.ChannelStatusEnabled, manualOverride: "status", wantStatus: common.ChannelStatusManuallyDisabled},
|
||||
{name: "tag disable after exhaustion is preserved", initialStatus: common.ChannelStatusEnabled, manualOverride: "tag", wantStatus: common.ChannelStatusManuallyDisabled},
|
||||
} {
|
||||
t.Run(fmt.Sprintf("cache=%t/%s/%s", cacheEnabled, action, tc.name), func(t *testing.T) {
|
||||
common.MemoryCacheEnabled = cacheEnabled
|
||||
tag := t.Name()
|
||||
channel := &model.Channel{Name: t.Name(), Type: 1, Key: "key-one\nkey-two", Status: tc.initialStatus, Models: "test-model", Group: "default", Tag: &tag,
|
||||
ChannelInfo: model.ChannelInfo{IsMultiKey: true, MultiKeySize: 2, MultiKeyStatusList: map[int]int{1: common.ChannelStatusManuallyDisabled}},
|
||||
}
|
||||
require.NoError(t, channel.Insert())
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, channel.Delete())
|
||||
model.InitChannelCache()
|
||||
})
|
||||
for _, operation := range []string{"disable_key", action} {
|
||||
if operation == action {
|
||||
if tc.manualOverride == "status" {
|
||||
model.UpdateChannelStatus(channel.Id, "", common.ChannelStatusManuallyDisabled, "manual operation")
|
||||
} else if tc.manualOverride == "tag" {
|
||||
require.NoError(t, model.DisableChannelByTag(tag))
|
||||
}
|
||||
}
|
||||
payload, err := common.Marshal(MultiKeyManageRequest{ChannelId: channel.Id, Action: operation, KeyIndex: common.GetPointer(0)})
|
||||
require.NoError(t, err)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Set("id", root.Id)
|
||||
c.Set("role", common.RoleRootUser)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/channel/multi_key", bytes.NewReader(payload))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
ManageMultiKeys(c)
|
||||
var result struct {
|
||||
Success bool `json:"success"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, recorder.Body.String())
|
||||
}
|
||||
loaded, err := model.GetChannelById(channel.Id, true)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tc.wantStatus, loaded.Status)
|
||||
assert.NotContains(t, loaded.ChannelInfo.MultiKeyStatusList, 0)
|
||||
assert.NotContains(t, loaded.ChannelInfo.MultiKeyDisabledReason, 0)
|
||||
assert.NotContains(t, loaded.ChannelInfo.MultiKeyDisabledTime, 0)
|
||||
var ability model.Ability
|
||||
require.NoError(t, database.Where("channel_id = ?", channel.Id).First(&ability).Error)
|
||||
assert.Equal(t, tc.wantStatus == common.ChannelStatusEnabled, ability.Enabled)
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -2,13 +2,25 @@ package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"slices"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
@@ -18,7 +30,7 @@ func TestShouldRetryHonorsPinRetryMode(t *testing.T) {
|
||||
openaiErr := types.NewOpenAIError(errors.New("upstream"), types.ErrorCodeBadResponseStatusCode, http.StatusInternalServerError)
|
||||
|
||||
c := newPinRetryContext()
|
||||
assert.True(t, shouldRetry(c, openaiErr, 1))
|
||||
assert.True(t, service.ShouldRetryRelayError(c, openaiErr, 1))
|
||||
|
||||
origin := newPinRetryContext()
|
||||
service.GetChannelConstraints(origin).AddPin(dto.ChannelPin{
|
||||
@@ -27,7 +39,7 @@ func TestShouldRetryHonorsPinRetryMode(t *testing.T) {
|
||||
Rank: dto.PinRankOriginTask,
|
||||
RetryMode: dto.PinRetrySameChannel,
|
||||
})
|
||||
assert.True(t, shouldRetry(origin, openaiErr, 1), "origin pin retries on the same channel")
|
||||
assert.True(t, service.ShouldRetryRelayError(origin, openaiErr, 1), "origin pin retries on the same channel")
|
||||
|
||||
token := newPinRetryContext()
|
||||
service.GetChannelConstraints(token).AddPin(dto.ChannelPin{
|
||||
@@ -36,14 +48,14 @@ func TestShouldRetryHonorsPinRetryMode(t *testing.T) {
|
||||
Rank: dto.PinRankToken,
|
||||
RetryMode: dto.PinRetrySingleAttempt,
|
||||
})
|
||||
assert.False(t, shouldRetry(token, openaiErr, 1), "token pin suppresses retry")
|
||||
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "pinned_channel", Source: "channel_constraint"}, service.DecideRelayRetry(token, openaiErr, 1), "token pin suppresses retry")
|
||||
}
|
||||
|
||||
func TestShouldRetryTaskRelayHonorsPinRetryMode(t *testing.T) {
|
||||
taskErr := &dto.TaskError{StatusCode: http.StatusInternalServerError}
|
||||
|
||||
c := newPinRetryContext()
|
||||
assert.True(t, shouldRetryTaskRelay(c, 1, taskErr, 1))
|
||||
assert.Equal(t, "retry", decideTaskRetry(c, taskErr, 1).Action)
|
||||
|
||||
origin := newPinRetryContext()
|
||||
service.GetChannelConstraints(origin).AddPin(dto.ChannelPin{
|
||||
@@ -52,7 +64,7 @@ func TestShouldRetryTaskRelayHonorsPinRetryMode(t *testing.T) {
|
||||
Rank: dto.PinRankOriginTask,
|
||||
RetryMode: dto.PinRetrySameChannel,
|
||||
})
|
||||
assert.True(t, shouldRetryTaskRelay(origin, 2, taskErr, 1))
|
||||
assert.Equal(t, "retry", decideTaskRetry(origin, taskErr, 1).Action)
|
||||
|
||||
token := newPinRetryContext()
|
||||
service.GetChannelConstraints(token).AddPin(dto.ChannelPin{
|
||||
@@ -61,7 +73,7 @@ func TestShouldRetryTaskRelayHonorsPinRetryMode(t *testing.T) {
|
||||
Rank: dto.PinRankToken,
|
||||
RetryMode: dto.PinRetrySingleAttempt,
|
||||
})
|
||||
assert.False(t, shouldRetryTaskRelay(token, 1, taskErr, 1))
|
||||
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "pinned_channel", Source: "channel_constraint"}, decideTaskRetry(token, taskErr, 1))
|
||||
}
|
||||
|
||||
func TestSameChannelPinsMergeToStricterRetryMode(t *testing.T) {
|
||||
@@ -84,7 +96,7 @@ func TestSameChannelPinsMergeToStricterRetryMode(t *testing.T) {
|
||||
assert.Equal(t, 7, pin.ChannelId)
|
||||
assert.Equal(t, dto.PinRetrySingleAttempt, pin.RetryMode)
|
||||
assert.Empty(t, overridden)
|
||||
assert.False(t, shouldRetry(c, types.NewOpenAIError(errors.New("upstream"), types.ErrorCodeBadResponseStatusCode, http.StatusInternalServerError), 1))
|
||||
assert.False(t, service.ShouldRetryRelayError(c, types.NewOpenAIError(errors.New("upstream"), types.ErrorCodeBadResponseStatusCode, http.StatusInternalServerError), 1))
|
||||
}
|
||||
|
||||
func newPinRetryContext() *gin.Context {
|
||||
@@ -93,3 +105,115 @@ func newPinRetryContext() *gin.Context {
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
return c
|
||||
}
|
||||
|
||||
func TestRequestPolicyConfigReturnsSettingsWithoutMigration(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
GetRequestPolicy(ctx)
|
||||
assert.Equal(t, http.StatusOK, recorder.Code)
|
||||
var response struct {
|
||||
Success bool `json:"success"`
|
||||
Data map[string]any `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
assert.True(t, response.Success)
|
||||
assert.Contains(t, response.Data, "options")
|
||||
assert.NotContains(t, response.Data, "migration")
|
||||
assert.NotContains(t, response.Data, "differences")
|
||||
}
|
||||
|
||||
func TestRequestPolicyRoutingDatabaseMatrix(t *testing.T) {
|
||||
require.NoError(t, i18n.Init())
|
||||
for _, dialect := range []struct{ kind, env string }{{"sqlite", ""}, {"mysql", "TEST_MYSQL_DSN"}, {"postgres", "TEST_POSTGRES_DSN"}} {
|
||||
t.Run(dialect.kind, func(t *testing.T) {
|
||||
dsn := ""
|
||||
if dialect.env != "" {
|
||||
dsn = os.Getenv(dialect.env)
|
||||
if dsn == "" {
|
||||
t.Skipf("%s not configured", dialect.env)
|
||||
}
|
||||
}
|
||||
db := modelManagementDB(t, dialect.kind, dsn)
|
||||
previousGroups := setting.UserUsableGroups2JSONString()
|
||||
previousRatios, err := common.Marshal(ratio_setting.GetGroupRatioCopy())
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`))
|
||||
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1}`))
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(previousGroups))
|
||||
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(string(previousRatios)))
|
||||
})
|
||||
// Channel 1 holds the session binding but is no longer usable.
|
||||
channels := []model.Channel{
|
||||
{Id: 1, Name: "A", Type: 1, Key: "test-only", Status: common.ChannelStatusAutoDisabled, Models: "policy-test", Group: "default", Priority: common.GetPointer(int64(10))},
|
||||
{Id: 2, Name: "B", Type: 1, Key: "test-only", Status: common.ChannelStatusEnabled, Models: "policy-test", Group: "default", Priority: common.GetPointer(int64(5))},
|
||||
}
|
||||
for i := range channels {
|
||||
require.NoError(t, db.Create(&channels[i]).Error)
|
||||
require.NoError(t, channels[i].AddAbilities(db))
|
||||
}
|
||||
affinity := operation_setting.GetChannelAffinitySetting()
|
||||
previousAffinity := *affinity
|
||||
t.Cleanup(func() { *affinity = previousAffinity })
|
||||
for _, cached := range []bool{false, true} {
|
||||
for _, keep := range []bool{false, true} {
|
||||
for _, tc := range []struct {
|
||||
globalMode, ruleMode string
|
||||
blocked bool
|
||||
}{
|
||||
{"strict", "inherit", true},
|
||||
{"prefer", "strict", true},
|
||||
{"prefer", "inherit", false},
|
||||
{"strict", "prefer", false},
|
||||
} {
|
||||
t.Run(fmt.Sprintf("cache=%t/keep=%t/global=%s/rule=%s", cached, keep, tc.globalMode, tc.ruleMode), func(t *testing.T) {
|
||||
common.MemoryCacheEnabled = cached
|
||||
model.InitChannelCache()
|
||||
snapshot, err := model.BuildRequestPolicy(map[string]string{
|
||||
"channel_affinity_setting.enabled": "true",
|
||||
"channel_affinity_setting.session_mode": tc.globalMode,
|
||||
"channel_affinity_setting.keep_on_channel_disabled": fmt.Sprint(keep),
|
||||
"channel_affinity_setting.rules": fmt.Sprintf(`[{"name":"session","model_regex":[".*"],"key_sources":[{"type":"request_header","key":"X-Session"}],"session_mode":%q}]`, tc.ruleMode),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
*affinity = snapshot.Affinity
|
||||
seed := newPinRetryContext()
|
||||
seed.Request.Header.Set("X-Session", t.Name())
|
||||
_, found := service.GetPreferredChannelByAffinity(seed, "policy-test", "default")
|
||||
require.False(t, found)
|
||||
seed.Set("channel_id", 1)
|
||||
service.RecordChannelAffinity(seed, 1)
|
||||
t.Cleanup(func() { service.ClearCurrentChannelAffinityCache(seed) })
|
||||
bound, found := service.GetPreferredChannelByAffinity(seed, "policy-test", "default")
|
||||
require.True(t, found)
|
||||
require.Equal(t, 1, bound)
|
||||
|
||||
request := newPinRetryContext()
|
||||
request.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"policy-test"}`))
|
||||
request.Request.Header.Set("Content-Type", "application/json")
|
||||
request.Request.Header.Set("X-Session", t.Name())
|
||||
common.SetContextKey(request, constant.ContextKeyUsingGroup, "default")
|
||||
middleware.Distribute()(request)
|
||||
events := service.RequestPolicy(request).Events()
|
||||
assert.True(t, slices.ContainsFunc(events, func(event service.PolicyEvent) bool { return event.Decision.Reason == "session_rule_matched" }), "decision events are recorded for every request")
|
||||
bound, found = service.GetPreferredChannelByAffinity(seed, "policy-test", "default")
|
||||
if tc.blocked {
|
||||
assert.Equal(t, http.StatusServiceUnavailable, request.Writer.Status())
|
||||
assert.True(t, request.IsAborted())
|
||||
assert.Equal(t, keep, found, "binding retention is independent of strict request handling")
|
||||
if found {
|
||||
assert.Equal(t, 1, bound)
|
||||
}
|
||||
return
|
||||
}
|
||||
assert.False(t, request.IsAborted())
|
||||
assert.Equal(t, 2, common.GetContextKeyInt(request, constant.ContextKeyChannelId), "prefer falls back to the next eligible channel")
|
||||
require.True(t, found)
|
||||
assert.Equal(t, 2, bound, "a successful fallback rebinds the session")
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -23,6 +23,33 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGetChannelDefaultBaseURLsUsesBuiltInDefaults(t *testing.T) {
|
||||
originalBaseURLs := constant.ChannelBaseURLs
|
||||
constant.ChannelBaseURLs = append([]string(nil), originalBaseURLs...)
|
||||
constant.ChannelBaseURLs[constant.ChannelTypeDeepSeek] = "https://deepseek.server.example"
|
||||
t.Cleanup(func() {
|
||||
constant.ChannelBaseURLs = originalBaseURLs
|
||||
})
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/api/channel/default_base_urls", nil)
|
||||
GetChannelDefaultBaseURLs(c)
|
||||
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
var response struct {
|
||||
Success bool `json:"success"`
|
||||
Data map[int]string `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
require.True(t, response.Success)
|
||||
assert.Equal(t, "https://deepseek.server.example", response.Data[constant.ChannelTypeDeepSeek])
|
||||
assert.Equal(t, "https://api.openai.com", response.Data[constant.ChannelTypeOpenAI])
|
||||
assert.NotContains(t, response.Data, constant.ChannelTypeAzure)
|
||||
assert.NotContains(t, response.Data, constant.ChannelTypeNewAPI)
|
||||
assert.NotContains(t, response.Data, constant.ChannelTypeTaskPlugin)
|
||||
}
|
||||
|
||||
func TestValidateChannelProxy(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
@@ -398,7 +398,7 @@ func fetchChannelUpstreamModelIDs(channel *model.Channel) ([]string, error) {
|
||||
return normalizeModelNames(models), nil
|
||||
}
|
||||
|
||||
if channel.Type == constant.ChannelTypeAdvancedCustom {
|
||||
if constant.IsAdvancedCustomChannel(channel.Type) {
|
||||
return fetchAdvancedCustomUpstreamModelIDs(channel, baseURL)
|
||||
}
|
||||
|
||||
@@ -417,10 +417,13 @@ func fetchChannelUpstreamModelIDs(channel *model.Channel) ([]string, error) {
|
||||
url = fmt.Sprintf("%s/api/paas/v4/models", baseURL)
|
||||
}
|
||||
case constant.ChannelTypeVolcEngine:
|
||||
// 火山方舟 OpenAI 兼容 API 根路径是 /api/v3(chat/embeddings 均为
|
||||
// {base}/api/v3/...,见 relay/channel/volcengine 的 GetRequestURL),
|
||||
// 不存在 /v1/models 端点。原拼接会导致「获取模型列表」固定 404 失败。
|
||||
if plan, ok := constant.ChannelSpecialBases[baseURL]; ok && plan.OpenAIBaseURL != "" {
|
||||
url = fmt.Sprintf("%s/v1/models", plan.OpenAIBaseURL)
|
||||
} else {
|
||||
url = fmt.Sprintf("%s/v1/models", baseURL)
|
||||
url = fmt.Sprintf("%s/api/v3/models", baseURL)
|
||||
}
|
||||
case constant.ChannelTypeMoonshot:
|
||||
if plan, ok := constant.ChannelSpecialBases[baseURL]; ok && plan.OpenAIBaseURL != "" {
|
||||
@@ -473,7 +476,7 @@ func fetchAdvancedCustomUpstreamModelIDs(channel *model.Channel, baseURL string)
|
||||
RelayMode: relayconstant.RelayModeUnknown,
|
||||
RequestURLPath: dto.AdvancedCustomModelListPath,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelType: constant.ChannelTypeAdvancedCustom,
|
||||
ChannelType: channel.Type,
|
||||
ChannelBaseUrl: baseURL,
|
||||
ApiKey: key,
|
||||
ChannelOtherSettings: channel.GetOtherSettings(),
|
||||
|
||||
@@ -0,0 +1,498 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
|
||||
"github.com/QuantumNous/new-api/relay"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
var detectionKeyPattern = regexp.MustCompile(`^[a-zA-Z0-9_-]{16,96}$`)
|
||||
var detectionCalls sync.Map
|
||||
|
||||
type detectionRequest struct {
|
||||
Group string `json:"group"`
|
||||
Kind string `json:"kind"`
|
||||
Key string `json:"key"`
|
||||
PriceVersion string `json:"price_version"`
|
||||
}
|
||||
|
||||
func detectionInput(c *gin.Context) (detectionRequest, *model.User, service.DetectionPrice, int, error) {
|
||||
var input detectionRequest
|
||||
if err := c.ShouldBindJSON(&input); err != nil {
|
||||
return input, nil, service.DetectionPrice{}, 0, errors.New("invalid detection request")
|
||||
}
|
||||
if input.Kind != "pelican" && input.Kind != "fingerprint" {
|
||||
return input, nil, service.DetectionPrice{}, 0, errors.New("invalid detection type")
|
||||
}
|
||||
user, err := model.GetUserById(c.GetInt("id"), false)
|
||||
if err != nil {
|
||||
return input, nil, service.DetectionPrice{}, 0, err
|
||||
}
|
||||
if !service.IsUserSelectableGroup(user.Group, input.Group) {
|
||||
return input, nil, service.DetectionPrice{}, 0, errors.New("group unavailable")
|
||||
}
|
||||
price, reserve, err := service.DetectionPricing(user.Group, input.Group, input.Kind)
|
||||
return input, user, price, reserve, err
|
||||
}
|
||||
|
||||
func EstimateDetection(c *gin.Context) {
|
||||
input, _, price, reserve, err := detectionInput(c)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
encoded, _ := common.Marshal(price)
|
||||
calls := 1
|
||||
if input.Kind == "fingerprint" {
|
||||
calls = 3
|
||||
}
|
||||
c.JSON(200, gin.H{"success": true, "data": gin.H{"reserved": reserve, "ratio": price.Snapshot.GroupRatio, "calls": calls, "price_version": model.DetectionSecretHash(string(encoded))}})
|
||||
}
|
||||
|
||||
func CreateUserDetection(c *gin.Context) {
|
||||
input, user, price, reserve, err := detectionInput(c)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if !detectionKeyPattern.MatchString(input.Key) {
|
||||
common.ApiError(c, errors.New("invalid idempotency key"))
|
||||
return
|
||||
}
|
||||
var previous []model.DetectionRun
|
||||
key := fmt.Sprintf("user:%d:%s", user.Id, input.Key)
|
||||
if err = model.DB.Where("request_key = ?", key).Find(&previous).Error; err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if len(previous) > 0 {
|
||||
if previous[0].GroupName != input.Group || previous[0].Kind != input.Kind {
|
||||
common.ApiError(c, errors.New("idempotency key conflict"))
|
||||
return
|
||||
}
|
||||
c.JSON(200, gin.H{"success": true, "data": gin.H{"id": previous[0].ID}})
|
||||
return
|
||||
}
|
||||
encoded, _ := common.Marshal(price)
|
||||
if input.PriceVersion != model.DetectionSecretHash(string(encoded)) {
|
||||
common.ApiError(c, errors.New("pricing changed; request a new estimate"))
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
price.Request.At = &now
|
||||
encoded, _ = common.Marshal(price)
|
||||
channel, err := model.GetRandomSatisfiedChannel(input.Group, fingerprint.Model, 0, nil)
|
||||
if err != nil || channel == nil || !service.ChannelSupportsAstra(channel) {
|
||||
common.ApiError(c, errors.New("no Astra channel available"))
|
||||
return
|
||||
}
|
||||
id, err := common.GenerateRandomCharsKey(32)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
run := &model.DetectionRun{ID: id, UserID: user.Id, RequestKey: fmt.Sprintf("user:%d:%s", user.Id, input.Key), ChannelID: channel.Id, GroupName: input.Group, Kind: input.Kind, Status: "pending", CreatedAt: time.Now().Unix(), Reserved: reserve, Price: string(encoded)}
|
||||
if err = model.CreateDetection(run); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(200, gin.H{"success": true, "data": gin.H{"id": run.ID}})
|
||||
}
|
||||
|
||||
func GetUserDetection(c *gin.Context) {
|
||||
var run model.DetectionRun
|
||||
query := model.DB.Where("user_id = ?", c.GetInt("id"))
|
||||
if c.Param("id") == "latest" {
|
||||
query = query.Where("group_name = ?", c.Query("group")).Order("created_at desc")
|
||||
} else {
|
||||
query = query.Where("id = ?", c.Param("id"))
|
||||
}
|
||||
if err := query.First(&run).Error; err != nil {
|
||||
if c.Param("id") == "latest" {
|
||||
c.JSON(200, gin.H{"success": true, "data": nil})
|
||||
} else {
|
||||
c.JSON(404, gin.H{"success": false, "message": "Detection not found"})
|
||||
}
|
||||
return
|
||||
}
|
||||
var calls []model.DetectionCall
|
||||
if err := model.DB.Where("run_id = ?", run.ID).Find(&calls).Error; err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
usage := []any{}
|
||||
for _, call := range calls {
|
||||
var parsed any
|
||||
_ = common.UnmarshalJsonStr(call.Usage, &parsed)
|
||||
if parsed != nil {
|
||||
usage = append(usage, parsed)
|
||||
}
|
||||
if call.Error != "" && run.Error == "" {
|
||||
run.Error = "usage_or_generation_incomplete"
|
||||
}
|
||||
}
|
||||
c.JSON(200, gin.H{"success": true, "data": gin.H{"id": run.ID, "kind": run.Kind, "status": run.Status, "quality": run.Quality, "reserved": run.Reserved, "charged": run.Charged, "settled": run.Settled, "usage": usage, "error": run.Error}})
|
||||
}
|
||||
|
||||
func GetDetectionArtwork(c *gin.Context) {
|
||||
getDetectionArtwork(c, false)
|
||||
}
|
||||
|
||||
func GetSiteDetectionArtwork(c *gin.Context) {
|
||||
getDetectionArtwork(c, true)
|
||||
}
|
||||
|
||||
func getDetectionArtwork(c *gin.Context, allGroups bool) {
|
||||
var run model.DetectionRun
|
||||
if err := model.DB.Select("html", "channel_id", "created_at").Where("id = ? AND created_at > ?", c.Param("id"), time.Now().Add(-6*time.Hour).Unix()).First(&run).Error; err != nil || run.HTML == "" {
|
||||
c.Status(404)
|
||||
return
|
||||
}
|
||||
settings, err := model.GetStatusPageGroups(c.Request.Context())
|
||||
if err != nil {
|
||||
c.Status(http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
channels, err := model.GetStatusChannels(c.Request.Context())
|
||||
if err != nil {
|
||||
c.Status(http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
usable := service.GetUserUsableGroups("")
|
||||
visible := false
|
||||
for _, channel := range channels {
|
||||
if channel.Id != run.ChannelID {
|
||||
continue
|
||||
}
|
||||
for group := range strings.SplitSeq(channel.Group, ",") {
|
||||
group = strings.TrimSpace(group)
|
||||
if _, ok := usable[group]; (allGroups || (ok && !settings[group].Hidden)) && run.CreatedAt > settings[group].ResetBefore {
|
||||
visible = true
|
||||
}
|
||||
}
|
||||
}
|
||||
if !visible {
|
||||
c.Status(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
c.Header("Content-Security-Policy", "sandbox allow-scripts; default-src 'none'; script-src 'unsafe-inline'; style-src 'unsafe-inline'; img-src data:; form-action 'none'; base-uri 'none'; frame-ancestors 'self'")
|
||||
c.Header("X-Content-Type-Options", "nosniff")
|
||||
c.Header("Cache-Control", "no-store")
|
||||
c.Data(200, "text/html; charset=utf-8", []byte(sanitizeDetectionHTML(run.HTML)))
|
||||
}
|
||||
|
||||
// Decode only generation fields: external credentials cannot forward tools,
|
||||
// previous responses, billing modifiers or arbitrary provider parameters.
|
||||
func detectionProxyRequest(data []byte) (*dto.OpenAIResponsesRequest, error) {
|
||||
var fields map[string]any
|
||||
if len(data) > 64<<10 || common.Unmarshal(data, &fields) != nil {
|
||||
return nil, errors.New("invalid request")
|
||||
}
|
||||
if name, ok := fields["model"].(string); !ok || name != fingerprint.Model {
|
||||
return nil, errors.New("model not allowed")
|
||||
}
|
||||
for key := range fields {
|
||||
switch key {
|
||||
case "model", "input", "instructions", "stream", "max_output_tokens", "reasoning", "store", "tools", "tool_choice":
|
||||
default:
|
||||
return nil, errors.New("unsupported request field")
|
||||
}
|
||||
}
|
||||
input, err := common.Marshal(fields["input"])
|
||||
if err != nil || len(input) > 16384 {
|
||||
return nil, errors.New("input too large")
|
||||
}
|
||||
var checkText func(any) bool
|
||||
checkText = func(v any) bool {
|
||||
switch value := v.(type) {
|
||||
case string:
|
||||
return true
|
||||
case []any:
|
||||
for _, item := range value {
|
||||
if !checkText(item) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
case map[string]any:
|
||||
for k, item := range value {
|
||||
switch k {
|
||||
case "role":
|
||||
role, ok := item.(string)
|
||||
if !ok || (role != "user" && role != "system" && role != "developer") {
|
||||
return false
|
||||
}
|
||||
case "type":
|
||||
kind, ok := item.(string)
|
||||
if !ok || (kind != "input_text" && kind != "message") {
|
||||
return false
|
||||
}
|
||||
case "content", "text":
|
||||
if !checkText(item) {
|
||||
return false
|
||||
}
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
if !checkText(fields["input"]) {
|
||||
return nil, errors.New("text input required")
|
||||
}
|
||||
instructions, _ := fields["instructions"].(string)
|
||||
if len(input)+len(instructions) > 16384 {
|
||||
return nil, errors.New("input too large")
|
||||
}
|
||||
// Service-selected bounds cannot be overridden by the bearer.
|
||||
fields = map[string]any{"model": fingerprint.Model, "input": fields["input"], "instructions": instructions, "stream": true, "max_output_tokens": 16384, "reasoning": map[string]string{"effort": "high"}, "store": false, "tools": []any{}}
|
||||
encoded, _ := common.Marshal(fields)
|
||||
var request dto.OpenAIResponsesRequest
|
||||
err = common.Unmarshal(encoded, &request)
|
||||
return &request, err
|
||||
}
|
||||
|
||||
func DetectionProxy(c *gin.Context) {
|
||||
secret := strings.TrimPrefix(c.GetHeader("Authorization"), "Bearer ")
|
||||
if !strings.HasPrefix(secret, "det_") {
|
||||
c.AbortWithStatus(401)
|
||||
return
|
||||
}
|
||||
data, err := io.ReadAll(io.LimitReader(c.Request.Body, (64<<10)+1))
|
||||
if err != nil {
|
||||
c.AbortWithStatus(400)
|
||||
return
|
||||
}
|
||||
request, err := detectionProxyRequest(data)
|
||||
if err != nil {
|
||||
c.JSON(400, gin.H{"error": gin.H{"message": err.Error()}})
|
||||
return
|
||||
}
|
||||
call, err := model.ClaimDetectionCall(model.DetectionSecretHash(secret))
|
||||
if err != nil {
|
||||
c.AbortWithStatus(401)
|
||||
return
|
||||
}
|
||||
var run model.DetectionRun
|
||||
if err = model.DB.First(&run, "id = ?", call.RunID).Error; err != nil {
|
||||
c.AbortWithStatus(503)
|
||||
return
|
||||
}
|
||||
ctx, cancel := context.WithDeadline(c.Request.Context(), time.Unix(run.Deadline-10, 0))
|
||||
detectionCalls.Store(call.ID, cancel)
|
||||
defer func() { cancel(); detectionCalls.Delete(call.ID) }()
|
||||
channel, err := model.GetChannelById(run.ChannelID, true)
|
||||
var raw []byte
|
||||
var contentType string
|
||||
// Flush SSE comments while generation is buffered for accounting. Without
|
||||
// these bytes, the public edge times out a valid long reasoning request.
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("X-Accel-Buffering", "no")
|
||||
c.Writer.WriteHeader(http.StatusOK)
|
||||
_, _ = c.Writer.WriteString(": detection started\n\n")
|
||||
c.Writer.Flush()
|
||||
if err == nil && channel.Status == common.ChannelStatusEnabled {
|
||||
raw, contentType, err = waitDetectionResponse(c, ctx, channel, request)
|
||||
} else {
|
||||
err = errors.New("channel unavailable")
|
||||
}
|
||||
success := err == nil
|
||||
call.Available = &success
|
||||
if err != nil {
|
||||
call.Error = "channel_call_failed"
|
||||
if ctx.Err() != nil {
|
||||
call.Available = nil
|
||||
call.Error = "interrupted_usage_unknown"
|
||||
}
|
||||
} else {
|
||||
usage, complete := detectionResponseUsage(raw, contentType)
|
||||
if !complete {
|
||||
call.Error = "incomplete_response"
|
||||
call.Available = nil
|
||||
} else {
|
||||
call.Usage = usage.Raw
|
||||
quota, quotaErr := service.DetectionUsageQuota(run.Price, usage.Get("input_tokens").Int(), usage.Get("output_tokens").Int(), usage.Get("input_tokens_details.cached_tokens").Int(), max(usage.Get("input_tokens_details.cached_creation_tokens").Int(), usage.Get("input_tokens_details.cache_creation_tokens").Int()))
|
||||
if quotaErr != nil || quota > run.Reserved {
|
||||
call.Error = "settlement_unavailable"
|
||||
} else if run.UserID > 0 {
|
||||
call.Quota = quota
|
||||
}
|
||||
}
|
||||
}
|
||||
if success {
|
||||
if html := detectionHTML(raw, contentType); html != "" {
|
||||
if saveErr := model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Update("html", html).Error; saveErr != nil {
|
||||
call.Error = "artwork_save_failed"
|
||||
success = false
|
||||
}
|
||||
}
|
||||
}
|
||||
if err = model.FinishDetectionCall(call); err != nil {
|
||||
success = false
|
||||
}
|
||||
if !success {
|
||||
_, _ = c.Writer.WriteString("event: error\ndata: {\"type\":\"error\",\"error\":{\"message\":\"Detection generation failed\"}}\n\n")
|
||||
c.Writer.Flush()
|
||||
return
|
||||
}
|
||||
if strings.Contains(contentType, "text/event-stream") {
|
||||
_, _ = c.Writer.Write(raw)
|
||||
} else {
|
||||
event, _ := common.Marshal(map[string]any{"type": "response.completed", "response": gjson.ParseBytes(raw).Value()})
|
||||
_, _ = fmt.Fprintf(c.Writer, "event: response.completed\ndata: %s\n\n", event)
|
||||
}
|
||||
c.Writer.Flush()
|
||||
}
|
||||
|
||||
func waitDetectionResponse(c *gin.Context, ctx context.Context, channel *model.Channel, request *dto.OpenAIResponsesRequest) ([]byte, string, error) {
|
||||
type response struct {
|
||||
raw []byte
|
||||
contentType string
|
||||
err error
|
||||
}
|
||||
done := make(chan response, 1)
|
||||
go func() {
|
||||
raw, contentType, err := runDetectionResponse(ctx, channel, request)
|
||||
done <- response{raw, contentType, err}
|
||||
}()
|
||||
ticker := time.NewTicker(15 * time.Second)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case result := <-done:
|
||||
return result.raw, result.contentType, result.err
|
||||
case <-ticker.C:
|
||||
_, _ = c.Writer.WriteString(": waiting for generation\n\n")
|
||||
c.Writer.Flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func runDetectionResponse(ctx context.Context, channel *model.Channel, request *dto.OpenAIResponsesRequest) ([]byte, string, error) {
|
||||
encoded, err := common.Marshal(request)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequestWithContext(ctx, http.MethodPost, "/v1/responses", bytes.NewReader(encoded))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
if middleware.SetupContextForSelectedChannel(c, channel, fingerprint.Model) != nil {
|
||||
return nil, "", errors.New("channel setup failed")
|
||||
}
|
||||
info, err := relaycommon.GenRelayInfo(c, types.RelayFormatOpenAIResponses, request, nil)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
info.IsChannelTest = true
|
||||
info.DisablePing = true
|
||||
info.AstraDetectionDisabled = true
|
||||
info.AstraAutoSwitch = false
|
||||
adaptor, body, closer, apiErr := relay.PrepareResponsesRequest(c, info, request)
|
||||
if closer != nil {
|
||||
defer closer.Close()
|
||||
}
|
||||
if apiErr != nil {
|
||||
return nil, "", errors.New("preparation failed")
|
||||
}
|
||||
body, boundedCloser, err := boundDetectionBody(body, request)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
defer boundedCloser.Close()
|
||||
response, err := adaptor.DoRequest(c, info, body)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
resp, ok := response.(*http.Response)
|
||||
if !ok || resp == nil {
|
||||
return nil, "", errors.New("unsupported transport")
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != 200 {
|
||||
return nil, "", errors.New("upstream request failed")
|
||||
}
|
||||
raw, err := io.ReadAll(io.LimitReader(resp.Body, (8<<20)+1))
|
||||
if len(raw) > 8<<20 {
|
||||
return nil, "", errors.New("response too large")
|
||||
}
|
||||
return raw, resp.Header.Get("Content-Type"), err
|
||||
}
|
||||
|
||||
func detectionResponseUsage(raw []byte, contentType string) (gjson.Result, bool) {
|
||||
result := gjson.ParseBytes(raw)
|
||||
if strings.Contains(contentType, "text/event-stream") {
|
||||
result = gjson.Result{}
|
||||
for line := range strings.SplitSeq(string(raw), "\n") {
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
event := gjson.Parse(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
|
||||
switch event.Get("type").String() {
|
||||
case "response.completed", "response.incomplete", "response.failed":
|
||||
result = event.Get("response")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
usage := result.Get("usage")
|
||||
status := result.Get("status").String()
|
||||
return usage, (status == "completed" || status == "incomplete" || status == "failed") && usage.Get("input_tokens").Type == gjson.Number && usage.Get("output_tokens").Type == gjson.Number
|
||||
}
|
||||
|
||||
// Apply billing and safety limits AFTER channel parameter overrides. Never let
|
||||
// an administrative override turn a bounded paid test into unbounded tool work.
|
||||
func boundDetectionBody(body common.ReplayableBody, request *dto.OpenAIResponsesRequest) (common.ReplayableBody, io.Closer, error) {
|
||||
raw, err := io.ReadAll(io.LimitReader(body, 128<<10))
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
var prepared map[string]any
|
||||
if common.Unmarshal(raw, &prepared) != nil {
|
||||
return nil, nil, errors.New("invalid prepared detection request")
|
||||
}
|
||||
var input any
|
||||
if common.Unmarshal(request.Input, &input) != nil {
|
||||
return nil, nil, errors.New("invalid detection input")
|
||||
}
|
||||
prepared["input"] = input
|
||||
var instructions any
|
||||
_ = common.Unmarshal(request.Instructions, &instructions)
|
||||
prepared["instructions"] = instructions
|
||||
prepared["max_output_tokens"] = *request.MaxOutputTokens
|
||||
if *request.MaxOutputTokens == 16384 {
|
||||
prepared["reasoning"] = map[string]string{"effort": "high"}
|
||||
}
|
||||
prepared["tools"] = []any{}
|
||||
prepared["tool_choice"] = "none"
|
||||
prepared["store"] = false
|
||||
prepared["stream"] = true
|
||||
for _, key := range []string{"previous_response_id", "conversation", "background", "max_tokens", "max_completion_tokens"} {
|
||||
delete(prepared, key)
|
||||
}
|
||||
encoded, err := common.Marshal(prepared)
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
return relaycommon.NewOutboundJSONBody(encoded)
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"strings"
|
||||
|
||||
"golang.org/x/net/html"
|
||||
)
|
||||
|
||||
// Artwork scripts run only inside an opaque-origin sandbox. CSP blocks network
|
||||
// access; navigation, forms, and embedded documents are removed here.
|
||||
func sanitizeDetectionHTML(source string) string {
|
||||
root, err := html.Parse(strings.NewReader(source))
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
var walk func(*html.Node)
|
||||
walk = func(n *html.Node) {
|
||||
for child := n.FirstChild; child != nil; {
|
||||
next := child.NextSibling
|
||||
switch strings.ToLower(child.Data) {
|
||||
case "iframe", "frame", "frameset", "object", "embed", "base", "meta", "link", "form", "input", "textarea", "select", "audio", "video", "source", "foreignobject":
|
||||
if child.Type == html.ElementNode {
|
||||
n.RemoveChild(child)
|
||||
child = next
|
||||
continue
|
||||
}
|
||||
}
|
||||
if child.Type == html.ElementNode {
|
||||
attrs := child.Attr[:0]
|
||||
for _, attr := range child.Attr {
|
||||
name := strings.ToLower(attr.Key)
|
||||
value := strings.TrimSpace(attr.Val)
|
||||
if strings.HasPrefix(name, "on") {
|
||||
continue
|
||||
}
|
||||
switch name {
|
||||
case "href":
|
||||
if !strings.HasPrefix(value, "#") {
|
||||
continue
|
||||
}
|
||||
case "src":
|
||||
if child.Data != "img" || !(strings.HasPrefix(value, "data:image/png;") || strings.HasPrefix(value, "data:image/jpeg;") || strings.HasPrefix(value, "data:image/webp;")) {
|
||||
continue
|
||||
}
|
||||
case "srcset", "srcdoc", "action", "formaction", "ping", "target", "download", "http-equiv":
|
||||
continue
|
||||
}
|
||||
attrs = append(attrs, attr)
|
||||
}
|
||||
child.Attr = attrs
|
||||
}
|
||||
walk(child)
|
||||
child = next
|
||||
}
|
||||
}
|
||||
walk(root)
|
||||
var out bytes.Buffer
|
||||
if html.Render(&out, root) != nil {
|
||||
return ""
|
||||
}
|
||||
return out.String()
|
||||
}
|
||||
@@ -0,0 +1,171 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestDetectionProxyFlushesBeforeGenerationAndConsumesCredentialOnce(t *testing.T) {
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
connection, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
connection.SetMaxOpenConns(1)
|
||||
oldDB := model.DB
|
||||
model.DB = db
|
||||
t.Cleanup(func() { model.DB = oldDB; _ = connection.Close() })
|
||||
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}))
|
||||
release := make(chan struct{})
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
select {
|
||||
case <-release:
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = io.WriteString(w, `{"status":"completed","usage":{"input_tokens":10,"output_tokens":5},"output":[{"type":"message","content":[{"type":"output_text","text":"<svg></svg>"}]}]}`)
|
||||
}))
|
||||
defer upstream.Close()
|
||||
url := upstream.URL
|
||||
channel := model.Channel{Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: fingerprint.Model}
|
||||
require.NoError(t, db.Create(&channel).Error)
|
||||
price, _ := common.Marshal(service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: "p+c", ExprHash: billingexpr.ExprHashString("p+c"), GroupRatio: 1, QuotaPerUnit: 500000, ExprVersion: 1}})
|
||||
run := model.DetectionRun{ID: "stream-test", RequestKey: "stream-test", Kind: "pelican", ChannelID: channel.Id, Status: "running", Deadline: time.Now().Unix() + 60, Reserved: 100, Price: string(price)}
|
||||
require.NoError(t, db.Create(&run).Error)
|
||||
call := model.DetectionCall{ID: "stream-call", RunID: run.ID, SecretHash: model.DetectionSecretHash("det_test"), State: "ready", ExpiresAt: time.Now().Unix() + 60}
|
||||
require.NoError(t, db.Create(&call).Error)
|
||||
router := gin.New()
|
||||
router.POST("/responses", DetectionProxy)
|
||||
gateway := httptest.NewServer(router)
|
||||
defer gateway.Close()
|
||||
client := http.Client{Timeout: 5 * time.Second}
|
||||
request, err := http.NewRequest(http.MethodPost, gateway.URL+"/responses", strings.NewReader(`{"model":"gpt-6-astra","input":"draw"}`))
|
||||
require.NoError(t, err)
|
||||
request.Header.Set("Authorization", "Bearer det_test")
|
||||
response, err := client.Do(request)
|
||||
close(release)
|
||||
require.NoError(t, err)
|
||||
defer response.Body.Close()
|
||||
reader := bufio.NewReader(response.Body)
|
||||
first, err := reader.ReadString('\n')
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, ": detection started\n", first)
|
||||
raw, err := io.ReadAll(reader)
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, string(raw), "response.completed")
|
||||
require.NoError(t, db.First(&call, "id = ?", call.ID).Error)
|
||||
assert.Equal(t, "finished", call.State)
|
||||
assert.Contains(t, call.Usage, "input_tokens")
|
||||
replay := httptest.NewRecorder()
|
||||
request = httptest.NewRequest(http.MethodPost, "/responses", strings.NewReader(`{"model":"gpt-6-astra","input":"draw"}`))
|
||||
request.Header.Set("Authorization", "Bearer det_test")
|
||||
router.ServeHTTP(replay, request)
|
||||
assert.Equal(t, 401, replay.Code)
|
||||
}
|
||||
|
||||
func TestDetectionProxyRejectsEscapesAndEnforcesBounds(t *testing.T) {
|
||||
for _, body := range []string{
|
||||
`{"model":[],"input":"hi"}`,
|
||||
`{"model":"gpt-6-luna","input":"hi"}`,
|
||||
`{"model":"gpt-6-astra","input":[{"role":[],"content":"hi"}]}`,
|
||||
`{"model":"gpt-6-astra","input":"hi","previous_response_id":"x"}`,
|
||||
`{"model":"gpt-6-astra","input":[{"type":"input_image","image_url":"file:///etc/passwd"}]}`,
|
||||
} {
|
||||
_, err := detectionProxyRequest([]byte(body))
|
||||
require.Error(t, err)
|
||||
}
|
||||
request, err := detectionProxyRequest([]byte(`{"model":"gpt-6-astra","input":"draw a pelican","max_output_tokens":1000000,"tools":[{"type":"web_search"}]}`))
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, request.MaxOutputTokens)
|
||||
assert.EqualValues(t, 16384, *request.MaxOutputTokens)
|
||||
assert.Equal(t, "[]", string(request.Tools))
|
||||
}
|
||||
|
||||
func TestDetectionOnlyBillsCompleteReportedUsage(t *testing.T) {
|
||||
answer, parseErr := fingerprint.ReadResponse(strings.NewReader(`{"status":"incomplete","usage":{"input_tokens":100,"output_tokens":10,"input_tokens_details":{"cached_tokens":40,"cache_creation_tokens":5}}}`), "application/json")
|
||||
require.ErrorIs(t, parseErr, fingerprint.ErrIncomplete)
|
||||
assert.True(t, answer.HasUsage)
|
||||
assert.EqualValues(t, 5, answer.CacheWriteTokens)
|
||||
assert.Empty(t, answer.Text)
|
||||
_, ok := detectionResponseUsage([]byte(`{"status":"completed"}`), "application/json")
|
||||
assert.False(t, ok)
|
||||
_, ok = detectionResponseUsage([]byte("data: {\"type\":\"response.output_text.delta\",\"delta\":\"hello\"}\n\n"), "text/event-stream")
|
||||
assert.False(t, ok)
|
||||
usage, ok := detectionResponseUsage([]byte("data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"usage\":{\"input_tokens\":100,\"output_tokens\":10}}}\n\n"), "text/event-stream")
|
||||
require.True(t, ok)
|
||||
assert.EqualValues(t, 100, usage.Get("input_tokens").Int())
|
||||
expression := `p * 2 + cr * 0.2 + c * 10`
|
||||
price := service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: expression, ExprHash: billingexpr.ExprHashString(expression), GroupRatio: 0.18, QuotaPerUnit: 500000, ExprVersion: 1}}
|
||||
encoded, err := common.Marshal(price)
|
||||
require.NoError(t, err)
|
||||
quota, err := service.DetectionUsageQuota(string(encoded), 100, 10, 40)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 21, quota)
|
||||
_, err = service.DetectionUsageQuota(string(encoded), 10, 20, 11)
|
||||
assert.Error(t, err)
|
||||
at := time.Date(2026, 9, 24, 9, 0, 0, 0, time.UTC)
|
||||
price.Snapshot.ExprString = `hour("UTC") == 9 ? p * 2 : p * 100`
|
||||
price.Snapshot.ExprHash = billingexpr.ExprHashString(price.Snapshot.ExprString)
|
||||
price.Request.At = &at
|
||||
encoded, err = common.Marshal(price)
|
||||
require.NoError(t, err)
|
||||
quota, err = service.DetectionUsageQuota(string(encoded), 100, 0, 0)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 18, quota)
|
||||
}
|
||||
|
||||
func TestDetectionArtworkCannotNavigate(t *testing.T) {
|
||||
cleaned := sanitizeDetectionHTML(`<meta http-equiv="refresh" content="0;url=https://evil.test"><script>alert(1)</script><a href="https://evil.test" onclick="alert(1)">x</a><iframe srcdoc="bad"></iframe><svg><circle cx="20" cy="20" r="10"/></svg><style>circle{fill:red}</style>`)
|
||||
for _, bad := range []string{"evil.test", "onclick", "<iframe", "http-equiv"} {
|
||||
assert.NotContains(t, cleaned, bad)
|
||||
}
|
||||
assert.Contains(t, cleaned, "<circle")
|
||||
assert.Contains(t, cleaned, "<style>")
|
||||
}
|
||||
|
||||
func TestDetectionPriceBoundIncludesIntermediateTiers(t *testing.T) {
|
||||
bound, err := billingexpr.DetectionCostBound(`len < 100 ? p * 2 : len < 1000 ? p * 100 : p * 1`, 32768, 16384)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, float64(3276800), bound)
|
||||
_, err = billingexpr.DetectionCostBound(`p / (c + 1)`, 32768, 16384)
|
||||
require.Error(t, err)
|
||||
bound, err = billingexpr.DetectionCostBound(`len <= 272000 ? tier("standard", p * 10 + c * 50 + cr * 1 + cc * 12) : tier("long", p * 20 + c * 100 + cr * 2 + cc * 25)`, 32768, 16384)
|
||||
require.NoError(t, err)
|
||||
assert.Greater(t, bound, float64(16384*100))
|
||||
}
|
||||
|
||||
func TestDetectionLatestPerChannelStatus(t *testing.T) {
|
||||
now := time.Now().Unix()
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
ids []int
|
||||
evidence map[int]service.ChannelDetectionEvidence
|
||||
want int
|
||||
}{
|
||||
{"successful fingerprint", []int{1}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: true}}, 1},
|
||||
{"failed availability", []int{1}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: false}}, 0},
|
||||
{"partial coverage", []int{1, 2}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: true}}, 2},
|
||||
{"failure and unknown", []int{1, 2}, map[int]service.ChannelDetectionEvidence{1: {At: now, Success: false}}, -1},
|
||||
{"expired success", []int{1}, map[int]service.ChannelDetectionEvidence{1: {At: now - 7201, Success: true}}, -1},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
status, _ := service.DetectionFallbackStatus(tc.ids, tc.evidence, now)
|
||||
assert.Equal(t, tc.want, status)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,410 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/hmac"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
type detectionHandler struct{}
|
||||
|
||||
func pelicanCredential(id string) string {
|
||||
mac := hmac.New(sha256.New, []byte(common.SessionSecret))
|
||||
mac.Write([]byte("newapi-pelican-credential-v1:" + id))
|
||||
return "det_" + hex.EncodeToString(mac.Sum(nil))
|
||||
}
|
||||
|
||||
func (detectionHandler) Type() string { return "group_detection" }
|
||||
func (detectionHandler) Enabled() bool {
|
||||
return common.GetEnvOrDefaultBool("PELICAN_DETECTION_ENABLED", true)
|
||||
}
|
||||
func (detectionHandler) Interval() time.Duration { return 15 * time.Second }
|
||||
func (detectionHandler) NewPayload() any { return nil }
|
||||
|
||||
func (detectionHandler) Run(ctx context.Context, task *model.SystemTask, runner string) {
|
||||
_ = model.DeliverDetectionLogs()
|
||||
// Keep accounting facts, but expire bulky generated artwork after the public window.
|
||||
_ = model.DB.Model(&model.DetectionRun{}).Where("settled = ? AND created_at < ? AND html <> ?", true, time.Now().Add(-6*time.Hour).Unix(), "").Update("html", "").Error
|
||||
scheduleErr := schedulePelicanDetections(ctx)
|
||||
var err error
|
||||
{
|
||||
var runs []model.DetectionRun
|
||||
err = model.DB.Where("settled = ?", false).Order("created_at asc").Find(&runs).Error
|
||||
if err == nil {
|
||||
jobs := make(chan model.DetectionRun)
|
||||
var workers sync.WaitGroup
|
||||
for range 2 {
|
||||
workers.Go(func() {
|
||||
for run := range jobs {
|
||||
processDetection(ctx, run)
|
||||
}
|
||||
})
|
||||
}
|
||||
for _, run := range runs {
|
||||
select {
|
||||
case jobs <- run:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
close(jobs)
|
||||
workers.Wait()
|
||||
}
|
||||
}
|
||||
if err == nil {
|
||||
err = scheduleErr
|
||||
}
|
||||
status := model.SystemTaskStatusSucceeded
|
||||
if err != nil {
|
||||
status = model.SystemTaskStatusFailed
|
||||
}
|
||||
finishSystemTaskHandler(task, runner, status, nil, err)
|
||||
}
|
||||
|
||||
func schedulePelicanDetections(ctx context.Context) error {
|
||||
channels, schedules, err := service.AutomaticDetectionChannelsDue(ctx, "pelican", time.Now())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
for _, channel := range channels {
|
||||
groups := channel.GetGroups()
|
||||
if len(groups) == 0 {
|
||||
continue
|
||||
}
|
||||
price, reserve, priceErr := service.DetectionPricing("", groups[0], "pelican")
|
||||
if priceErr != nil {
|
||||
return priceErr
|
||||
}
|
||||
created := time.Unix(now, 0)
|
||||
price.Request.At = &created
|
||||
encoded, _ := common.Marshal(price)
|
||||
id, err := common.GenerateRandomCharsKey(32)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
run := &model.DetectionRun{ID: id, RequestKey: fmt.Sprintf("adaptive:%d:%s", channel.Id, schedules[channel.Id].LastResultID), ChannelID: channel.Id, GroupName: groups[0], Kind: "pelican", Status: "pending", CreatedAt: now, Price: string(encoded), Reserved: reserve}
|
||||
if err = model.CreateDetection(run); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func processDetection(ctx context.Context, run model.DetectionRun) {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
if run.Deadline == 0 {
|
||||
run.Deadline = time.Now().Unix() + 600
|
||||
if model.DB.Model(&model.DetectionRun{}).Where("id = ? AND deadline = 0", run.ID).Updates(map[string]any{"deadline": run.Deadline, "status": "running"}).Error != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
taskCtx, taskCancel := context.WithDeadline(ctx, time.Unix(run.Deadline, 0))
|
||||
defer taskCancel()
|
||||
if time.Now().Unix() >= run.Deadline {
|
||||
var calls []model.DetectionCall
|
||||
if model.DB.Where("run_id = ?", run.ID).Find(&calls).Error != nil {
|
||||
return
|
||||
}
|
||||
for _, call := range calls {
|
||||
if call.State != "running" {
|
||||
continue
|
||||
}
|
||||
if cancel, ok := detectionCalls.Load(call.ID); ok {
|
||||
cancel.(context.CancelFunc)()
|
||||
return
|
||||
}
|
||||
// All proxy contexts expire before the task deadline. A vanished process
|
||||
// cannot resume or recharge this consumed credential.
|
||||
call.Error = "interrupted_usage_unknown"
|
||||
call.Quota = 0
|
||||
_ = model.FinishDetectionCall(&call)
|
||||
}
|
||||
if run.RemoteID != "" {
|
||||
cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
|
||||
_, _ = manxueRunRequest(cancelCtx, &run, http.MethodDelete, "/api/v1/tests/"+run.RemoteID, nil, "")
|
||||
cancel()
|
||||
}
|
||||
finishDetection(run, "failed", "detection_timeout")
|
||||
return
|
||||
}
|
||||
if run.Status == "succeeded" || run.Status == "failed" {
|
||||
settleDetectionAndSchedule(run)
|
||||
return
|
||||
}
|
||||
if run.Kind == "fingerprint" {
|
||||
processPaidFingerprint(ctx, run)
|
||||
return
|
||||
}
|
||||
if run.RemoteID == localPelicanID {
|
||||
processLocalPelican(taskCtx, run)
|
||||
return
|
||||
}
|
||||
if run.RemoteID == "" {
|
||||
var existing int64
|
||||
if model.DB.Model(&model.DetectionCall{}).Where("run_id = ?", run.ID).Count(&existing).Error != nil {
|
||||
return
|
||||
}
|
||||
// A domain-separated PRF allows safe retry after a submit/commit crash;
|
||||
// only its digest is stored, and the random run ID prevents enumeration.
|
||||
secret := pelicanCredential(run.ID)
|
||||
call := model.DetectionCall{ID: run.ID, RunID: run.ID, SecretHash: model.DetectionSecretHash(secret), ExpiresAt: run.Deadline + 300, State: "ready"}
|
||||
if existing == 0 && model.DB.Create(&call).Error != nil {
|
||||
return
|
||||
}
|
||||
payload := map[string]any{"benchmark": "pelican", "base_url": "https://ai.shrinken.pw/api/detection-proxy/v1", "api_key": secret, "model": fingerprint.Model, "protocol": "responses", "reasoning_effort": "high"}
|
||||
var result gjson.Result
|
||||
var err error
|
||||
for taskCtx.Err() == nil {
|
||||
result, err = manxueRunRequest(taskCtx, &run, http.MethodPost, "/api/v1/tests", payload, run.ID)
|
||||
if errors.Is(err, errDetectionServiceUnavailable) && taskCtx.Err() == nil {
|
||||
processLocalPelican(taskCtx, run)
|
||||
return
|
||||
}
|
||||
if err == nil {
|
||||
break
|
||||
}
|
||||
select {
|
||||
case <-taskCtx.Done():
|
||||
case <-time.After(5 * time.Second):
|
||||
}
|
||||
}
|
||||
if err != nil || taskCtx.Err() != nil {
|
||||
if ctx.Err() == nil {
|
||||
processDetection(ctx, run)
|
||||
}
|
||||
return
|
||||
}
|
||||
remoteID := result.Get("id").String()
|
||||
if remoteID == "" || len(remoteID) > 128 || strings.ContainsAny(remoteID, "/?#\\") {
|
||||
finishDetection(run, "failed", "invalid_detection_response")
|
||||
return
|
||||
}
|
||||
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"remote_id": remoteID, "status": "running"}).Error != nil {
|
||||
return
|
||||
}
|
||||
run.RemoteID = remoteID
|
||||
}
|
||||
// Each worker retains its slot until completion, limiting remote generation
|
||||
// concurrency as well as local submission concurrency.
|
||||
for time.Now().Unix() < run.Deadline {
|
||||
result, err := manxueRunRequest(taskCtx, &run, http.MethodGet, "/api/v1/tests/"+run.RemoteID, nil, "")
|
||||
if errors.Is(err, errDetectionServiceUnavailable) && taskCtx.Err() == nil {
|
||||
processLocalPelican(taskCtx, run)
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-time.After(5 * time.Second):
|
||||
}
|
||||
continue
|
||||
}
|
||||
switch result.Get("status").String() {
|
||||
case "succeeded":
|
||||
quality := result.Get("assessment.quality").String()
|
||||
if quality != "normal" && quality != "degraded" {
|
||||
quality = "unknown"
|
||||
}
|
||||
html := result.Get("result.html").String()
|
||||
if len(html) > 2<<20 {
|
||||
html = ""
|
||||
}
|
||||
metadata, _ := common.Marshal(map[string]string{"source_url": manxueArtworkURL(result)})
|
||||
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"quality": quality, "html": html, "result": string(metadata)}).Error != nil {
|
||||
return
|
||||
}
|
||||
finishDetection(run, "succeeded", "")
|
||||
return
|
||||
case "failed", "cancelled":
|
||||
finishDetection(run, "failed", "detection_service_failed")
|
||||
return
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-time.After(5 * time.Second):
|
||||
}
|
||||
}
|
||||
processDetection(ctx, run)
|
||||
}
|
||||
|
||||
func finishDetection(run model.DetectionRun, status, reason string) {
|
||||
updates := map[string]any{"status": status, "error": reason}
|
||||
if status == "failed" && run.Kind == "pelican" {
|
||||
updates["result"] = ""
|
||||
}
|
||||
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(updates).Error != nil {
|
||||
return
|
||||
}
|
||||
settleDetectionAndSchedule(run)
|
||||
}
|
||||
|
||||
func settleDetectionAndSchedule(run model.DetectionRun) {
|
||||
if model.SettleDetection(run.ID) != nil || run.UserID != 0 || run.Kind != "pelican" {
|
||||
return
|
||||
}
|
||||
// Reconciliation also repairs the clock if a process exits after settlement.
|
||||
if err := model.DB.Select("id", "channel_id", "finished_at", "status", "quality").First(&run, "id = ?", run.ID).Error; err != nil {
|
||||
return
|
||||
}
|
||||
outcome := "unknown"
|
||||
if run.Status == "succeeded" {
|
||||
outcome = run.Quality
|
||||
}
|
||||
_ = model.ApplyDetectionScheduleResult(context.Background(), run.ChannelID, "pelican", run.ID, run.FinishedAt, outcome)
|
||||
}
|
||||
|
||||
func manxueRequest(ctx context.Context, method, path string, payload any, key string) (gjson.Result, error) {
|
||||
client := &http.Client{Timeout: 35 * time.Second, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
|
||||
return manxueRequestWithClient(ctx, client, method, path, payload, key)
|
||||
}
|
||||
|
||||
func manxueRequestWithClient(ctx context.Context, client *http.Client, method, path string, payload any, key string) (gjson.Result, error) {
|
||||
var encoded []byte
|
||||
if payload != nil {
|
||||
var err error
|
||||
encoded, err = common.Marshal(payload)
|
||||
if err != nil {
|
||||
return gjson.Result{}, err
|
||||
}
|
||||
}
|
||||
for attempt := 0; attempt < 3; attempt++ {
|
||||
req, err := http.NewRequestWithContext(ctx, method, "https://manxue.ai"+path, bytes.NewReader(encoded))
|
||||
if err != nil {
|
||||
return gjson.Result{}, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", "Mozilla/5.0")
|
||||
if strings.HasPrefix(path, "/api/visitor/") {
|
||||
req.Header.Set("Origin", "https://manxue.ai")
|
||||
}
|
||||
if key != "" {
|
||||
req.Header.Set("Idempotency-Key", key)
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return gjson.Result{}, errDetectionServiceUnavailable
|
||||
}
|
||||
body, readErr := io.ReadAll(io.LimitReader(resp.Body, (3<<20)+1))
|
||||
resp.Body.Close()
|
||||
if readErr != nil || resp.StatusCode >= 500 {
|
||||
return gjson.Result{}, errDetectionServiceUnavailable
|
||||
}
|
||||
if resp.StatusCode == 429 || resp.StatusCode == 503 {
|
||||
seconds, _ := strconv.Atoi(resp.Header.Get("Retry-After"))
|
||||
if at, err := http.ParseTime(resp.Header.Get("Retry-After")); err == nil {
|
||||
seconds = int(time.Until(at).Seconds()) + 1
|
||||
}
|
||||
seconds = max(seconds, 2<<attempt)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return gjson.Result{}, ctx.Err()
|
||||
case <-time.After(time.Duration(seconds) * time.Second):
|
||||
}
|
||||
continue
|
||||
}
|
||||
if readErr != nil || len(body) > 3<<20 || resp.StatusCode < 200 || resp.StatusCode >= 300 || !gjson.ValidBytes(body) {
|
||||
return gjson.Result{}, errors.New("invalid detection service response")
|
||||
}
|
||||
return gjson.ParseBytes(body), nil
|
||||
}
|
||||
return gjson.Result{}, errDetectionServiceUnavailable
|
||||
}
|
||||
|
||||
func processPaidFingerprint(ctx context.Context, run model.DetectionRun) {
|
||||
channel, err := model.GetChannelById(run.ChannelID, true)
|
||||
if err != nil || channel.Status != common.ChannelStatusEnabled {
|
||||
finishDetection(run, "failed", "channel_unavailable")
|
||||
return
|
||||
}
|
||||
samples := []fingerprint.Sample{}
|
||||
for index := 0; index < 3; index++ {
|
||||
id := fmt.Sprintf("%s_%d", run.ID, index)
|
||||
var previous []model.DetectionCall
|
||||
if model.DB.Where("id = ?", id).Find(&previous).Error != nil {
|
||||
return
|
||||
}
|
||||
if len(previous) > 0 {
|
||||
if previous[0].State == "running" {
|
||||
return
|
||||
}
|
||||
var sample fingerprint.Sample
|
||||
if previous[0].Sample != "" && common.UnmarshalJsonStr(previous[0].Sample, &sample) == nil {
|
||||
samples = append(samples, sample)
|
||||
}
|
||||
continue
|
||||
}
|
||||
call := model.DetectionCall{ID: id, RunID: run.ID, SecretHash: model.DetectionSecretHash(id), State: "running", StartedAt: time.Now().Unix()}
|
||||
if model.DB.Create(&call).Error != nil {
|
||||
return
|
||||
}
|
||||
probeCtx, cancel := context.WithDeadline(ctx, time.Unix(run.Deadline-10, 0))
|
||||
detectionCalls.Store(call.ID, cancel)
|
||||
answer, count, probeErr := runAstraFingerprintProbe(probeCtx, channel, true)
|
||||
interrupted := probeCtx.Err() != nil
|
||||
cancel()
|
||||
available := probeErr == nil || answer.HasUsage
|
||||
call.Available = &available
|
||||
if !answer.HasUsage && (errors.Is(probeErr, fingerprint.ErrInvalidSample) || errors.Is(probeErr, fingerprint.ErrIncomplete) || interrupted) {
|
||||
call.Available = nil
|
||||
}
|
||||
if probeErr != nil {
|
||||
call.Error = "channel_call_failed"
|
||||
}
|
||||
if answer.HasUsage {
|
||||
usage, _ := common.Marshal(answer)
|
||||
call.Usage = string(usage)
|
||||
quota, quotaErr := service.DetectionUsageQuota(run.Price, answer.InputTokens, answer.OutputTokens, answer.CachedTokens, answer.CacheWriteTokens)
|
||||
if quotaErr != nil || quota > run.Reserved/3 {
|
||||
call.Error = "settlement_unavailable"
|
||||
} else {
|
||||
call.Quota = quota
|
||||
}
|
||||
} else {
|
||||
call.Error = "usage_unavailable"
|
||||
}
|
||||
if probeErr == nil && fingerprint.ValidArray(answer.Text, count) {
|
||||
sample := fingerprint.Sample{Text: answer.Text, Count: count}
|
||||
samples = append(samples, sample)
|
||||
encoded, _ := common.Marshal(sample)
|
||||
call.Sample = string(encoded)
|
||||
}
|
||||
if model.FinishDetectionCall(&call) != nil {
|
||||
detectionCalls.Delete(call.ID)
|
||||
return
|
||||
}
|
||||
detectionCalls.Delete(call.ID)
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
result, err := fingerprint.Score(samples)
|
||||
if err != nil {
|
||||
finishDetection(run, "failed", "scoring_failed")
|
||||
return
|
||||
}
|
||||
encoded, _ := common.Marshal(result)
|
||||
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"quality": result.Status, "result": string(encoded)}).Error != nil {
|
||||
return
|
||||
}
|
||||
finishDetection(run, "succeeded", "")
|
||||
}
|
||||
@@ -0,0 +1,28 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestSuccessfulUserRequestRestoresStatus(t *testing.T) {
|
||||
now := time.Now().Unix()
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
success, probe int64
|
||||
disabled bool
|
||||
want int
|
||||
}{
|
||||
{"fresh user overrides failed probe", now, now - 10, false, 1},
|
||||
{"newer failed probe wins", now - 10, now, false, 0},
|
||||
{"expired success", now - 901, now - 1000, false, 0},
|
||||
{"disabled remains disabled", now, now - 10, true, 0},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
m := Monitor{Status: 0, ChannelDisabled: tc.disabled}
|
||||
applyUserSuccess(&m, tc.success, tc.probe, now)
|
||||
assert.Equal(t, tc.want, m.Status)
|
||||
})
|
||||
}
|
||||
}
|
||||
+4
-4
@@ -22,7 +22,7 @@ func GetAllLogs(c *gin.Context) {
|
||||
group := c.Query("group")
|
||||
requestId := c.Query("request_id")
|
||||
upstreamRequestId := c.Query("upstream_request_id")
|
||||
logs, total, err := model.GetAllLogs(logType, startTimestamp, endTimestamp, modelName, username, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), channel, group, requestId, upstreamRequestId)
|
||||
logs, total, err := model.GetAllLogs(logType, startTimestamp, endTimestamp, modelName, username, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), channel, group, requestId, upstreamRequestId, c.Query("correlation_key"), c.Query("attempt_status"))
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
@@ -49,7 +49,7 @@ func GetUserLogs(c *gin.Context) {
|
||||
group := c.Query("group")
|
||||
requestId := c.Query("request_id")
|
||||
upstreamRequestId := c.Query("upstream_request_id")
|
||||
logs, total, err := model.GetUserLogs(userId, logType, startTimestamp, endTimestamp, modelName, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), group, requestId, upstreamRequestId)
|
||||
logs, total, err := model.GetUserLogs(userId, logType, startTimestamp, endTimestamp, modelName, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), group, requestId, upstreamRequestId,c.Query("attempt_status"))
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
@@ -109,7 +109,7 @@ func GetLogsStat(c *gin.Context) {
|
||||
modelName := c.Query("model_name")
|
||||
channel, _ := strconv.Atoi(c.Query("channel"))
|
||||
group := c.Query("group")
|
||||
stat, err := model.SumUsedQuota(logType, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group)
|
||||
stat, err := model.SumUsedQuota(logType, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group, c.Query("correlation_key"), c.Query("attempt_status"))
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
@@ -136,7 +136,7 @@ func GetLogsSelfStat(c *gin.Context) {
|
||||
modelName := c.Query("model_name")
|
||||
channel, _ := strconv.Atoi(c.Query("channel"))
|
||||
group := c.Query("group")
|
||||
quotaNum, err := model.SumUsedQuota(logType, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group)
|
||||
quotaNum, err := model.SumUsedQuota(logType, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group,"",c.Query("attempt_status"))
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
|
||||
@@ -30,17 +30,18 @@ func VerifyLogin(c *gin.Context) {
|
||||
writeSecurityOperationError(c, service.ErrProofMethod)
|
||||
return
|
||||
}
|
||||
bundle, err := service.VerifyLoginCode(request.FlowToken, request.Code, c.ClientIP(), c.Request.UserAgent())
|
||||
bundle, migration, err := service.VerifyLoginCode(request.FlowToken, request.Code, c.ClientIP(), c.Request.UserAgent())
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
completeVerifiedLoginResponse(c, bundle, service.VerificationMethodTwoFA)
|
||||
completeVerifiedLoginResponse(c, bundle, service.VerificationMethodTwoFA, migration)
|
||||
}
|
||||
|
||||
func LoginPasskeyBegin(c *gin.Context) {
|
||||
var request struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
RPID string `json:"rp_id"`
|
||||
}
|
||||
if common.DecodeJson(c.Request.Body, &request) != nil || request.FlowToken == "" {
|
||||
common.ApiErrorMsg(c, "参数错误")
|
||||
@@ -56,7 +57,11 @@ func LoginPasskeyBegin(c *gin.Context) {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
credentialRPID := ""
|
||||
if credential.RPID != nil {
|
||||
credentialRPID = *credential.RPID
|
||||
}
|
||||
wa, rpIDs, err := passkeysvc.BuildLoginWebAuthn(c.Request, request.RPID, credentialRPID)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
@@ -75,7 +80,7 @@ func LoginPasskeyBegin(c *gin.Context) {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, gin.H{"flow_token": token, "expires_at": expiresAt, "options": options})
|
||||
common.ApiSuccess(c, gin.H{"flow_token": token, "expires_at": expiresAt, "options": options, "rp_ids": rpIDs})
|
||||
}
|
||||
|
||||
func LoginPasskeyFinish(c *gin.Context) {
|
||||
@@ -104,6 +109,7 @@ func LoginPasskeyFinish(c *gin.Context) {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
c.Set("passkey_rp_id", sessionData.RelyingPartyID)
|
||||
if security.LoginFlowID != verification.Flow.Id || sessionData.UserVerification != protocol.VerificationRequired {
|
||||
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
|
||||
return
|
||||
@@ -113,29 +119,33 @@ func LoginPasskeyFinish(c *gin.Context) {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
wa, err := passkeysvc.BuildWebAuthnForRPID(c.Request, sessionData.RelyingPartyID)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if credential.RPID != nil && *credential.RPID != "" && *credential.RPID != sessionData.RelyingPartyID {
|
||||
writeSecurityOperationError(c, service.ErrVerificationFailed)
|
||||
return
|
||||
}
|
||||
validated, err := wa.ValidateLogin(passkeysvc.NewWebAuthnUser(&model.User{Id: identity.UserID}, credential), *sessionData, parsed)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
if err := model.UpdatePasskeyAssertionState(identity.UserID, validated, time.Now()); err != nil {
|
||||
if err := model.UpdatePasskeyAssertionState(identity.UserID, validated, time.Now(), sessionData.RelyingPartyID); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
bundle, err := service.CompleteLoginVerification(request.FlowToken, verification, service.VerificationMethodPasskey, c.ClientIP(), c.Request.UserAgent())
|
||||
bundle, migration, err := service.CompleteLoginVerification(request.FlowToken, verification, service.VerificationMethodPasskey, c.ClientIP(), c.Request.UserAgent())
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
completeVerifiedLoginResponse(c, bundle, service.VerificationMethodPasskey)
|
||||
completeVerifiedLoginResponse(c, bundle, service.VerificationMethodPasskey, migration)
|
||||
}
|
||||
|
||||
func completeVerifiedLoginResponse(c *gin.Context, bundle *service.AuthBundle, method string) {
|
||||
func completeVerifiedLoginResponse(c *gin.Context, bundle *service.AuthBundle, method string, migration *service.LegacyGitHubMigration) {
|
||||
identity, err := service.ParseAccessToken(bundle.AccessToken)
|
||||
if err != nil {
|
||||
writeAuthSessionError(c, err)
|
||||
@@ -147,5 +157,13 @@ func completeVerifiedLoginResponse(c *gin.Context, bundle *service.AuthBundle, m
|
||||
return
|
||||
}
|
||||
c.Set("login_verification_method", method)
|
||||
if migration != nil {
|
||||
// The legacy GitHub binding was rewritten together with this session.
|
||||
notificationFailed := service.NotifyAccountSecurityChange(user.Email, "Login account linked: GitHub") != nil
|
||||
recordLegacyGitHubBindingAudit(c, user, true, map[string]any{
|
||||
"legacy_id": migration.LegacyID, "provider_user_id": migration.GitHubID,
|
||||
"verification_method": method, "notification_failed": notificationFailed,
|
||||
})
|
||||
}
|
||||
writeLoginResponse(c, user, bundle)
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"github.com/tidwall/gjson"
|
||||
"net/url"
|
||||
)
|
||||
|
||||
func manxueArtworkURL(result gjson.Result) string {
|
||||
for _, key := range []string{"share_url", "public_url", "result_url", "result.share_url", "result.public_url"} {
|
||||
raw := result.Get(key).String()
|
||||
u, err := url.Parse(raw)
|
||||
if err == nil && u.Scheme == "https" && u.Host == "manxue.ai" && u.User == nil && u.Query().Get("result") != "" {
|
||||
return u.String()
|
||||
}
|
||||
}
|
||||
id := result.Get("public_id").String()
|
||||
if id == "" {
|
||||
id = result.Get("result.public_id").String()
|
||||
}
|
||||
if id == "" && result.Get("public").Bool() {
|
||||
id = result.Get("result.id").String()
|
||||
}
|
||||
if id == "" {
|
||||
return ""
|
||||
}
|
||||
return "https://manxue.ai/?" + url.Values{"result": {id}}.Encode()
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/tidwall/gjson"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestManxueArtworkURL(t *testing.T) {
|
||||
assert.Empty(t, manxueArtworkURL(gjson.Parse(`{"public":false,"result":{"id":"private"}}`)))
|
||||
assert.Equal(t, "https://manxue.ai/?result=community-one", manxueArtworkURL(gjson.Parse(`{"public":true,"result":{"id":"community-one"}}`)))
|
||||
assert.Empty(t, manxueArtworkURL(gjson.Parse(`{"share_url":"https://evil.test/?result=one"}`)))
|
||||
}
|
||||
func TestArtworkPreservesScriptDrawnLegs(t *testing.T) {
|
||||
result := sanitizeDetectionHTML(`<svg><path id="leg"></path></svg><button id="toggle">pause</button><script>document.getElementById('leg').setAttribute('d','M0 0L10 10')</script>`)
|
||||
assert.Contains(t, result, "<script>")
|
||||
assert.Contains(t, result, "setAttribute")
|
||||
assert.Contains(t, result, "<button")
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/cookiejar"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// Visitor sessions belong to a single detection. Keep them in its private task
|
||||
// metadata until completion so polling can resume after a process restart.
|
||||
type manxueVisitorSession struct {
|
||||
Cookies []*http.Cookie `json:"manxue_cookies,omitempty"`
|
||||
}
|
||||
|
||||
func manxueRunRequest(ctx context.Context, run *model.DetectionRun, method, path string, payload any, key string) (gjson.Result, error) {
|
||||
var session manxueVisitorSession
|
||||
_ = common.UnmarshalJsonStr(run.Result, &session)
|
||||
if len(session.Cookies) == 0 && !(method == http.MethodPost && run.RemoteID == "") {
|
||||
return manxueRequest(ctx, method, path, payload, key)
|
||||
}
|
||||
jar, _ := cookiejar.New(nil)
|
||||
origin, _ := url.Parse("https://manxue.ai/api/visitor/")
|
||||
client := &http.Client{Timeout: 35 * time.Second, Jar: jar, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
|
||||
if len(session.Cookies) == 0 {
|
||||
if _, err := manxueRequestWithClient(ctx, client, http.MethodGet, "/api/visitor/config", nil, ""); err != nil {
|
||||
return gjson.Result{}, errDetectionServiceUnavailable
|
||||
}
|
||||
session.Cookies = jar.Cookies(origin)
|
||||
if len(session.Cookies) == 0 {
|
||||
return gjson.Result{}, errDetectionServiceUnavailable
|
||||
}
|
||||
encoded, err := common.Marshal(session)
|
||||
if err != nil {
|
||||
return gjson.Result{}, err
|
||||
}
|
||||
if err = model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Update("result", string(encoded)).Error; err != nil {
|
||||
return gjson.Result{}, err
|
||||
}
|
||||
run.Result = string(encoded)
|
||||
} else {
|
||||
jar.SetCookies(origin, session.Cookies)
|
||||
}
|
||||
if method == http.MethodPost {
|
||||
fields, ok := payload.(map[string]any)
|
||||
if !ok {
|
||||
return gjson.Result{}, errors.New("invalid pelican payload")
|
||||
}
|
||||
public := make(map[string]any, len(fields)+1)
|
||||
for k, v := range fields {
|
||||
public[k] = v
|
||||
}
|
||||
public["public"] = true
|
||||
payload = public
|
||||
}
|
||||
return manxueRequestWithClient(ctx, client, method, strings.Replace(path, "/api/v1/tests", "/api/visitor/tests", 1), payload, key)
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type manxueTestTransport func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f manxueTestTransport) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
|
||||
|
||||
func TestManxuePublicDetectionResumesAndReturnsPublicURL(t *testing.T) {
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
conn, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
conn.SetMaxOpenConns(1)
|
||||
oldDB, oldTransport := model.DB, http.DefaultTransport
|
||||
model.DB = db
|
||||
t.Cleanup(func() { model.DB = oldDB; http.DefaultTransport = oldTransport; _ = conn.Close() })
|
||||
require.NoError(t, db.AutoMigrate(&model.DetectionRun{}))
|
||||
configCalls, createCalls := 0, 0
|
||||
http.DefaultTransport = manxueTestTransport(func(r *http.Request) (*http.Response, error) {
|
||||
w := httptest.NewRecorder()
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
switch r.URL.Path {
|
||||
case "/api/visitor/config":
|
||||
configCalls++
|
||||
http.SetCookie(w, &http.Cookie{Name: "gallery_visitor", Value: "test-session", Path: "/api/visitor/", Secure: true, HttpOnly: true})
|
||||
_, _ = w.WriteString(`{"prompt":"pelican"}`)
|
||||
case "/api/visitor/tests":
|
||||
createCalls++
|
||||
cookie, err := r.Cookie("gallery_visitor")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "test-session", cookie.Value)
|
||||
require.Equal(t, "stable-key", r.Header.Get("Idempotency-Key"))
|
||||
var body map[string]any
|
||||
require.NoError(t, json.NewDecoder(r.Body).Decode(&body))
|
||||
require.Equal(t, true, body["public"])
|
||||
_, _ = w.WriteString(`{"id":"remote-one","status":"running"}`)
|
||||
case "/api/visitor/tests/remote-one":
|
||||
_, err := r.Cookie("gallery_visitor")
|
||||
require.NoError(t, err)
|
||||
_, _ = w.WriteString(`{"status":"succeeded","public_id":"community-one","result":{"html":"<html>pelican</html>"}}`)
|
||||
case "/api/v1/tests/legacy":
|
||||
require.Empty(t, r.Header.Get("Cookie"))
|
||||
_, _ = w.WriteString(`{"status":"succeeded"}`)
|
||||
default:
|
||||
t.Fatalf("unexpected path: %s", r.URL.Path)
|
||||
}
|
||||
return w.Result(), nil
|
||||
})
|
||||
run := model.DetectionRun{ID: "public-test", RequestKey: "public-test", Kind: "pelican"}
|
||||
require.NoError(t, db.Create(&run).Error)
|
||||
input := map[string]any{"benchmark": "pelican"}
|
||||
_, err = manxueRunRequest(context.Background(), &run, http.MethodPost, "/api/v1/tests", input, "stable-key")
|
||||
require.NoError(t, err)
|
||||
require.NotContains(t, input, "public")
|
||||
var resumed model.DetectionRun
|
||||
require.NoError(t, db.First(&resumed, "id = ?", run.ID).Error)
|
||||
resumed.RemoteID = "remote-one"
|
||||
result, err := manxueRunRequest(context.Background(), &resumed, http.MethodGet, "/api/v1/tests/remote-one", nil, "")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "https://manxue.ai/?result=community-one", manxueArtworkURL(result))
|
||||
require.Equal(t, 1, configCalls)
|
||||
require.Equal(t, 1, createCalls)
|
||||
legacy := model.DetectionRun{RemoteID: "legacy"}
|
||||
_, err = manxueRunRequest(context.Background(), &legacy, http.MethodGet, "/api/v1/tests/legacy", nil, "")
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, manxueArtworkURL(gjson.Parse(`{"id":"private","public":false}`)))
|
||||
}
|
||||
+3
-3
@@ -44,10 +44,10 @@ func TestStatus(c *gin.Context) {
|
||||
func GetStatus(c *gin.Context) {
|
||||
|
||||
cs := console_setting.GetConsoleSetting()
|
||||
passkeySetting := system_setting.PasskeySettingsSnapshot()
|
||||
common.OptionMapRWMutex.RLock()
|
||||
defer common.OptionMapRWMutex.RUnlock()
|
||||
|
||||
passkeySetting := system_setting.GetPasskeySettings()
|
||||
legalSetting := system_setting.GetLegalSettings()
|
||||
|
||||
data := gin.H{
|
||||
@@ -117,8 +117,8 @@ func GetStatus(c *gin.Context) {
|
||||
"oidc_display_name": system_setting.GetOIDCSettings().GetEffectiveDisplayName(),
|
||||
"passkey_login": passkeySetting.Enabled,
|
||||
"passkey_display_name": passkeySetting.RPDisplayName,
|
||||
"passkey_rp_id": passkeySetting.RPID,
|
||||
"passkey_origins": passkeySetting.Origins,
|
||||
"passkey_rp_id": passkeySetting.EffectiveRPID(),
|
||||
"passkey_rp_ids": passkeySetting.RelyingPartyIDs(),
|
||||
"passkey_allow_insecure": passkeySetting.AllowInsecureOrigin,
|
||||
"passkey_user_verification": passkeySetting.UserVerification,
|
||||
"passkey_attachment": passkeySetting.AttachmentPreference,
|
||||
|
||||
@@ -195,6 +195,18 @@ func getModelListGroups(c *gin.Context) (modelListGroups, error) {
|
||||
return modelListGroups{}, err
|
||||
}
|
||||
}
|
||||
if service.RequestRoutingPolicy(c) != nil {
|
||||
entries := service.ResolveRequestPool(c, "")
|
||||
ownerGroups := make([]string, 0, len(entries))
|
||||
for _, entry := range entries {
|
||||
ownerGroups = append(ownerGroups, entry.Group)
|
||||
}
|
||||
return modelListGroups{
|
||||
userGroup: userGroup,
|
||||
tokenGroup: tokenGroup,
|
||||
ownerGroups: ownerGroups,
|
||||
}, nil
|
||||
}
|
||||
|
||||
if tokenGroup == "auto" {
|
||||
return modelListGroups{
|
||||
|
||||
@@ -493,6 +493,44 @@ func TestListModelsTokenLimitUsesResolvedCustomAutoGroups(t *testing.T) {
|
||||
require.Empty(t, anthropicResponse.LastID)
|
||||
}
|
||||
|
||||
func TestListModelsUsesTokenRoutingPolicyGroups(t *testing.T) {
|
||||
withSelfUseModeEnabled(t)
|
||||
originalUsableGroups := setting.UserUsableGroups2JSONString()
|
||||
originalRatios := ratio_setting.GroupRatio2JSONString()
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"gpt":"GPT","domestic":"Domestic"}`))
|
||||
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"gpt":1,"domestic":1}`))
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups))
|
||||
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios))
|
||||
})
|
||||
|
||||
db := setupModelListControllerTestDB(t)
|
||||
require.NoError(t, db.Create(&[]model.Ability{
|
||||
{Group: "gpt", Model: "zz-gpt-model", ChannelId: 1, Enabled: true},
|
||||
{Group: "domestic", Model: "zz-deepseek-model", ChannelId: 2, Enabled: true},
|
||||
}).Error)
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
common.SetContextKey(ctx, constant.ContextKeyUserGroup, "gpt")
|
||||
common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "gpt")
|
||||
ctx.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{
|
||||
Version: 1,
|
||||
Mode: "custom",
|
||||
Groups: []model.TokenGroupEntry{
|
||||
{Group: "gpt", Priority: 0},
|
||||
{Group: "domestic", Priority: 1},
|
||||
},
|
||||
})
|
||||
|
||||
ListModels(ctx, constant.ChannelTypeOpenAI)
|
||||
|
||||
ids := decodeListModelsResponse(t, recorder)
|
||||
require.Contains(t, ids, "zz-gpt-model")
|
||||
require.Contains(t, ids, "zz-deepseek-model")
|
||||
}
|
||||
|
||||
func TestSetupLoginDoesNotTouchPasswordWhenPasswordFieldOmitted(t *testing.T) {
|
||||
db := setupModelListControllerTestDB(t)
|
||||
require.NoError(t, db.AutoMigrate(&model.Log{}, &model.AuditLog{}, &model.UserSession{}, &model.TwoFA{}, &model.PasskeyCredential{}))
|
||||
@@ -517,7 +555,7 @@ func TestSetupLoginDoesNotTouchPasswordWhenPasswordFieldOmitted(t *testing.T) {
|
||||
Role: user.Role,
|
||||
Status: user.Status,
|
||||
Group: user.Group,
|
||||
}, c)
|
||||
}, nil, c)
|
||||
})
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
|
||||
@@ -14,8 +14,13 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
"github.com/QuantumNous/new-api/setting/config"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
@@ -63,7 +68,7 @@ func modelManagementDB(t *testing.T, kind, dsn string) *gorm.DB {
|
||||
for _, value := range restoreRatios {
|
||||
require.NoError(t, value.restore("{}"))
|
||||
}
|
||||
config.UpdateConfigFromMap(config.GlobalConfig.Get("billing_setting"), map[string]string{"billing_mode": "{}", "billing_expr": "{}"})
|
||||
config.UpdateConfigFromMap(config.GlobalConfig.Get("billing_setting"), map[string]string{"billing_mode": "{}", "billing_expr": "{}", "plugin_billing_expr": "{}"})
|
||||
var version string
|
||||
query := "SELECT version()"
|
||||
if kind == "sqlite" {
|
||||
@@ -75,7 +80,7 @@ func modelManagementDB(t *testing.T, kind, dsn string) *gorm.DB {
|
||||
for _, value := range restoreRatios {
|
||||
require.NoError(t, value.restore(value.value))
|
||||
}
|
||||
config.UpdateConfigFromMap(config.GlobalConfig.Get("billing_setting"), map[string]string{"billing_mode": previousConfig["billing_setting.billing_mode"], "billing_expr": previousConfig["billing_setting.billing_expr"]})
|
||||
config.UpdateConfigFromMap(config.GlobalConfig.Get("billing_setting"), map[string]string{"billing_mode": previousConfig["billing_setting.billing_mode"], "billing_expr": previousConfig["billing_setting.billing_expr"], "plugin_billing_expr": previousConfig[billing_setting.PluginBillingExprOption]})
|
||||
common.OptionMap = previousOptions
|
||||
common.IsMasterNode, common.SQLitePath = previousMaster, previousSQLite
|
||||
common.RedisEnabled, common.MemoryCacheEnabled = previousRedis, previousMemory
|
||||
@@ -104,6 +109,369 @@ func modelManagementRequest(t *testing.T, handler gin.HandlerFunc, method, path
|
||||
return recorder
|
||||
}
|
||||
|
||||
func TestModelPricingConversionDatabaseMatrix(t *testing.T) {
|
||||
previousQuota := common.QuotaPerUnit
|
||||
common.QuotaPerUnit = 500000
|
||||
t.Cleanup(func() { common.QuotaPerUnit = previousQuota })
|
||||
for _, dialect := range []struct{ kind, env string }{{"sqlite", ""}, {"mysql", "TEST_MYSQL_DSN"}, {"postgres", "TEST_POSTGRES_DSN"}} {
|
||||
t.Run(dialect.kind, func(t *testing.T) {
|
||||
if dialect.env != "" && os.Getenv(dialect.env) == "" {
|
||||
t.Skip("set " + dialect.env + " to run this database")
|
||||
}
|
||||
db := modelManagementDB(t, dialect.kind, os.Getenv(dialect.env))
|
||||
contract, err := os.ReadFile("../pkg/billingexpr/testdata/frontend_simulation.json")
|
||||
require.NoError(t, err)
|
||||
var fixtures []struct {
|
||||
Name, Expression string
|
||||
Conversion *struct {
|
||||
ModelName string `json:"model_name"`
|
||||
Pricing model.PricingValues `json:"pricing"`
|
||||
}
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(contract, &fixtures))
|
||||
for _, fixture := range fixtures {
|
||||
if fixture.Conversion == nil {
|
||||
continue
|
||||
}
|
||||
t.Run("frontend_contract/"+fixture.Name, func(t *testing.T) {
|
||||
var response struct {
|
||||
Success bool
|
||||
Data model.ModelPricingConversion
|
||||
}
|
||||
modelManagementRequest(t, PreviewModelPricingConversion, http.MethodPost, "/api/option/model_pricing/convert", fixture.Conversion, &response)
|
||||
require.True(t, response.Success)
|
||||
assert.Equal(t, fixture.Expression, response.Data.Expression)
|
||||
})
|
||||
}
|
||||
// The preview must resolve the draft, even when the running process
|
||||
// still has a different saved output multiplier for that model.
|
||||
require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"conversion-defaults":99}`))
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
draft model.PricingValues
|
||||
expression, reason string
|
||||
}{
|
||||
{"conversion-defaults", model.PricingValues{"ModelRatio": float64(2)}, `tier("base", p * 4 + c * 4)`, ""},
|
||||
{"gpt-4-32k", model.PricingValues{"ModelRatio": 30.0, "CacheRatio": 1.0}, `tier("base", p * 60 + c * 120)`, ""},
|
||||
{"conversion-free-cache", model.PricingValues{"ModelRatio": 0.0, "CacheRatio": 0.0}, `tier("base", p * 0 + c * 0 + cr * 0)`, ""},
|
||||
{"claude-3-7-sonnet-20250219", model.PricingValues{"ModelRatio": float64(1.5), "CacheRatio": float64(0.1), "CreateCacheRatio": float64(1.25)}, `tier("base", p * 3 + c * 15 + cr * 0.3 + cc * 3.75 + cc1h * 6)`, ""},
|
||||
{"conversion-image-default", model.PricingValues{"ModelRatio": float64(2), "ImageRatio": float64(1)}, `tier("base", p * 4 + c * 4)`, ""},
|
||||
{"conversion-image-free", model.PricingValues{"ModelRatio": float64(2), "ImageRatio": float64(0)}, `tier("base", p * 4 + c * 4 + cr * 4 + img * 0)`, ""},
|
||||
{"deepseek-chat", model.PricingValues{"ModelRatio": float64(0.135), "CacheRatio": float64(0.25)}, `tier("base", p * 0.27 + c * 0.27 + cr * 0.0675)`, ""},
|
||||
{"gpt-4o-custom", model.PricingValues{"ModelRatio": float64(2)}, `tier("base", p * 4 + c * 16)`, ""},
|
||||
{"gemini-2.5-pro-custom", model.PricingValues{"ModelRatio": float64(2)}, `tier("base", p * 4 + c * 32)`, ""},
|
||||
{"vendor/claude-sonnet-4", model.PricingValues{"ModelRatio": float64(2), "CompletionRatio": float64(0)}, `tier("base", p * 4 + c * 0 + cr * 4 + cc * 5 + cc1h * 8)`, ""},
|
||||
{"conversion-custom", model.PricingValues{"ModelRatio": float64(2), "CompletionRatio": float64(3), "CacheRatio": float64(0), "CreateCacheRatio": float64(1.5), "ImageRatio": float64(2)}, `tier("base", p * 4 + c * 12 + cr * 0 + cc * 6 + img * 8)`, ""},
|
||||
{"deepseek-reasoner", model.PricingValues{"ModelRatio": 0.275, "CacheRatio": 0.25}, `tier("base", p * 0.55 + c * 0.55 + cr * 0.1375)`, ""},
|
||||
{"gpt-5.6-sol", model.PricingValues{"ModelRatio": float64(2), "CompletionRatio": float64(2)}, `tier("base", p * 4 + c * 8)`, ""},
|
||||
{"gpt-5.5", model.PricingValues{"ModelRatio": float64(2), "CompletionRatio": float64(2)}, `tier("base", p * 4 + c * 8)`, ""},
|
||||
{"gpt-6-astra", model.PricingValues{"ModelRatio": float64(2), "CompletionRatio": float64(2)}, `tier("base", p * 4 + c * 8)`, ""},
|
||||
{"conversion-free", model.PricingValues{"ModelPrice": float64(0)}, `tier("request", fixed(0))`, ""},
|
||||
{"conversion-fixed", model.PricingValues{"ModelPrice": float64(0.25), "ModelRatio": float64(7)}, `tier("request", fixed(0.25))`, ""},
|
||||
{"gpt-4o-2024-05-13", model.PricingValues{"ModelRatio": float64(2), "CompletionRatio": float64(99)}, `tier("base", p * 4 + c * 12)`, ""},
|
||||
{"gpt-image-2", model.PricingValues{"ModelPrice": float64(1)}, `tier("image", fixed(1)) * image_count`, ""},
|
||||
{"qwen-image-3.0-pro", model.PricingValues{"ModelPrice": float64(1)}, `tier("image", fixed(1)) * image_count`, ""},
|
||||
{"wan2.7-image-pro", model.PricingValues{"ModelPrice": float64(1)}, "", "Task pricing must be converted manually using the task usage schema."},
|
||||
{"video-summarizer", model.PricingValues{"ModelPrice": float64(1)}, `tier("request", fixed(1))`, ""},
|
||||
{"sora-transcript-helper", model.PricingValues{"ModelPrice": float64(1)}, `tier("request", fixed(1))`, ""},
|
||||
{"gpt-realtime", model.PricingValues{"ModelRatio": float64(1)}, "", "Realtime pricing must be converted manually."},
|
||||
{"conversion-audio", model.PricingValues{"ModelRatio": float64(1), "AudioRatio": float64(2)}, `tier("base", p * 2 + c * 2 + ai * 4 + ao * 4)`, ""},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
before, err := model.GetModelPricingSnapshot([]string{tc.name})
|
||||
require.NoError(t, err)
|
||||
var response struct {
|
||||
Success bool
|
||||
Data model.ModelPricingConversion
|
||||
}
|
||||
modelManagementRequest(t, PreviewModelPricingConversion, http.MethodPost, "/api/option/model_pricing/convert", map[string]any{"model_name": tc.name, "pricing": tc.draft}, &response)
|
||||
require.True(t, response.Success)
|
||||
assert.Equal(t, tc.expression, response.Data.Expression)
|
||||
assert.Equal(t, tc.reason, response.Data.UnsupportedReason)
|
||||
if tc.name == "claude-3-7-sonnet-20250219" {
|
||||
// Both providers include images in input; Claude reports cache
|
||||
// separately and does not report an image token breakdown.
|
||||
for _, usage := range []dto.Usage{
|
||||
{PromptTokens: 1000, CompletionTokens: 100, UsageSemantic: "openai", PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 200, CachedCreationTokens: 50, ImageTokens: 300}},
|
||||
{PromptTokens: 750, CompletionTokens: 100, UsageSemantic: "anthropic", ClaudeCacheCreation5mTokens: 50, PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 200, CachedCreationTokens: 50}},
|
||||
} {
|
||||
tokens := service.BuildTieredTokenParams(&usage, usage.UsageSemantic == "anthropic", billingexpr.UsedVars(response.Data.Expression))
|
||||
cost, _, err := billingexpr.RunExpr(response.Data.Expression, tokens)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, float64(750), tokens.P)
|
||||
assert.Equal(t, 3997.5, cost, "image input must remain billed exactly once")
|
||||
}
|
||||
usage := dto.Usage{PromptTokens: 750, CompletionTokens: 100, UsageSemantic: "anthropic", ClaudeCacheCreation5mTokens: 30, ClaudeCacheCreation1hTokens: 20, PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 200, CachedCreationTokens: 50}}
|
||||
tokens := service.BuildTieredTokenParams(&usage, true, billingexpr.UsedVars(response.Data.Expression))
|
||||
cost, _, err := billingexpr.RunExpr(response.Data.Expression, tokens)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 4042.5, cost, "mixed TTLs preserve both legacy cache-write prices")
|
||||
}
|
||||
if tc.name == "deepseek-chat" {
|
||||
for _, usage := range []dto.Usage{
|
||||
{PromptTokens: 1000, CompletionTokens: 100, PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 200}},
|
||||
{PromptTokens: 800, CompletionTokens: 100, UsageSemantic: "anthropic", PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 200}},
|
||||
} {
|
||||
tokens := service.BuildTieredTokenParams(&usage, usage.UsageSemantic == "anthropic", billingexpr.UsedVars(response.Data.Expression))
|
||||
cost, _, err := billingexpr.RunExpr(response.Data.Expression, tokens)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 256.5, cost, "cache reads remain billed across both response formats")
|
||||
}
|
||||
}
|
||||
after, err := model.GetModelPricingSnapshot([]string{tc.name})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, before.Entries, after.Entries, "preview must not write pricing")
|
||||
if tc.expression == "" {
|
||||
return
|
||||
}
|
||||
pricing := tc.draft
|
||||
pricing["billing_setting.billing_mode"] = "tiered_expr"
|
||||
pricing["billing_setting.billing_expr"] = tc.expression
|
||||
change := model.ModelPricingChange{ModelName: tc.name, ExpectedVersion: before.Entries[0].Version, Pricing: pricing}
|
||||
require.NoError(t, model.UpdateModelPricing([]model.ModelPricingChange{change}))
|
||||
after, err = model.GetModelPricingSnapshot([]string{tc.name})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, pricing, after.Entries[0].Configured)
|
||||
assert.ErrorIs(t, model.UpdateModelPricing([]model.ModelPricingChange{change}), model.ErrModelPricingConflict)
|
||||
change.ExpectedVersion = after.Entries[0].Version
|
||||
change.Pricing["billing_setting.billing_mode"] = "ratio"
|
||||
require.NoError(t, model.UpdateModelPricing([]model.ModelPricingChange{change}))
|
||||
after, err = model.GetModelPricingSnapshot([]string{tc.name})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "ratio", after.Entries[0].Effective["billing_setting.billing_mode"])
|
||||
if _, fixed := pricing["ModelPrice"]; !fixed {
|
||||
var preview struct {
|
||||
Success bool
|
||||
Data model.ModelPricingDescription
|
||||
}
|
||||
modelManagementRequest(t, PreviewModelPricing, http.MethodPost, "/api/option/model_pricing/preview", map[string]any{"model_name": tc.name, "pricing": pricing}, &preview)
|
||||
require.True(t, preview.Success)
|
||||
assert.Equal(t, after.Entries[0].Effective, preview.Data.Effective)
|
||||
require.NotEmpty(t, response.Data.Effective)
|
||||
for field, value := range response.Data.Effective {
|
||||
assert.Equal(t, value, preview.Data.Effective[field], "conversion preview field %s", field)
|
||||
}
|
||||
cache, _ := ratio_setting.GetCacheRatio(tc.name)
|
||||
createCache, _ := ratio_setting.GetCreateCacheRatio(tc.name)
|
||||
image, _ := ratio_setting.GetImageRatio(tc.name)
|
||||
assert.Equal(t, ratio_setting.GetCompletionRatio(tc.name), preview.Data.Effective["CompletionRatio"])
|
||||
assert.Equal(t, cache, preview.Data.Effective["CacheRatio"])
|
||||
assert.Equal(t, createCache, preview.Data.Effective["CreateCacheRatio"])
|
||||
assert.Equal(t, image, preview.Data.Effective["ImageRatio"])
|
||||
}
|
||||
})
|
||||
}
|
||||
require.NoError(t, (&model.Channel{Name: "Mapped image", Models: "conversion-alias", Type: 1, ModelMapping: common.GetPointer(`{"conversion-alias":"conversion-hop","conversion-hop":"gpt-image-2"}`)}).Insert())
|
||||
preview, err := model.PreviewModelPricingConversion("conversion-alias", model.PricingValues{"ModelPrice": float64(1)})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, `tier("image", fixed(1)) * image_count`, preview.Expression)
|
||||
assert.Empty(t, preview.UnsupportedReason)
|
||||
|
||||
t.Run("cache_writes_follow_configured_prices_and_claude_names", func(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
configured *float64
|
||||
mode model.CacheWriteMode
|
||||
}{
|
||||
{"unrecognized-model", nil, model.CacheWriteNone},
|
||||
{"gpt-5.6-terra", nil, model.CacheWriteNone},
|
||||
{"gpt-5.6-terra", common.GetPointer(1.25), model.CacheWriteStandard},
|
||||
{"vendor/claude-sonnet-4-6-high", nil, model.CacheWriteClaudeTTL},
|
||||
{"custom-CLAUDE-alias", common.GetPointer(0.0), model.CacheWriteClaudeTTL},
|
||||
{"deepseek-reasoner", common.GetPointer(0.0), model.CacheWriteStandard},
|
||||
{"custom-cache", common.GetPointer(1.5), model.CacheWriteStandard},
|
||||
} {
|
||||
t.Run(tc.name+"/"+string(tc.mode), func(t *testing.T) {
|
||||
draft := model.PricingValues{"ModelRatio": 2.0, "CompletionRatio": 2.0, "CacheRatio": 0.25, "billing_setting.billing_mode": "ratio"}
|
||||
if tc.configured != nil {
|
||||
draft["CreateCacheRatio"] = *tc.configured
|
||||
}
|
||||
var response struct {
|
||||
Success bool
|
||||
Data model.ModelPricingConversion
|
||||
}
|
||||
modelManagementRequest(t, PreviewModelPricingConversion, http.MethodPost, "/api/option/model_pricing/convert", map[string]any{"model_name": tc.name, "pricing": draft}, &response)
|
||||
require.True(t, response.Success)
|
||||
require.NotEmpty(t, response.Data.Expression)
|
||||
assert.Equal(t, tc.mode, response.Data.CacheWriteMode)
|
||||
var prices struct {
|
||||
Success bool
|
||||
Data model.ModelPricingDescription
|
||||
}
|
||||
modelManagementRequest(t, PreviewModelPricing, http.MethodPost, "/api/option/model_pricing/preview", map[string]any{"model_name": tc.name, "pricing": draft}, &prices)
|
||||
require.True(t, prices.Success)
|
||||
assert.Equal(t, response.Data.Effective, prices.Data.Effective)
|
||||
assert.Equal(t, response.Data.CacheWriteMode, prices.Data.CacheWriteMode)
|
||||
vars := billingexpr.UsedVars(response.Data.Expression)
|
||||
assert.Equal(t, tc.mode != model.CacheWriteNone, vars["cc"])
|
||||
assert.Equal(t, tc.mode == model.CacheWriteClaudeTTL, vars["cc1h"])
|
||||
if tc.configured != nil && *tc.configured == 0 {
|
||||
cost, _, err := billingexpr.RunExpr(response.Data.Expression, billingexpr.TokenParams{CC: 100, CC1h: 100})
|
||||
require.NoError(t, err)
|
||||
assert.Zero(t, cost, "explicit zero cache prices must stay free")
|
||||
}
|
||||
})
|
||||
}
|
||||
// The billing name remains authoritative for cache terms.
|
||||
require.NoError(t, (&model.Channel{Name: "Cache alias", Type: 1, Models: "opaque-cache-alias", ModelMapping: common.GetPointer(`{"opaque-cache-alias":"claude-sonnet-4-6"}`)}).Insert())
|
||||
alias, err := model.PreviewModelPricingConversion("opaque-cache-alias", model.PricingValues{"ModelRatio": 2.0})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, `tier("base", p * 4 + c * 4)`, alias.Expression)
|
||||
})
|
||||
|
||||
t.Run("same_price_cache_reads_preserve_overlapping_categories", func(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
field, term string
|
||||
ratio float64
|
||||
details dto.InputTokenDetails
|
||||
cost, unpricedCacheCost float64
|
||||
}{
|
||||
{"CreateCacheRatio", "cc * 75", 1.25, dto.InputTokenDetails{CachedTokens: 800, CacheWriteTokens: 500}, 97500, 79500},
|
||||
{"ImageRatio", "img * 120", 2, dto.InputTokenDetails{CachedTokens: 600, ImageTokens: 500}, 108000, 102000},
|
||||
} {
|
||||
t.Run(tc.field, func(t *testing.T) {
|
||||
converted, err := model.PreviewModelPricingConversion("overlapping-cache", model.PricingValues{"ModelRatio": 30.0, "CompletionRatio": 2.0, "CacheRatio": 1.0, tc.field: tc.ratio})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, `tier("base", p * 60 + c * 120 + cr * 60 + `+tc.term+`)`, converted.Expression)
|
||||
usage := dto.Usage{PromptTokens: 1000, CompletionTokens: 100, PromptTokensDetails: tc.details}
|
||||
params := service.BuildTieredTokenParams(&usage, false, billingexpr.UsedVars(converted.Expression))
|
||||
cost, _, err := billingexpr.RunExpr(converted.Expression, params)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tc.cost, cost, "preserve the existing overlapping-category charge")
|
||||
unpriced := `tier("base", p * 60 + c * 120 + ` + tc.term + `)`
|
||||
params = service.BuildTieredTokenParams(&usage, false, billingexpr.UsedVars(unpriced))
|
||||
cost, _, err = billingexpr.RunExpr(unpriced, params)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, tc.unpricedCacheCost, cost, "existing expressions without cr retain their OpenAI normalization")
|
||||
})
|
||||
}
|
||||
})
|
||||
t.Run("media_prices_and_request_adjustments", func(t *testing.T) {
|
||||
for _, quantity := range []string{"129", "-1", "0.5", "18446744073709551615"} {
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", strings.NewReader(`{"model":"z-image","parameters":{"n":`+quantity+`}}`))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
Relay(ctx, types.RelayFormatOpenAIImage)
|
||||
assert.Equal(t, http.StatusBadRequest, recorder.Code, recorder.Body.String())
|
||||
}
|
||||
conflicting, err := model.PreviewModelPricingConversion("gemini-2.5-flash-conflict", model.PricingValues{"ModelRatio": 0.15, "AudioRatio": 10.0})
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, conflicting.UnsupportedReason, "Gemini and OpenAI audio prices differ")
|
||||
aliName := "z-image-media-conversion"
|
||||
require.NoError(t, (&model.Channel{Name: aliName, Models: aliName, Type: 17, Status: common.ChannelStatusEnabled}).Insert())
|
||||
aliConversion, err := model.PreviewModelPricingConversion(aliName, model.PricingValues{"ModelPrice": 0.04})
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, aliConversion.UnsupportedReason)
|
||||
quantity := 3
|
||||
aliCost, _, err := billingexpr.RunExprWithRequest(aliConversion.Expression, billingexpr.TokenParams{}, billingexpr.RequestInput{ImageCount: &quantity, Body: []byte(`{"parameters":{"prompt_extend":true}}`)})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 240000.0, aliCost)
|
||||
require.NoError(t, (&model.Channel{Name: "Different image rules", Models: aliName, Type: 1, Status: common.ChannelStatusEnabled}).Insert())
|
||||
conflicting, err = model.PreviewModelPricingConversion(aliName, model.PricingValues{"ModelPrice": 0.04})
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, conflicting.UnsupportedReason, "different image request multipliers")
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
pricing model.PricingValues
|
||||
usage dto.Usage
|
||||
want float64
|
||||
}{
|
||||
{"gemini-2.5-flash", model.PricingValues{"ModelRatio": 0.15, "CompletionRatio": 2.0, "CacheRatio": 0.1}, dto.Usage{PromptTokens: 1000, CompletionTokens: 100, PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 100, AudioTokens: 200}}, 473},
|
||||
{"gemini-2.5-flash-zero", model.PricingValues{"ModelRatio": 0.0}, dto.Usage{PromptTokens: 1000, PromptTokensDetails: dto.InputTokenDetails{AudioTokens: 200}}, 200},
|
||||
{"media-audio", model.PricingValues{"ModelRatio": 1.0, "CompletionRatio": 3.0, "AudioRatio": 2.0, "AudioCompletionRatio": 4.0}, dto.Usage{PromptTokens: 1000, CompletionTokens: 100, PromptTokensDetails: dto.InputTokenDetails{TextTokens: 800, AudioTokens: 200}, CompletionTokenDetails: dto.OutputTokenDetails{TextTokens: 50, AudioTokens: 50}}, 3500},
|
||||
{"media-free-audio", model.PricingValues{"ModelRatio": 1.0, "AudioRatio": 0.0}, dto.Usage{PromptTokens: 1000, CompletionTokens: 100, PromptTokensDetails: dto.InputTokenDetails{AudioTokens: 200}, CompletionTokenDetails: dto.OutputTokenDetails{AudioTokens: 100}}, 1600},
|
||||
{"media-audio-cache", model.PricingValues{"ModelRatio": 1.0, "CompletionRatio": 3.0, "AudioRatio": 2.0, "CacheRatio": 0.1}, dto.Usage{PromptTokens: 1000, CompletionTokens: 100, PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 400, AudioTokens: 200}, CompletionTokenDetails: dto.OutputTokenDetails{AudioTokens: 50}}, 2900},
|
||||
{"media-text-cache", model.PricingValues{"ModelRatio": 1.0, "CompletionRatio": 3.0, "AudioRatio": 2.0, "CacheRatio": 0.1}, dto.Usage{PromptTokens: 1000, CompletionTokens: 100, PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 400}}, 1880},
|
||||
{"gemini-2.5-flash-overlap", model.PricingValues{"ModelRatio": 1.0}, dto.Usage{PromptTokens: 1000, PromptTokensDetails: dto.InputTokenDetails{CachedTokens: 800, AudioTokens: 500}}, 2100},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
converted, err := model.PreviewModelPricingConversion(tc.name, tc.pricing)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, converted.UnsupportedReason)
|
||||
require.NotNil(t, converted.BillingDetails.AudioInputPrice)
|
||||
params := service.BuildTieredTokenParams(&tc.usage, false, billingexpr.UsedVars(converted.Expression))
|
||||
cost, _, err := billingexpr.RunExpr(converted.Expression, params)
|
||||
require.NoError(t, err)
|
||||
assert.InDelta(t, tc.want, cost, 1e-9)
|
||||
})
|
||||
}
|
||||
converted, err := model.PreviewModelPricingConversion("dall-e-3", model.PricingValues{"ModelPrice": 0.04})
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, converted.UnsupportedReason)
|
||||
assert.NotContains(t, converted.Expression, "256x256")
|
||||
assert.NotContains(t, converted.Expression, "512x512")
|
||||
smallImage, err := model.PreviewModelPricingConversion("dall-e-2", model.PricingValues{"ModelPrice": 0.02})
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, smallImage.UnsupportedReason)
|
||||
assert.NotContains(t, smallImage.Expression, "1792")
|
||||
smallCost, _, err := billingexpr.RunExprWithRequest(smallImage.Expression, billingexpr.TokenParams{}, billingexpr.RequestInput{Body: []byte(`{"size":"256x256"}`), ImageCount: &quantity})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, 24000.0, smallCost)
|
||||
for _, tc := range []struct {
|
||||
size, quality string
|
||||
count int
|
||||
cost float64
|
||||
}{
|
||||
{"1024x1024", "standard", 3, 120000},
|
||||
{"1024x1024", "hd", 2, 160000},
|
||||
{"1024x1792", "hd", 3, 360000},
|
||||
} {
|
||||
body, err := common.Marshal(map[string]any{"size": tc.size, "quality": tc.quality})
|
||||
require.NoError(t, err)
|
||||
cost, _, err := billingexpr.RunExprWithRequest(converted.Expression, billingexpr.TokenParams{}, billingexpr.RequestInput{Body: body, ImageCount: &tc.count})
|
||||
require.NoError(t, err)
|
||||
assert.InDelta(t, tc.cost, cost, 1e-9)
|
||||
}
|
||||
for _, channelType := range []int{20, 58} {
|
||||
name := fmt.Sprintf("gemini-3.5-flash-media-%d", channelType)
|
||||
require.NoError(t, (&model.Channel{Name: name, Models: name, Type: channelType}).Insert())
|
||||
require.NoError(t, db.Create(&model.Model{ModelName: name, Endpoints: `{"custom-text":"/v1/custom"}`}).Error)
|
||||
converted, err := model.PreviewModelPricingConversion(name, model.PricingValues{"ModelRatio": 0.75, "CompletionRatio": 6.0, "CacheRatio": 0.1})
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, converted.UnsupportedReason)
|
||||
assert.Equal(t, `tier("base", p * 1.5 + c * 9 + cr * 0.15)`, converted.Expression)
|
||||
assert.Nil(t, converted.BillingDetails.AudioInputPrice)
|
||||
}
|
||||
})
|
||||
require.NoError(t, db.Create(&model.Model{ModelName: "conversion-metadata", Endpoints: `{"image-generation":"/v1/images/generations"}`}).Error)
|
||||
preview, err = model.PreviewModelPricingConversion("conversion-metadata", model.PricingValues{"ModelPrice": float64(1)})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, `tier("image", fixed(1)) * image_count`, preview.Expression)
|
||||
assert.Empty(t, preview.UnsupportedReason)
|
||||
for _, tc := range []struct {
|
||||
name, rule string
|
||||
nameRule int
|
||||
endpoint constant.EndpointType
|
||||
}{
|
||||
{"catalog-image-prefix-opaque", "catalog-image-prefix-", model.NameRulePrefix, constant.EndpointTypeImageGeneration},
|
||||
{"opaque-catalog-image-suffix", "-catalog-image-suffix", model.NameRuleSuffix, constant.EndpointTypeImageGeneration},
|
||||
{"before-catalog-image-middle-after", "catalog-image-middle", model.NameRuleContains, constant.EndpointTypeImageGeneration},
|
||||
{"opaque-catalog-video", "opaque-catalog-video", model.NameRuleExact, constant.EndpointTypeOpenAIVideo},
|
||||
} {
|
||||
t.Run("endpoint_metadata/"+tc.name, func(t *testing.T) {
|
||||
endpoints, err := common.Marshal(map[string]string{string(tc.endpoint): "/v1/fixture"})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.Create(&model.Model{ModelName: tc.rule, NameRule: tc.nameRule, Endpoints: string(endpoints)}).Error)
|
||||
converted, err := model.PreviewModelPricingConversion(tc.name, model.PricingValues{"ModelPrice": 1.0})
|
||||
require.NoError(t, err)
|
||||
if tc.endpoint == constant.EndpointTypeOpenAIVideo {
|
||||
assert.Empty(t, converted.Expression)
|
||||
assert.Equal(t, "Video pricing must be converted manually.", converted.UnsupportedReason)
|
||||
} else {
|
||||
assert.Equal(t, `tier("image", fixed(1)) * image_count`, converted.Expression)
|
||||
}
|
||||
})
|
||||
}
|
||||
sora := &model.Channel{Name: "Opaque video route", Models: "opaque-video-route", Type: constant.ChannelTypeSora, Group: "a,b"}
|
||||
require.NoError(t, sora.Insert())
|
||||
video, err := model.PreviewModelPricingConversion("opaque-video-route", model.PricingValues{"ModelPrice": 1.0})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Video pricing must be converted manually.", video.UnsupportedReason)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelManagementDatabaseMatrix(t *testing.T) {
|
||||
_, err := jsplugin.DefaultRegistry.Register(`
|
||||
export const meta = {apiVersion: 1, key: "model-management-task", name: "Management task fixture", version: "1.0.0", author: {name: "Test"}, models: ["matrix-task"], fetchMode: "per_task", usageSchema: {seconds: {type: "number", unit: "second"}}};
|
||||
@@ -1134,3 +1502,168 @@ func TestModelDeletionDatabaseMatrix(t *testing.T) {
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedModelPluginPricingDatabaseMatrix(t *testing.T) {
|
||||
const name = "shared-model::priced"
|
||||
const base = `tier("base", u("seconds") * 0.4)`
|
||||
const variant = `tier("beta", u("credits") * 2)`
|
||||
for _, dialect := range []struct{ kind, env string }{{"sqlite", ""}, {"mysql", "TEST_MYSQL_DSN"}, {"postgres", "TEST_POSTGRES_DSN"}} {
|
||||
t.Run(dialect.kind, func(t *testing.T) {
|
||||
if dialect.env != "" && os.Getenv(dialect.env) == "" {
|
||||
t.Skip("set " + dialect.env + " to run this database")
|
||||
}
|
||||
db := modelManagementDB(t, dialect.kind, os.Getenv(dialect.env))
|
||||
for _, spec := range []struct{ key, field, unit string }{{"matrix-alpha", "seconds", "second"}, {"matrix-beta", "credits", "credit"}} {
|
||||
source := fmt.Sprintf(`
|
||||
export const meta = {apiVersion:1,key:%q,name:%q,version:"1.0.0",author:{name:"Test"},models:[%q],fetchMode:"per_task",usageSchema:{%s:{type:"number",unit:%q}}};
|
||||
export function buildSubmitRequest(){return {};}
|
||||
export function parseSubmitResponse(){return {};}
|
||||
export function buildQueryRequest(){return {};}
|
||||
export function parseTaskResult(){return {};}
|
||||
`, spec.key, spec.key, name, spec.field, spec.unit)
|
||||
_, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(spec.key) })
|
||||
}
|
||||
|
||||
// Representative existing options have no plugin-expression row.
|
||||
baseJSON, err := common.Marshal(map[string]string{name: base})
|
||||
require.NoError(t, err)
|
||||
modeJSON, err := common.Marshal(map[string]string{name: "tiered_expr"})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.Create(&[]model.Option{
|
||||
{Key: "ModelPrice", Value: `{"unrelated-model":0.75}`},
|
||||
{Key: "billing_setting.billing_expr", Value: string(baseJSON)},
|
||||
{Key: "billing_setting.billing_mode", Value: string(modeJSON)},
|
||||
}).Error)
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{
|
||||
"billing_setting.billing_expr": string(baseJSON), "billing_setting.billing_mode": string(modeJSON),
|
||||
}))
|
||||
before, err := model.GetModelPricingSnapshot([]string{name})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, before.Entries, 1)
|
||||
require.Len(t, before.Entries[0].PluginVariants, 2)
|
||||
assert.True(t, before.Entries[0].PluginVariants[0].Compatible)
|
||||
assert.False(t, before.Entries[0].PluginVariants[1].Compatible)
|
||||
assert.Equal(t, base, before.Entries[0].PluginVariants[1].Effective)
|
||||
change := model.ModelPricingChange{ModelName: name, ExpectedVersion: before.Entries[0].Version, Pricing: model.PricingValues{
|
||||
"billing_setting.billing_mode": "tiered_expr", "billing_setting.billing_expr": base,
|
||||
billing_setting.PluginBillingExprOption: map[string]any{"matrix-beta": variant},
|
||||
}}
|
||||
require.NoError(t, model.UpdateModelPricing([]model.ModelPricingChange{change}))
|
||||
loaded, err := model.GetModelPricingSnapshot([]string{name})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, change.Pricing, loaded.Entries[0].Configured)
|
||||
require.Len(t, loaded.Entries[0].PluginVariants, 2)
|
||||
assert.True(t, loaded.Entries[0].PluginVariants[0].Compatible)
|
||||
assert.Equal(t, base, loaded.Entries[0].PluginVariants[0].Effective)
|
||||
assert.True(t, loaded.Entries[0].PluginVariants[1].Compatible)
|
||||
assert.Equal(t, variant, loaded.Entries[0].PluginVariants[1].Configured)
|
||||
var flat map[string]string
|
||||
require.NoError(t, common.UnmarshalJsonStr(loaded.Options[billing_setting.PluginBillingExprOption], &flat))
|
||||
assert.Equal(t, map[string]string{"matrix-beta::" + name: variant}, flat)
|
||||
all, err := model.GetModelPricingSnapshot(nil)
|
||||
require.NoError(t, err)
|
||||
for _, entry := range all.Entries {
|
||||
assert.NotEqual(t, "matrix-beta::"+name, entry.ModelName)
|
||||
}
|
||||
unrelated, err := model.GetModelPricingSnapshot([]string{"unrelated-model"})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, float64(0.75), unrelated.Entries[0].Configured["ModelPrice"])
|
||||
assert.NotContains(t, billing_setting.GetPricingSyncData(map[string]any{}), "plugin_billing_expr")
|
||||
assert.NotContains(t, billing_setting.GetPricingSyncData(map[string]any{})["billing_expr"], "matrix-beta::"+name)
|
||||
// Restart/re-read is idempotent; a second save based on its version succeeds.
|
||||
require.NoError(t, db.AutoMigrate(&model.Option{}))
|
||||
change.ExpectedVersion = loaded.Entries[0].Version
|
||||
require.NoError(t, model.UpdateModelPricing([]model.ModelPricingChange{change}))
|
||||
repeated, err := model.GetModelPricingSnapshot([]string{name})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, loaded.Entries[0].Version, repeated.Entries[0].Version)
|
||||
// Changing only the provider price changes the model version.
|
||||
change.Pricing[billing_setting.PluginBillingExprOption] = map[string]any{"matrix-beta": `tier("free", u("credits") * 0)`}
|
||||
require.NoError(t, model.UpdateModelPricing([]model.ModelPricingChange{change}))
|
||||
updated, err := model.GetModelPricingSnapshot([]string{name})
|
||||
require.NoError(t, err)
|
||||
assert.NotEqual(t, loaded.Entries[0].Version, updated.Entries[0].Version)
|
||||
assert.ErrorIs(t, model.UpdateModelPricing([]model.ModelPricingChange{change}), model.ErrModelPricingConflict)
|
||||
// The legacy option API validates the full draft and cannot drop a required override.
|
||||
response := modelManagementRequest(t, UpdateOption, http.MethodPut, "/api/option/", OptionUpdateRequest{Key: billing_setting.PluginBillingExprOption, Value: `{}`}, nil)
|
||||
assert.Contains(t, response.Body.String(), `"success":false`)
|
||||
assert.Contains(t, response.Body.String(), "matrix-beta")
|
||||
afterFailure, err := model.GetModelPricingSnapshot([]string{name})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, updated.Entries[0].Version, afterFailure.Entries[0].Version)
|
||||
// Model-level legacy saves also skip providers with a stored override.
|
||||
response = modelManagementRequest(t, UpdateOption, http.MethodPut, "/api/option/", OptionUpdateRequest{Key: "billing_setting.billing_expr", Value: string(baseJSON)}, nil)
|
||||
assert.Contains(t, response.Body.String(), `"success":true`)
|
||||
flat["matrix-beta::"+name] = variant
|
||||
raw, err := common.Marshal(flat)
|
||||
require.NoError(t, err)
|
||||
response = modelManagementRequest(t, UpdateOption, http.MethodPut, "/api/option/", OptionUpdateRequest{Key: billing_setting.PluginBillingExprOption, Value: string(raw)}, nil)
|
||||
assert.Contains(t, response.Body.String(), `"success":true`)
|
||||
final, err := model.GetModelPricingSnapshot([]string{name})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, variant, final.Entries[0].PluginVariants[1].Effective)
|
||||
// A provider may disappear without making every legacy price save fail.
|
||||
require.NoError(t, jsplugin.DefaultRegistry.Unregister("matrix-beta"))
|
||||
stale, err := model.GetModelPricingSnapshot([]string{name})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, stale.Entries[0].PluginVariants, 2)
|
||||
assert.True(t, stale.Entries[0].PluginVariants[1].Stale)
|
||||
assert.Equal(t, variant, stale.Entries[0].PluginVariants[1].Configured)
|
||||
assert.Empty(t, stale.Entries[0].PluginVariants[1].Effective)
|
||||
_, err = model.PreviewModelPricing(name, stale.Entries[0].Configured)
|
||||
require.NoError(t, err)
|
||||
ratioJSON, err := common.Marshal(map[string]float64{name: 2})
|
||||
require.NoError(t, err)
|
||||
response = modelManagementRequest(t, UpdateOption, http.MethodPut, "/api/option/", OptionUpdateRequest{Key: "ModelRatio", Value: string(ratioJSON)}, nil)
|
||||
assert.Contains(t, response.Body.String(), `"success":true`)
|
||||
response = modelManagementRequest(t, UpdateOption, http.MethodPut, "/api/option/", OptionUpdateRequest{Key: "billing_setting.billing_expr", Value: string(baseJSON)}, nil)
|
||||
assert.Contains(t, response.Body.String(), `"success":true`)
|
||||
final, err = model.GetModelPricingSnapshot([]string{name})
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, variant, final.Entries[0].PluginVariants[1].Configured)
|
||||
// The versioned save also preserves an unchanged stale expression.
|
||||
change = model.ModelPricingChange{ModelName: name, ExpectedVersion: final.Entries[0].Version, Pricing: final.Entries[0].Configured}
|
||||
require.NoError(t, model.UpdateModelPricing([]model.ModelPricingChange{change}))
|
||||
change.Pricing[billing_setting.PluginBillingExprOption] = map[string]any{"matrix-beta": "1"}
|
||||
// Even if this process's cache is ahead of storage, a changed stale
|
||||
// override is validated against the locked database value.
|
||||
tamperedJSON, err := common.Marshal(map[string]string{"matrix-beta::" + name: "1"})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{billing_setting.PluginBillingExprOption: string(tamperedJSON)}))
|
||||
require.ErrorContains(t, model.UpdateModelPricing([]model.ModelPricingChange{change}), "does not declare this model")
|
||||
require.NoError(t, config.GlobalConfig.LoadFromDB(map[string]string{billing_setting.PluginBillingExprOption: final.Options[billing_setting.PluginBillingExprOption]}))
|
||||
// Updating a plugin to stop declaring this model also leaves the
|
||||
// stored override visible and inert, even with no active providers.
|
||||
_, err = jsplugin.DefaultRegistry.Register(`
|
||||
export const meta = {apiVersion:1,key:"matrix-beta",name:"Beta updated",version:"2.0.0",author:{name:"Test"},models:["replacement-model"],fetchMode:"per_task",usageSchema:{credits:{type:"number",unit:"credit"}}};
|
||||
export function buildSubmitRequest(){return {};}
|
||||
export function parseSubmitResponse(){return {};}
|
||||
export function buildQueryRequest(){return {};}
|
||||
export function parseTaskResult(){return {};}
|
||||
`, jsplugin.Options{})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, jsplugin.DefaultRegistry.Unregister("matrix-alpha"))
|
||||
stale, err = model.GetModelPricingSnapshot([]string{name})
|
||||
require.NoError(t, err)
|
||||
require.Len(t, stale.Entries[0].PluginVariants, 1)
|
||||
assert.True(t, stale.Entries[0].PluginVariants[0].Stale)
|
||||
assert.Equal(t, "Beta updated", stale.Entries[0].PluginVariants[0].PluginName)
|
||||
require.NoError(t, model.UpdateModelPricingOptions(map[string]string{"ModelRatio": string(ratioJSON)}))
|
||||
// Removing just the stale override succeeds and leaves other prices.
|
||||
response = modelManagementRequest(t, UpdateOption, http.MethodPut, "/api/option/", OptionUpdateRequest{Key: billing_setting.PluginBillingExprOption, Value: `{}`}, nil)
|
||||
assert.Contains(t, response.Body.String(), `"success":true`)
|
||||
final, err = model.GetModelPricingSnapshot([]string{name})
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, final.Entries[0].PluginVariants)
|
||||
assert.Equal(t, base, final.Entries[0].Configured["billing_setting.billing_expr"])
|
||||
// A complete reset removes all provider keys, including models containing ::.
|
||||
require.NoError(t, model.UpdateModelPricing([]model.ModelPricingChange{{ModelName: name, ExpectedVersion: final.Entries[0].Version, Reset: true}}))
|
||||
final, err = model.GetModelPricingSnapshot([]string{name})
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, final.Entries[0].Configured)
|
||||
assert.Equal(t, "{}", final.Options[billing_setting.PluginBillingExprOption])
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -18,6 +18,44 @@ func GetModelPricingConfig(c *gin.Context) {
|
||||
common.ApiSuccess(c, snapshot)
|
||||
}
|
||||
|
||||
func PreviewModelPricingConversion(c *gin.Context) {
|
||||
var request struct {
|
||||
ModelName string `json:"model_name"`
|
||||
Pricing model.PricingValues `json:"pricing"`
|
||||
}
|
||||
if err := common.DecodeJson(c.Request.Body, &request); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
|
||||
return
|
||||
}
|
||||
preview, err := model.PreviewModelPricingConversion(request.ModelName, request.Pricing)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, preview)
|
||||
}
|
||||
|
||||
func PreviewModelPricing(c *gin.Context) {
|
||||
var request struct {
|
||||
ModelName string `json:"model_name"`
|
||||
Pricing model.PricingValues `json:"pricing"`
|
||||
}
|
||||
if err := common.DecodeJson(c.Request.Body, &request); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
|
||||
return
|
||||
}
|
||||
preview, err := model.PreviewModelPricing(request.ModelName, request.Pricing)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, model.ModelPricingDescription{
|
||||
Effective: preview,
|
||||
CacheWriteMode: model.ResolveCacheWriteMode(request.ModelName, request.Pricing),
|
||||
BillingDetails: model.ResolveLegacyBillingDetails(request.ModelName, preview, request.Pricing),
|
||||
})
|
||||
}
|
||||
|
||||
func UpdateModelPricingConfig(c *gin.Context) {
|
||||
var request struct {
|
||||
Changes []model.ModelPricingChange `json:"changes"`
|
||||
|
||||
+89
-34
@@ -5,6 +5,7 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -292,7 +293,7 @@ func HandleOAuth(c *gin.Context) {
|
||||
|
||||
switch flow.Intent {
|
||||
case model.AuthFlowIntentLogin:
|
||||
handleOAuthLogin(c, provider, oauthUser, flow)
|
||||
handleOAuthLogin(c, provider, oauthUser, token, flow)
|
||||
case model.AuthFlowIntentVerify:
|
||||
handleOAuthVerification(c, providerName, oauthUser, flow)
|
||||
}
|
||||
@@ -314,14 +315,14 @@ func handleOAuthVerification(c *gin.Context, provider string, oauthUser *oauth.O
|
||||
common.ApiSuccess(c, proof)
|
||||
}
|
||||
|
||||
func handleOAuthLogin(c *gin.Context, provider oauth.Provider, oauthUser *oauth.OAuthUser, flow *model.AuthFlow) {
|
||||
func handleOAuthLogin(c *gin.Context, provider oauth.Provider, oauthUser *oauth.OAuthUser, token *oauth.OAuthToken, flow *model.AuthFlow) {
|
||||
// 7. Find or create user
|
||||
var payload oauthFlowPayload
|
||||
if err := common.UnmarshalJsonStr(flow.Payload, &payload); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
user, err := findOrCreateOAuthUser(c, provider, oauthUser, payload.AffiliateCode)
|
||||
user, migration, err := findOrCreateOAuthUser(c, provider, oauthUser, token, payload.AffiliateCode)
|
||||
if err != nil {
|
||||
if errors.Is(err, model.ErrEmailAlreadyTaken) {
|
||||
common.ApiErrorI18n(c, i18n.MsgUserEmailAlreadyTaken)
|
||||
@@ -334,6 +335,8 @@ func handleOAuthLogin(c *gin.Context, provider oauth.Provider, oauthUser *oauth.
|
||||
common.ApiErrorI18n(c, i18n.MsgUserRegisterDisabled)
|
||||
case *OAuthEmailAlreadyTakenError:
|
||||
common.ApiErrorI18n(c, i18n.MsgUserEmailAlreadyTaken)
|
||||
case *OAuthLegacyBindingNotConfirmedError:
|
||||
common.ApiErrorI18n(c, i18n.MsgOAuthNotAutoLinked, providerParams(provider.GetName()))
|
||||
default:
|
||||
writeSecurityOperationError(c, err)
|
||||
}
|
||||
@@ -347,7 +350,7 @@ func handleOAuthLogin(c *gin.Context, provider oauth.Provider, oauthUser *oauth.
|
||||
}
|
||||
|
||||
// 9. Setup login
|
||||
setupLogin(user, c)
|
||||
setupLogin(user, migration, c)
|
||||
}
|
||||
|
||||
// handleOAuthBind handles binding OAuth account to existing user
|
||||
@@ -377,10 +380,6 @@ func handleOAuthBind(c *gin.Context, providerName string, provider oauth.Provide
|
||||
common.ApiErrorI18n(c, i18n.MsgOAuthAlreadyBound, providerParams(provider.GetName()))
|
||||
return false, false
|
||||
}
|
||||
if legacyID, ok := oauthUser.Extra["legacy_id"].(string); ok && legacyID != "" && provider.IsUserIDTaken(legacyID) {
|
||||
common.ApiErrorI18n(c, i18n.MsgOAuthAlreadyBound, providerParams(provider.GetName()))
|
||||
return false, false
|
||||
}
|
||||
_, err = model.ConsumeAuthFlowWithAction(state, match, func(tx *gorm.DB, _ *model.AuthFlow) error {
|
||||
if providerName == "telegram" {
|
||||
return model.BindTelegramForSessionWithTx(tx, identity, oauthUser.ProviderUserID)
|
||||
@@ -404,53 +403,91 @@ func handleOAuthBind(c *gin.Context, providerName string, provider oauth.Provide
|
||||
return true, notificationFailed
|
||||
}
|
||||
|
||||
// findOrCreateOAuthUser finds existing user or creates new user
|
||||
func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *oauth.OAuthUser, affiliateCode string) (*model.User, error) {
|
||||
// findOrCreateOAuthUser finds the existing user or creates a new one. For a
|
||||
// legacy GitHub binding that still waits for the login verification, it also
|
||||
// returns the rewrite to carry into the challenge.
|
||||
func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *oauth.OAuthUser, token *oauth.OAuthToken, affiliateCode string) (*model.User, *service.LegacyGitHubMigration, error) {
|
||||
user := &model.User{}
|
||||
if provider.ProviderUserIDColumn() == "telegram_id" {
|
||||
err := provider.FillUserByProviderID(user, oauthUser.ProviderUserID)
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return nil, oauth.ErrTelegramAccountNotBound
|
||||
return nil, nil, oauth.ErrTelegramAccountNotBound
|
||||
}
|
||||
return user, err
|
||||
return user, nil, err
|
||||
}
|
||||
|
||||
// Check if user already exists with new ID
|
||||
if provider.IsUserIDTaken(oauthUser.ProviderUserID) {
|
||||
err := provider.FillUserByProviderID(user, oauthUser.ProviderUserID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return nil, nil, err
|
||||
}
|
||||
// Check if user has been deleted
|
||||
if user.Id == 0 {
|
||||
return nil, &OAuthUserDeletedError{}
|
||||
return nil, nil, &OAuthUserDeletedError{}
|
||||
}
|
||||
return user, nil
|
||||
return user, nil, nil
|
||||
}
|
||||
|
||||
// Try to find user with legacy ID (for GitHub migration from login to numeric ID)
|
||||
if legacyID, ok := oauthUser.Extra["legacy_id"].(string); ok && legacyID != "" {
|
||||
if provider.IsUserIDTaken(legacyID) {
|
||||
err := provider.FillUserByProviderID(user, legacyID)
|
||||
// Legacy GitHub bindings stored the login name, which only points at a
|
||||
// candidate account. Values that are all digits are already numeric account
|
||||
// IDs and never take part in the comparison.
|
||||
legacyID, _ := oauthUser.Extra["legacy_id"].(string)
|
||||
if strings.ContainsFunc(legacyID, func(r rune) bool { return r < '0' || r > '9' }) && provider.IsUserIDTaken(legacyID) {
|
||||
if err := provider.FillUserByProviderID(user, legacyID); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if user.Id != 0 {
|
||||
state, err := model.GetUserVerificationState(user.Id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return nil, nil, err
|
||||
}
|
||||
if user.Id != 0 {
|
||||
// Found user with legacy ID, migrate to new ID
|
||||
common.SysLog(fmt.Sprintf("[OAuth] Migrating user %d from legacy_id=%s to new_id=%s",
|
||||
user.Id, legacyID, oauthUser.ProviderUserID))
|
||||
if err := user.UpdateGitHubId(oauthUser.ProviderUserID); err != nil {
|
||||
common.SysError(fmt.Sprintf("[OAuth] Failed to migrate user %d: %s", user.Id, err.Error()))
|
||||
// Continue with login even if migration fails
|
||||
if state.HasTwoFA || state.HasPasskey {
|
||||
// The rewrite is written in the transaction that issues the session
|
||||
// once the login verification completes.
|
||||
return user, &service.LegacyGitHubMigration{GitHubID: oauthUser.ProviderUserID, LegacyID: legacyID}, nil
|
||||
}
|
||||
// Without a second factor, one of the addresses the provider has
|
||||
// confirmed must match the account email. The list is fetched only here
|
||||
// and never recorded.
|
||||
reason, matched := "no_matching_evidence", false
|
||||
if emailProvider, ok := provider.(oauth.VerifiedEmailProvider); ok && user.Email != "" {
|
||||
emails, err := emailProvider.GetVerifiedEmails(c.Request.Context(), token)
|
||||
if err != nil {
|
||||
common.SysError(fmt.Sprintf("[OAuth] Failed to load verified emails for user %d: %s", user.Id, err.Error()))
|
||||
reason = "verified_emails_unavailable"
|
||||
}
|
||||
return user, nil
|
||||
accountEmail := model.NormalizeEmail(user.Email)
|
||||
matched = slices.ContainsFunc(emails, func(email string) bool { return model.NormalizeEmail(email) == accountEmail })
|
||||
}
|
||||
if !matched {
|
||||
recordLegacyGitHubBindingAudit(c, user, false, map[string]any{"legacy_id": legacyID, "provider_user_id": oauthUser.ProviderUserID, "reason": reason})
|
||||
return nil, nil, &OAuthLegacyBindingNotConfirmedError{}
|
||||
}
|
||||
written := false
|
||||
err = model.DB.Transaction(func(tx *gorm.DB) error {
|
||||
var err error
|
||||
written, err = model.MigrateLegacyGitHubBindingWithTx(tx, user.Id, legacyID, oauthUser.ProviderUserID)
|
||||
return err
|
||||
})
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if written {
|
||||
user.GitHubId = oauthUser.ProviderUserID
|
||||
notificationFailed := service.NotifyAccountSecurityChange(user.Email, "Login account linked: "+provider.GetName()) != nil
|
||||
recordLegacyGitHubBindingAudit(c, user, true, map[string]any{
|
||||
"legacy_id": legacyID, "provider_user_id": oauthUser.ProviderUserID,
|
||||
"verified_email_matched": true, "notification_failed": notificationFailed,
|
||||
})
|
||||
}
|
||||
return user, nil, nil
|
||||
}
|
||||
}
|
||||
|
||||
// User doesn't exist, create new user if registration is enabled
|
||||
if !common.RegisterEnabled {
|
||||
return nil, &OAuthRegistrationDisabledError{}
|
||||
return nil, nil, &OAuthRegistrationDisabledError{}
|
||||
}
|
||||
|
||||
// Set up new user
|
||||
@@ -476,9 +513,9 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o
|
||||
user.Email = model.NormalizeEmail(oauthUser.Email)
|
||||
if err := model.EnsureEmailAvailable(user.Email, 0); err != nil {
|
||||
if errors.Is(err, model.ErrEmailAlreadyTaken) {
|
||||
return nil, &OAuthEmailAlreadyTakenError{}
|
||||
return nil, nil, &OAuthEmailAlreadyTakenError{}
|
||||
}
|
||||
return nil, err
|
||||
return nil, nil, err
|
||||
}
|
||||
}
|
||||
user.Role = common.RoleCommonUser
|
||||
@@ -512,7 +549,7 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
// Perform post-transaction tasks (logs, sidebar config, inviter rewards)
|
||||
@@ -541,14 +578,24 @@ func findOrCreateOAuthUser(c *gin.Context, provider oauth.Provider, oauthUser *o
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return nil, nil, err
|
||||
}
|
||||
|
||||
// Perform post-transaction tasks
|
||||
user.FinalizeOAuthUserCreation(inviterId)
|
||||
}
|
||||
|
||||
return user, nil
|
||||
return user, nil, nil
|
||||
}
|
||||
|
||||
// recordLegacyGitHubBindingAudit records the outcome of a legacy GitHub binding
|
||||
// rewrite as an account binding event. No session exists yet on the login path,
|
||||
// so the row carries the user's own role instead of the request context.
|
||||
func recordLegacyGitHubBindingAudit(c *gin.Context, user *model.User, success bool, params map[string]any) {
|
||||
params["provider"], params["legacy_migration"], params["success"] = "github", true, success
|
||||
model.RecordOperationAuditLog(user.Id, user.Role, auditContentEN("user.binding_bind", params), c.ClientIP(), "user.binding_bind", params, nil, &model.AuditRequestInfo{
|
||||
Method: c.Request.Method, Route: c.FullPath(), Path: c.FullPath(), Status: c.Writer.Status(), Success: success,
|
||||
}, c)
|
||||
}
|
||||
|
||||
// Error types for OAuth
|
||||
@@ -570,6 +617,14 @@ func (e *OAuthEmailAlreadyTakenError) Error() string {
|
||||
return "email is already in use"
|
||||
}
|
||||
|
||||
// OAuthLegacyBindingNotConfirmedError reports a legacy GitHub binding match that
|
||||
// neither a second factor nor a confirmed provider email backed.
|
||||
type OAuthLegacyBindingNotConfirmedError struct{}
|
||||
|
||||
func (e *OAuthLegacyBindingNotConfirmedError) Error() string {
|
||||
return "legacy binding was not confirmed"
|
||||
}
|
||||
|
||||
// handleOAuthError handles OAuth errors and returns translated message
|
||||
func handleOAuthError(c *gin.Context, err error) {
|
||||
switch e := err.(type) {
|
||||
|
||||
+89
-13
@@ -1,6 +1,7 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"slices"
|
||||
@@ -11,7 +12,6 @@ import (
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/billing_setting"
|
||||
@@ -136,6 +136,50 @@ type OptionUpdateRequest struct {
|
||||
Value any `json:"value"`
|
||||
}
|
||||
|
||||
func UpdatePasskeyDomains(c *gin.Context) {
|
||||
var request struct {
|
||||
RPID *string `json:"rp_id"`
|
||||
LegacyRPIDs *string `json:"legacy_rp_ids"`
|
||||
Origins *string `json:"origins"`
|
||||
Preview bool `json:"preview"`
|
||||
RemovalConfirmation string `json:"removal_confirmation"`
|
||||
}
|
||||
if err := common.DecodeJson(c.Request.Body, &request); err != nil || request.RPID == nil || request.LegacyRPIDs == nil || request.Origins == nil {
|
||||
common.ApiErrorI18n(c, i18n.MsgInvalidParams)
|
||||
return
|
||||
}
|
||||
change, err := model.UpdatePasskeyDomainOptions(map[string]string{
|
||||
"passkey.rp_id": *request.RPID, "passkey.legacy_rp_ids": *request.LegacyRPIDs, "passkey.origins": *request.Origins,
|
||||
}, request.Preview, request.RemovalConfirmation)
|
||||
if err != nil {
|
||||
writePasskeyDomainSettingsError(c, err)
|
||||
if !request.Preview {
|
||||
recordPasskeyDomainAudit(c, change, request.RemovalConfirmation != "", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if !request.Preview {
|
||||
recordPasskeyDomainAudit(c, change, request.RemovalConfirmation != "", nil)
|
||||
}
|
||||
common.ApiSuccess(c, change)
|
||||
}
|
||||
|
||||
func writePasskeyDomainSettingsError(c *gin.Context, err error) {
|
||||
var removal *model.PasskeyDomainRemovalError
|
||||
if errors.As(err, &removal) {
|
||||
c.JSON(http.StatusConflict, gin.H{
|
||||
"success": false, "code": "PASSKEY_RP_ID_REMOVAL_CONFIRMATION_REQUIRED",
|
||||
"message": i18n.T(c, i18n.MsgPasskeyRPIDRemovalConfirmation), "data": removal.Change,
|
||||
})
|
||||
return
|
||||
}
|
||||
if errors.Is(err, system_setting.ErrPasskeyRPIDInvalid) {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
common.ApiError(c, err)
|
||||
}
|
||||
|
||||
func UpdateOption(c *gin.Context) {
|
||||
var option OptionUpdateRequest
|
||||
err := common.DecodeJson(c.Request.Body, &option)
|
||||
@@ -359,25 +403,42 @@ func UpdateOption(c *gin.Context) {
|
||||
models = append(models, modelName)
|
||||
}
|
||||
sort.Strings(models)
|
||||
generation := jsplugin.DefaultRegistry.Generation()
|
||||
storedVariants := billing_setting.GetPluginBillingExprCopy()
|
||||
for _, modelName := range models {
|
||||
expression := expressions[modelName]
|
||||
if plugin, ok := generation.GetByModel(modelName); ok {
|
||||
err = billing_setting.SmokeTestTaskExpr(expression, plugin.Meta.UsageSchema)
|
||||
} else if target, resolved := model.ResolveTaskModelAlias(generation, modelName); resolved {
|
||||
if plugin, ok := generation.Get(target.PluginKey); ok {
|
||||
err = billing_setting.SmokeTestTaskExpr(expression, plugin.Meta.UsageSchema)
|
||||
} else {
|
||||
err = billing_setting.SmokeTestExpr(expression)
|
||||
variants := make(map[string]any)
|
||||
for key, expression := range storedVariants {
|
||||
if plugin, name, ok := billing_setting.SplitPluginBillingExprKey(key); ok && name == modelName {
|
||||
variants[plugin] = expression
|
||||
}
|
||||
} else {
|
||||
err = billing_setting.SmokeTestExpr(expression)
|
||||
}
|
||||
err = model.ValidateModelPricing(modelName, model.PricingValues{
|
||||
"billing_setting.billing_expr": expressions[modelName],
|
||||
billing_setting.PluginBillingExprOption: variants,
|
||||
})
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, fmt.Sprintf("模型 %s 的计费表达式无效: %v", modelName, err))
|
||||
return
|
||||
}
|
||||
}
|
||||
case billing_setting.PluginBillingExprOption:
|
||||
var expressions map[string]string
|
||||
if err = common.UnmarshalJsonStr(option.Value.(string), &expressions); err != nil || expressions == nil {
|
||||
common.ApiErrorMsg(c, "plugin billing expressions must be a JSON object")
|
||||
return
|
||||
}
|
||||
for key, expression := range expressions {
|
||||
plugin, name, valid := billing_setting.SplitPluginBillingExprKey(key)
|
||||
if !valid {
|
||||
common.ApiErrorMsg(c, "invalid plugin billing expression key: "+key)
|
||||
return
|
||||
}
|
||||
if err = model.ValidateModelPricing(name, model.PricingValues{
|
||||
billing_setting.PluginBillingExprOption: map[string]any{plugin: expression},
|
||||
}); err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
case "console_setting.api_info":
|
||||
err = console_setting.ValidateConsoleSettings(option.Value.(string), "ApiInfo")
|
||||
if err != nil {
|
||||
@@ -415,9 +476,24 @@ func UpdateOption(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
}
|
||||
if model.IsPasskeyDomainOption(option.Key) {
|
||||
change, updateErr := model.UpdatePasskeyDomainOptions(map[string]string{option.Key: option.Value.(string)}, false, "")
|
||||
if updateErr != nil {
|
||||
writePasskeyDomainSettingsError(c, updateErr)
|
||||
recordPasskeyDomainAudit(c, change, false, updateErr)
|
||||
return
|
||||
}
|
||||
recordPasskeyDomainAudit(c, change, false, nil)
|
||||
common.ApiSuccess(c, change)
|
||||
return
|
||||
}
|
||||
err = model.UpdateOption(option.Key, option.Value.(string))
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
if errors.Is(err, system_setting.ErrPasskeyRPIDInvalid) {
|
||||
writeSecurityOperationError(c, err)
|
||||
} else {
|
||||
common.ApiError(c, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
// 出于安全考虑只记录被修改的配置项名称,不记录配置值(可能含密钥等敏感信息)。
|
||||
|
||||
+53
-28
@@ -26,6 +26,7 @@ type passkeyFinishRequest struct {
|
||||
}
|
||||
|
||||
type passkeyVerifyBeginRequest struct {
|
||||
RPID string `json:"rp_id"`
|
||||
Scope string `json:"scope"`
|
||||
Context json.RawMessage `json:"context,omitempty"`
|
||||
}
|
||||
@@ -42,7 +43,7 @@ func parsePasskeyFinishRequest(c *gin.Context) (*passkeyFinishRequest, error) {
|
||||
}
|
||||
|
||||
func PasskeyRegisterBegin(c *gin.Context) {
|
||||
if !system_setting.GetPasskeySettings().Enabled {
|
||||
if !system_setting.PasskeySettingsSnapshot().Enabled {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "管理员未启用 Passkey 登录",
|
||||
@@ -118,7 +119,7 @@ func PasskeyRegisterBegin(c *gin.Context) {
|
||||
}
|
||||
|
||||
func PasskeyRegisterFinish(c *gin.Context) {
|
||||
if !system_setting.GetPasskeySettings().Enabled {
|
||||
if !system_setting.PasskeySettingsSnapshot().Enabled {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "管理员未启用 Passkey 登录",
|
||||
@@ -142,12 +143,6 @@ func PasskeyRegisterFinish(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
credentialRecord, err := model.GetPasskeyByUserID(user.Id)
|
||||
if err != nil && !errors.Is(err, model.ErrPasskeyNotFound) {
|
||||
writeSecurityOperationError(c, err)
|
||||
@@ -171,6 +166,12 @@ func PasskeyRegisterFinish(c *gin.Context) {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
wa, err := passkeysvc.BuildWebAuthnForRPID(c.Request, sessionData.RelyingPartyID)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
c.Set("passkey_rp_id", sessionData.RelyingPartyID)
|
||||
if sessionData.UserVerification != protocol.VerificationRequired {
|
||||
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
|
||||
return
|
||||
@@ -193,6 +194,7 @@ func PasskeyRegisterFinish(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
passkeyCredential.RPID = &sessionData.RelyingPartyID
|
||||
if err := model.RegisterPasskeyForSession(identity, passkeyCredential); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
@@ -281,7 +283,7 @@ func PasskeyStatus(c *gin.Context) {
|
||||
}
|
||||
|
||||
func PasskeyLoginBegin(c *gin.Context) {
|
||||
if !system_setting.GetPasskeySettings().Enabled {
|
||||
if !system_setting.PasskeySettingsSnapshot().Enabled {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "管理员未启用 Passkey 登录",
|
||||
@@ -289,7 +291,16 @@ func PasskeyLoginBegin(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
var request struct {
|
||||
RPID string `json:"rp_id"`
|
||||
}
|
||||
if c.Request.Body != nil && c.Request.ContentLength != 0 {
|
||||
if common.DecodeJson(c.Request.Body, &request) != nil {
|
||||
common.ApiErrorMsg(c, "无效的 Passkey 验证请求")
|
||||
return
|
||||
}
|
||||
}
|
||||
wa, rpIDs, err := passkeysvc.BuildLoginWebAuthn(c.Request, request.RPID, "")
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
@@ -316,6 +327,7 @@ func PasskeyLoginBegin(c *gin.Context) {
|
||||
"message": "",
|
||||
"data": gin.H{
|
||||
"options": assertion,
|
||||
"rp_ids": rpIDs,
|
||||
"flow_token": flowToken,
|
||||
"expires_at": expiresAt,
|
||||
},
|
||||
@@ -323,7 +335,7 @@ func PasskeyLoginBegin(c *gin.Context) {
|
||||
}
|
||||
|
||||
func PasskeyLoginFinish(c *gin.Context) {
|
||||
if !system_setting.GetPasskeySettings().Enabled {
|
||||
if !system_setting.PasskeySettingsSnapshot().Enabled {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "管理员未启用 Passkey 登录",
|
||||
@@ -342,12 +354,6 @@ func PasskeyLoginFinish(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
sessionData, _, err := passkeysvc.PopSessionDataFlow(
|
||||
request.FlowToken,
|
||||
model.AuthFlowPurposePasskeyLogin,
|
||||
@@ -357,6 +363,12 @@ func PasskeyLoginFinish(c *gin.Context) {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
wa, err := passkeysvc.BuildWebAuthnForRPID(c.Request, sessionData.RelyingPartyID)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
c.Set("passkey_rp_id", sessionData.RelyingPartyID)
|
||||
if sessionData.UserVerification != protocol.VerificationRequired {
|
||||
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
|
||||
return
|
||||
@@ -369,6 +381,10 @@ func PasskeyLoginFinish(c *gin.Context) {
|
||||
return nil, fmt.Errorf("未找到 Passkey 凭证: %w", err)
|
||||
}
|
||||
|
||||
if credential.RPID != nil && *credential.RPID != "" && *credential.RPID != sessionData.RelyingPartyID {
|
||||
return nil, service.ErrVerificationFailed
|
||||
}
|
||||
|
||||
// 通过凭证获取用户
|
||||
user := &model.User{Id: credential.UserID}
|
||||
if err := user.FillUserById(); err != nil {
|
||||
@@ -415,7 +431,7 @@ func PasskeyLoginFinish(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if err := model.UpdatePasskeyAssertionState(modelUser.Id, credential, time.Now()); err != nil {
|
||||
if err := model.UpdatePasskeyAssertionState(modelUser.Id, credential, time.Now(), sessionData.RelyingPartyID); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
@@ -474,7 +490,7 @@ func AdminResetPasskey(c *gin.Context) {
|
||||
}
|
||||
|
||||
func PasskeyVerifyBegin(c *gin.Context) {
|
||||
if !system_setting.GetPasskeySettings().Enabled {
|
||||
if !system_setting.PasskeySettingsSnapshot().Enabled {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "管理员未启用 Passkey 登录",
|
||||
@@ -516,7 +532,11 @@ func PasskeyVerifyBegin(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
credentialRPID := ""
|
||||
if credential.RPID != nil {
|
||||
credentialRPID = *credential.RPID
|
||||
}
|
||||
wa, rpIDs, err := passkeysvc.BuildLoginWebAuthn(c.Request, request.RPID, credentialRPID)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
@@ -544,6 +564,7 @@ func PasskeyVerifyBegin(c *gin.Context) {
|
||||
"message": "",
|
||||
"data": gin.H{
|
||||
"options": assertion,
|
||||
"rp_ids": rpIDs,
|
||||
"flow_token": flowToken,
|
||||
"expires_at": expiresAt,
|
||||
},
|
||||
@@ -551,7 +572,7 @@ func PasskeyVerifyBegin(c *gin.Context) {
|
||||
}
|
||||
|
||||
func PasskeyVerifyFinish(c *gin.Context) {
|
||||
if !system_setting.GetPasskeySettings().Enabled {
|
||||
if !system_setting.PasskeySettingsSnapshot().Enabled {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": false,
|
||||
"message": "管理员未启用 Passkey 登录",
|
||||
@@ -576,12 +597,6 @@ func PasskeyVerifyFinish(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
wa, err := passkeysvc.BuildWebAuthn(c.Request)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
credential, err := model.GetPasskeyByUserID(user.Id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
@@ -605,11 +620,21 @@ func PasskeyVerifyFinish(c *gin.Context) {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
wa, err := passkeysvc.BuildWebAuthnForRPID(c.Request, sessionData.RelyingPartyID)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
c.Set("passkey_rp_id", sessionData.RelyingPartyID)
|
||||
if sessionData.UserVerification != protocol.VerificationRequired {
|
||||
writeSecurityOperationError(c, model.ErrAuthFlowInvalid)
|
||||
return
|
||||
}
|
||||
|
||||
if credential.RPID != nil && *credential.RPID != "" && *credential.RPID != sessionData.RelyingPartyID {
|
||||
writeSecurityOperationError(c, service.ErrVerificationFailed)
|
||||
return
|
||||
}
|
||||
waUser := passkeysvc.NewWebAuthnUser(user, credential)
|
||||
validatedCredential, err := wa.ValidateLogin(waUser, *sessionData, parsedCredential)
|
||||
if err != nil {
|
||||
@@ -617,7 +642,7 @@ func PasskeyVerifyFinish(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if err := model.UpdatePasskeyAssertionState(user.Id, validatedCredential, time.Now()); err != nil {
|
||||
if err := model.UpdatePasskeyAssertionState(user.Id, validatedCredential, time.Now(), sessionData.RelyingPartyID); err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1,9 +1,15 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/ecdsa"
|
||||
"crypto/elliptic"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"maps"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
@@ -11,18 +17,743 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
passkeysvc "github.com/QuantumNous/new-api/service/passkey"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/go-webauthn/webauthn/protocol"
|
||||
"github.com/go-webauthn/webauthn/webauthn"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type passkeyTestBody struct {
|
||||
*strings.Reader
|
||||
}
|
||||
|
||||
type passkeyDomainBegin struct {
|
||||
FlowToken string `json:"flow_token"`
|
||||
RPIDs []string `json:"rp_ids"`
|
||||
Options struct {
|
||||
PublicKey struct {
|
||||
Challenge string `json:"challenge"`
|
||||
RPID string `json:"rpId"`
|
||||
RP struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"rp"`
|
||||
} `json:"publicKey"`
|
||||
} `json:"options"`
|
||||
}
|
||||
|
||||
func TestGetStatusDoesNotExposePasskeyOrigins(t *testing.T) {
|
||||
settings := system_setting.GetPasskeySettings()
|
||||
originalSettings := *settings
|
||||
originalOptionMap := common.OptionMap
|
||||
originalServerAddress := system_setting.ServerAddress
|
||||
t.Cleanup(func() {
|
||||
*settings = originalSettings
|
||||
common.OptionMap = originalOptionMap
|
||||
system_setting.ServerAddress = originalServerAddress
|
||||
})
|
||||
common.OptionMap = map[string]string{}
|
||||
system_setting.ServerAddress = "https://www.example.com"
|
||||
*settings = system_setting.PasskeySettings{Enabled: true, RPID: "example.com"}
|
||||
|
||||
for _, origins := range []string{"https://www.example.com,https://private.example.com", "", "[]"} {
|
||||
t.Run("origins="+origins, func(t *testing.T) {
|
||||
settings.Origins = origins
|
||||
response := httptest.NewRecorder()
|
||||
context, _ := gin.CreateTestContext(response)
|
||||
context.Request = httptest.NewRequest(http.MethodGet, "/api/status", nil)
|
||||
|
||||
GetStatus(context)
|
||||
|
||||
require.Equal(t, http.StatusOK, response.Code)
|
||||
var payload struct {
|
||||
Success bool `json:"success"`
|
||||
Data map[string]any `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &payload))
|
||||
require.True(t, payload.Success)
|
||||
assert.NotContains(t, payload.Data, "passkey_origins")
|
||||
assert.NotContains(t, response.Body.String(), "private.example.com")
|
||||
assert.Equal(t, true, payload.Data["passkey_login"])
|
||||
assert.Equal(t, origins, settings.Origins)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func passkeyDomainRequest(t *testing.T, path string, payload any, identity service.AuthIdentity, origin string, handler gin.HandlerFunc) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
body, err := common.Marshal(payload)
|
||||
require.NoError(t, err)
|
||||
return securityEnrollmentRequest(http.MethodPost, path, string(body), "", identity, func(c *gin.Context) {
|
||||
c.Request.Header.Set("Origin", origin)
|
||||
handler(c)
|
||||
})
|
||||
}
|
||||
|
||||
func decodePasskeyDomainBegin(t *testing.T, response *httptest.ResponseRecorder) passkeyDomainBegin {
|
||||
t.Helper()
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
var begin passkeyDomainBegin
|
||||
require.NoError(t, common.Unmarshal(result.Data, &begin))
|
||||
return begin
|
||||
}
|
||||
|
||||
// Reuse the existing software authenticator, replacing the domain-bound bytes
|
||||
// and recomputing the signature exactly as a browser/authenticator would.
|
||||
func passkeyDomainAssertion(t *testing.T, key *ecdsa.PrivateKey, challenge, rpID, origin string, userID int, verified bool) map[string]any {
|
||||
t.Helper()
|
||||
var assertion map[string]any
|
||||
require.NoError(t, common.Unmarshal(securityPasskeyResponse(t, key, challenge, false, 1, verified), &assertion))
|
||||
response := assertion["response"].(map[string]any)
|
||||
clientData, err := common.Marshal(map[string]string{"type": "webauthn.get", "challenge": challenge, "origin": origin})
|
||||
require.NoError(t, err)
|
||||
authData, err := base64.RawURLEncoding.DecodeString(response["authenticatorData"].(string))
|
||||
require.NoError(t, err)
|
||||
rpHash := sha256.Sum256([]byte(rpID))
|
||||
copy(authData, rpHash[:])
|
||||
clientHash := sha256.Sum256(clientData)
|
||||
signedHash := sha256.Sum256(append(append([]byte{}, authData...), clientHash[:]...))
|
||||
signature, err := ecdsa.SignASN1(rand.Reader, key, signedHash[:])
|
||||
require.NoError(t, err)
|
||||
response["clientDataJSON"] = base64.RawURLEncoding.EncodeToString(clientData)
|
||||
response["authenticatorData"] = base64.RawURLEncoding.EncodeToString(authData)
|
||||
response["signature"] = base64.RawURLEncoding.EncodeToString(signature)
|
||||
response["userHandle"] = base64.RawURLEncoding.EncodeToString([]byte(fmt.Sprint(userID)))
|
||||
return assertion
|
||||
}
|
||||
|
||||
func TestPasskeyDomainsPreserveCredentialsAcrossVerificationFlows(t *testing.T) {
|
||||
for _, kind := range []string{"direct", "login factor", "sensitive action", "historical mixed case"} {
|
||||
t.Run(kind, func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
key := newSecurityLoginPasskey(t, user.Id)
|
||||
settings := system_setting.GetPasskeySettings()
|
||||
legacyRPID := "www.example.com"
|
||||
if kind == "historical mixed case" {
|
||||
legacyRPID = "WWW.example.com"
|
||||
}
|
||||
settings.LegacyRPIDs = legacyRPID
|
||||
settings.Origins = "https://example.com,https://www.example.com"
|
||||
beginPath, finishPath := "/api/user/passkey/login/begin", "/api/user/passkey/login/finish"
|
||||
beginHandler, finishHandler := PasskeyLoginBegin, PasskeyLoginFinish
|
||||
request := map[string]any{"rp_id": legacyRPID}
|
||||
if kind == "login factor" {
|
||||
pending, err := service.StartLoginVerification(user, "password", nil)
|
||||
require.NoError(t, err)
|
||||
request["flow_token"] = pending.FlowToken
|
||||
beginPath, finishPath = "/api/user/login/passkey/begin", "/api/user/login/passkey/finish"
|
||||
beginHandler, finishHandler = LoginPasskeyBegin, LoginPasskeyFinish
|
||||
} else if kind == "sensitive action" {
|
||||
request["scope"] = service.VerificationScopeAccessTokenGenerate
|
||||
beginPath, finishPath = "/api/user/passkey/verify/begin", "/api/user/passkey/verify/finish"
|
||||
beginHandler, finishHandler = PasskeyVerifyBegin, PasskeyVerifyFinish
|
||||
}
|
||||
begin := decodePasskeyDomainBegin(t, passkeyDomainRequest(t, beginPath, request, identity, "https://www.example.com", beginHandler))
|
||||
assert.Equal(t, legacyRPID, begin.Options.PublicKey.RPID)
|
||||
assert.Equal(t, []string{"example.com", legacyRPID}, begin.RPIDs)
|
||||
before, err := model.GetPasskeyByUserID(user.Id)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, before.RPID)
|
||||
finish := map[string]any{"flow_token": begin.FlowToken, "credential": passkeyDomainAssertion(t, key, begin.Options.PublicKey.Challenge, legacyRPID, "https://www.example.com", user.Id, true)}
|
||||
if kind == "login factor" {
|
||||
finish["flow_token"], finish["passkey_flow_token"] = request["flow_token"], begin.FlowToken
|
||||
}
|
||||
response := passkeyDomainRequest(t, finishPath, finish, identity, "https://www.example.com", finishHandler)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
after, err := model.GetPasskeyByUserID(user.Id)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, after.RPID)
|
||||
assert.Equal(t, legacyRPID, *after.RPID)
|
||||
assert.Equal(t, before.CredentialID, after.CredentialID)
|
||||
assert.Equal(t, before.PublicKey, after.PublicKey)
|
||||
if kind == "sensitive action" {
|
||||
var proof service.SecurityProof
|
||||
require.NoError(t, common.Unmarshal(result.Data, &proof))
|
||||
assert.Equal(t, service.VerificationScopeAccessTokenGenerate, proof.Scope)
|
||||
}
|
||||
count, err := model.CountActiveUserSessions(user.Id, time.Now().Unix())
|
||||
require.NoError(t, err)
|
||||
replay := passkeyDomainRequest(t, finishPath, finish, identity, "https://www.example.com", finishHandler)
|
||||
require.NoError(t, common.Unmarshal(replay.Body.Bytes(), &result))
|
||||
assert.False(t, result.Success)
|
||||
afterCount, err := model.CountActiveUserSessions(user.Id, time.Now().Unix())
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, count, afterCount)
|
||||
// Once identified, a browser hint cannot override the credential's binding.
|
||||
wa, ids, err := passkeysvc.BuildLoginWebAuthn(httptest.NewRequest(http.MethodPost, beginPath, nil), "example.com", *after.RPID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, legacyRPID, wa.Config.RPID)
|
||||
assert.Equal(t, []string{legacyRPID}, ids)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPasskeyDomainsRejectInvalidAssertions(t *testing.T) {
|
||||
for _, failure := range []string{"wrong RP ID", "unlisted origin", "origin outside RP scope", "wrong user", "bad signature", "missing user verification", "expired challenge", "removed domain", "known different domain"} {
|
||||
t.Run(failure, func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
key := newSecurityLoginPasskey(t, user.Id)
|
||||
settings := system_setting.GetPasskeySettings()
|
||||
settings.LegacyRPIDs = "www.example.com"
|
||||
settings.Origins = "https://example.com,https://www.example.com"
|
||||
begin := decodePasskeyDomainBegin(t, passkeyDomainRequest(t, "/api/user/passkey/login/begin", map[string]string{"rp_id": "www.example.com"}, identity, "https://www.example.com", PasskeyLoginBegin))
|
||||
rpID, origin, handle := "www.example.com", "https://www.example.com", user.Id
|
||||
switch failure {
|
||||
case "wrong RP ID":
|
||||
rpID = "example.com"
|
||||
case "unlisted origin":
|
||||
origin = "https://evil.example.com"
|
||||
case "origin outside RP scope":
|
||||
origin = "https://example.com"
|
||||
case "wrong user":
|
||||
handle++
|
||||
case "expired challenge":
|
||||
flow, err := model.GetAuthFlow(begin.FlowToken, model.AuthFlowMatch{Purpose: model.AuthFlowPurposePasskeyLogin})
|
||||
require.NoError(t, err)
|
||||
var payload struct {
|
||||
SessionData webauthn.SessionData `json:"session_data"`
|
||||
Security passkeysvc.FlowSecurity `json:"security"`
|
||||
}
|
||||
require.NoError(t, common.UnmarshalJsonStr(flow.Payload, &payload))
|
||||
payload.SessionData.Expires = time.Now().Add(-time.Second)
|
||||
encoded, err := common.Marshal(payload)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, model.DB.Model(flow).Update("payload", string(encoded)).Error)
|
||||
case "removed domain":
|
||||
settings.LegacyRPIDs = ""
|
||||
case "known different domain":
|
||||
require.NoError(t, model.DB.Model(&model.PasskeyCredential{}).Where("user_id = ?", user.Id).Update("rp_id", "example.com").Error)
|
||||
}
|
||||
assertion := passkeyDomainAssertion(t, key, begin.Options.PublicKey.Challenge, rpID, origin, handle, failure != "missing user verification")
|
||||
if failure == "bad signature" {
|
||||
assertion["response"].(map[string]any)["signature"] = "AA"
|
||||
}
|
||||
response := passkeyDomainRequest(t, "/api/user/passkey/login/finish", map[string]any{"flow_token": begin.FlowToken, "credential": assertion}, identity, "https://www.example.com", PasskeyLoginFinish)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.False(t, result.Success, response.Body.String())
|
||||
assert.Empty(t, response.Header().Values("Set-Cookie"))
|
||||
count, err := model.CountActiveUserSessions(user.Id, time.Now().Unix())
|
||||
require.NoError(t, err)
|
||||
assert.EqualValues(t, 1, count)
|
||||
stored, err := model.GetPasskeyByUserID(user.Id)
|
||||
require.NoError(t, err)
|
||||
if failure == "known different domain" {
|
||||
assert.Equal(t, "example.com", *stored.RPID)
|
||||
} else {
|
||||
assert.Nil(t, stored.RPID)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPasskeyDomainChoicesRespectOriginAndConfiguration(t *testing.T) {
|
||||
_, identity := setupSecurityEnrollmentTest(t)
|
||||
settings := system_setting.GetPasskeySettings()
|
||||
settings.LegacyRPIDs = "www.example.com"
|
||||
settings.Origins = "https://example.com,https://www.example.com"
|
||||
for _, test := range []struct {
|
||||
origin, hint string
|
||||
success bool
|
||||
ids []string
|
||||
}{
|
||||
{"https://www.example.com", "", true, []string{"example.com", "www.example.com"}},
|
||||
{"https://example.com", "", true, []string{"example.com"}},
|
||||
{"https://example.com", "www.example.com", false, nil},
|
||||
{"https://www.example.com", "unconfigured.example.com", false, nil},
|
||||
{"https://attacker.test", "", false, nil},
|
||||
} {
|
||||
response := passkeyDomainRequest(t, "/api/user/passkey/login/begin", map[string]string{"rp_id": test.hint}, identity, test.origin, PasskeyLoginBegin)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.Equal(t, test.success, result.Success, "%+v", test)
|
||||
if test.success {
|
||||
assert.Equal(t, test.ids, decodePasskeyDomainBegin(t, response).RPIDs)
|
||||
} else {
|
||||
assert.Equal(t, "PASSKEY_RP_ID_UNAVAILABLE", result.Code)
|
||||
}
|
||||
}
|
||||
settings.RPID, settings.LegacyRPIDs, settings.Origins, settings.AllowInsecureOrigin = "localhost", "", "http://localhost:3000,http://localhost:3001", true
|
||||
for _, origin := range []string{"http://localhost:3000", "http://localhost:3001"} {
|
||||
begin := decodePasskeyDomainBegin(t, passkeyDomainRequest(t, "/api/user/passkey/login/begin", nil, identity, origin, PasskeyLoginBegin))
|
||||
assert.Equal(t, "localhost", begin.Options.PublicKey.RPID)
|
||||
}
|
||||
}
|
||||
|
||||
func setupPasskeyDomainOptions(t *testing.T) {
|
||||
t.Helper()
|
||||
require.NoError(t, model.DB.AutoMigrate(&model.Option{}))
|
||||
common.OptionMapRWMutex.RLock()
|
||||
previousOptions := maps.Clone(common.OptionMap)
|
||||
previousAddress := system_setting.ServerAddress
|
||||
common.OptionMapRWMutex.RUnlock()
|
||||
t.Cleanup(func() {
|
||||
common.OptionMapRWMutex.Lock()
|
||||
common.OptionMap = previousOptions
|
||||
system_setting.ServerAddress = previousAddress
|
||||
common.OptionMapRWMutex.Unlock()
|
||||
})
|
||||
}
|
||||
|
||||
func TestPasskeyDomainRemovalProtectsExistingAndUnknownCredentials(t *testing.T) {
|
||||
for _, binding := range []string{"www.example.com", ""} {
|
||||
t.Run("binding="+binding, func(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
setupPasskeyDomainOptions(t)
|
||||
settings := system_setting.GetPasskeySettings()
|
||||
settings.LegacyRPIDs = "www.example.com"
|
||||
settings.Origins = "https://example.com,https://www.example.com"
|
||||
credential := &model.PasskeyCredential{UserID: user.Id, CredentialID: "existing", PublicKey: "key"}
|
||||
if binding != "" {
|
||||
credential.RPID = &binding
|
||||
}
|
||||
require.NoError(t, model.DB.Create(credential).Error)
|
||||
assert.Error(t, model.UpdateOption("passkey.legacy_rp_ids", ""))
|
||||
assert.Equal(t, "www.example.com", system_setting.PasskeySettingsSnapshot().LegacyRPIDs)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPasskeyDomainRemovalWithoutAffectedCredentialsNeedsNoOverride(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
setupPasskeyDomainOptions(t)
|
||||
settings := system_setting.GetPasskeySettings()
|
||||
settings.LegacyRPIDs = "www.example.com,WWW.example.com"
|
||||
upper := "WWW.example.com"
|
||||
require.NoError(t, model.DB.Create(&model.PasskeyCredential{UserID: user.Id, RPID: &upper, CredentialID: "uppercase", PublicKey: "key"}).Error)
|
||||
change, err := model.UpdatePasskeyDomainOptions(map[string]string{"passkey.legacy_rp_ids": upper}, false, "")
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, []string{"www.example.com"}, change.RemovedRPIDs)
|
||||
assert.Zero(t, change.AffectedCredentials)
|
||||
assert.Zero(t, change.UnknownCredentials)
|
||||
assert.False(t, change.ConfirmationRequired)
|
||||
assert.Equal(t, upper, settings.LegacyRPIDs)
|
||||
}
|
||||
|
||||
func TestPasskeySnapshotDoesNotMaterializeDefaults(t *testing.T) {
|
||||
setupSecurityEnrollmentTest(t)
|
||||
setupPasskeyDomainOptions(t)
|
||||
settings := system_setting.GetPasskeySettings()
|
||||
settings.RPID, settings.Origins = "", ""
|
||||
system_setting.ServerAddress = "https://www.example.com"
|
||||
assert.Equal(t, "www.example.com", system_setting.PasskeySettingsSnapshot().EffectiveRPID())
|
||||
assert.Empty(t, settings.RPID)
|
||||
assert.Empty(t, settings.Origins)
|
||||
}
|
||||
|
||||
func TestPasskeyDomainPreviewConfirmationAndAudit(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
setupPasskeyDomainOptions(t)
|
||||
settings := system_setting.GetPasskeySettings()
|
||||
settings.LegacyRPIDs = "www.example.com,WWW.example.com"
|
||||
settings.Origins = "https://example.com,https://www.example.com"
|
||||
lower, upper, empty := "www.example.com", "WWW.example.com", ""
|
||||
for i, binding := range []*string{&lower, &upper, nil, &empty, &lower} {
|
||||
credential := &model.PasskeyCredential{UserID: user.Id + i, RPID: binding, CredentialID: fmt.Sprint(i), PublicKey: "key"}
|
||||
require.NoError(t, model.DB.Create(credential).Error)
|
||||
if i == 4 {
|
||||
require.NoError(t, model.DB.Delete(credential).Error)
|
||||
}
|
||||
}
|
||||
request := map[string]any{"rp_id": "example.com", "legacy_rp_ids": upper, "origins": settings.Origins, "preview": true}
|
||||
response := passkeyDomainRequest(t, "/api/option/passkey/domains", request, identity, "https://example.com", UpdatePasskeyDomains)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
var preview model.PasskeyDomainChange
|
||||
require.NoError(t, common.Unmarshal(result.Data, &preview))
|
||||
assert.Equal(t, []string{lower}, preview.RemovedRPIDs)
|
||||
assert.EqualValues(t, 1, preview.AffectedCredentials)
|
||||
assert.EqualValues(t, 2, preview.UnknownCredentials)
|
||||
assert.True(t, preview.ConfirmationRequired)
|
||||
require.NotEmpty(t, preview.RemovalConfirmation)
|
||||
options, err := model.AllOption()
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, options, "preview must not materialize default rows")
|
||||
assert.Equal(t, "www.example.com,WWW.example.com", settings.LegacyRPIDs)
|
||||
assert.ErrorIs(t, model.UpdateOptionsBulk(map[string]string{"passkey.legacy_rp_ids": upper, "Notice": "must roll back"}), model.ErrPasskeyDomainRemovalConfirmation)
|
||||
blocked := passkeyDomainRequest(t, "/api/option/", map[string]string{"key": "passkey.legacy_rp_ids", "value": upper}, identity, "https://example.com", UpdateOption)
|
||||
assert.Equal(t, http.StatusConflict, blocked.Code)
|
||||
assert.Contains(t, blocked.Body.String(), "PASSKEY_RP_ID_REMOVAL_CONFIRMATION_REQUIRED")
|
||||
request["preview"] = false
|
||||
request["removal_confirmation"] = preview.RemovalConfirmation
|
||||
// A new unknown credential makes the reviewed impact stale.
|
||||
require.NoError(t, model.DB.Create(&model.PasskeyCredential{UserID: user.Id + 10, CredentialID: "new", PublicKey: "key"}).Error)
|
||||
response = passkeyDomainRequest(t, "/api/option/passkey/domains", request, identity, "https://example.com", UpdatePasskeyDomains)
|
||||
assert.Equal(t, http.StatusConflict, response.Code)
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.NoError(t, common.Unmarshal(result.Data, &preview))
|
||||
assert.EqualValues(t, 3, preview.UnknownCredentials)
|
||||
options, err = model.AllOption()
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, options)
|
||||
request["removal_confirmation"] = preview.RemovalConfirmation
|
||||
response = passkeyDomainRequest(t, "/api/option/passkey/domains", request, identity, "https://example.com", UpdatePasskeyDomains)
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
assert.Equal(t, upper, settings.LegacyRPIDs)
|
||||
var audits []model.AuditLog
|
||||
require.NoError(t, model.LOG_DB.Find(&audits).Error)
|
||||
for _, audit := range audits {
|
||||
if audit.Action == "option.passkey_domains_blocked" {
|
||||
assert.False(t, audit.Success)
|
||||
assert.Equal(t, http.StatusConflict, audit.Status)
|
||||
}
|
||||
if audit.Action == "option.passkey_domains_confirmed" {
|
||||
assert.True(t, audit.Success)
|
||||
require.NotNil(t, audit.Other.Op)
|
||||
params, err := common.Marshal(audit.Other.Op.Params)
|
||||
require.NoError(t, err)
|
||||
var impact struct {
|
||||
Known, Unknown int64
|
||||
Confirmed bool
|
||||
RemovedRPIDs []string `json:"removed_rp_ids"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(params, &impact))
|
||||
assert.EqualValues(t, 1, impact.Known)
|
||||
assert.EqualValues(t, 3, impact.Unknown)
|
||||
assert.True(t, impact.Confirmed)
|
||||
assert.Equal(t, []string{lower}, impact.RemovedRPIDs)
|
||||
}
|
||||
}
|
||||
encoded, err := common.Marshal(audits)
|
||||
require.NoError(t, err)
|
||||
assert.Contains(t, string(encoded), "option.passkey_domains_confirmed")
|
||||
assert.Contains(t, string(encoded), "option.passkey_domains_blocked")
|
||||
assert.NotContains(t, string(encoded), preview.RemovalConfirmation)
|
||||
// Simulate an old node's cache after removal: persistence must use DB trust.
|
||||
settings.LegacyRPIDs = lower + "," + upper
|
||||
assert.ErrorIs(t, model.UpdatePasskeyAssertionState(user.Id, &webauthn.Credential{ID: []byte("0")}, time.Now(), lower), system_setting.ErrPasskeyRPIDUnavailable)
|
||||
assert.ErrorIs(t, model.RegisterPasskeyForSession(identity, &model.PasskeyCredential{UserID: user.Id, RPID: &lower, CredentialID: "replacement", PublicKey: "key"}), system_setting.ErrPasskeyRPIDUnavailable)
|
||||
}
|
||||
|
||||
func TestPasskeyDomainSettingsAcceptExactHTTPSInternalHosts(t *testing.T) {
|
||||
for _, test := range []struct {
|
||||
rpID, origins string
|
||||
valid bool
|
||||
}{
|
||||
{"intranet", "https://intranet:8443", true},
|
||||
{"intranet", "https://child.intranet", false},
|
||||
{"intranet", "http://intranet", false},
|
||||
{"com", "https://com", false},
|
||||
{"127.0.0.1", "https://127.0.0.1", false},
|
||||
{"intranet:8443", "https://intranet:8443", false},
|
||||
{"localhost", "http://localhost:3000", true},
|
||||
} {
|
||||
t.Run(test.rpID+"/"+test.origins, func(t *testing.T) {
|
||||
setupSecurityEnrollmentTest(t)
|
||||
setupPasskeyDomainOptions(t)
|
||||
change, err := model.UpdatePasskeyDomainOptions(map[string]string{"passkey.rp_id": test.rpID, "passkey.legacy_rp_ids": "", "passkey.origins": test.origins}, false, "")
|
||||
if !test.valid {
|
||||
assert.ErrorIs(t, err, system_setting.ErrPasskeyRPIDInvalid)
|
||||
return
|
||||
}
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, test.rpID, change.RPID)
|
||||
assert.Equal(t, "example.com", change.LegacyRPIDs)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPasskeyImplicitDomainChangesRetainPreviousTrust(t *testing.T) {
|
||||
setupSecurityEnrollmentTest(t)
|
||||
setupPasskeyDomainOptions(t)
|
||||
settings := system_setting.GetPasskeySettings()
|
||||
settings.RPID, settings.Origins = "", ""
|
||||
system_setting.ServerAddress = "https://www.example.com"
|
||||
require.NoError(t, model.UpdateOption("ServerAddress", "https://example.com"))
|
||||
assert.Equal(t, []string{"example.com", "www.example.com"}, system_setting.PasskeySettingsSnapshot().RelyingPartyIDs())
|
||||
assert.Equal(t, "https://www.example.com,https://example.com", settings.Origins)
|
||||
// Moving the retained RP into the primary is not a removal.
|
||||
change, err := model.UpdatePasskeyDomainOptions(map[string]string{"passkey.rp_id": "www.example.com", "passkey.legacy_rp_ids": ""}, false, "")
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, change.RemovedRPIDs)
|
||||
assert.Equal(t, "example.com", change.LegacyRPIDs)
|
||||
}
|
||||
|
||||
func TestPasskeyRegistrationAndDomainRemovalSerialize(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
setupPasskeyDomainOptions(t)
|
||||
system_setting.GetPasskeySettings().Origins = "https://example.com,https://www.example.com"
|
||||
require.NoError(t, model.UpdateOption("passkey.legacy_rp_ids", "www.example.com"))
|
||||
registering, release := make(chan struct{}), make(chan struct{})
|
||||
require.NoError(t, model.DB.Callback().Create().Before("gorm:create").Register("passkey_registration_barrier", func(tx *gorm.DB) {
|
||||
if _, ok := tx.Statement.Dest.(*model.PasskeyCredential); ok {
|
||||
close(registering)
|
||||
<-release
|
||||
}
|
||||
}))
|
||||
t.Cleanup(func() { _ = model.DB.Callback().Create().Remove("passkey_registration_barrier") })
|
||||
rpID := "www.example.com"
|
||||
registered, removed := make(chan error, 1), make(chan error, 1)
|
||||
go func() {
|
||||
registered <- model.RegisterPasskeyForSession(identity, &model.PasskeyCredential{UserID: user.Id, RPID: &rpID, CredentialID: "concurrent", PublicKey: "key"})
|
||||
}()
|
||||
<-registering
|
||||
removing := make(chan struct{})
|
||||
go func() {
|
||||
close(removing)
|
||||
removed <- model.UpdateOption("passkey.legacy_rp_ids", "")
|
||||
}()
|
||||
<-removing
|
||||
// A concurrent snapshot must remain read-only while the write is pending.
|
||||
assert.Equal(t, []string{"example.com", "www.example.com"}, system_setting.PasskeySettingsSnapshot().RelyingPartyIDs())
|
||||
close(release)
|
||||
require.NoError(t, <-registered)
|
||||
assert.ErrorIs(t, <-removed, model.ErrPasskeyDomainRemovalConfirmation)
|
||||
stored, err := model.GetPasskeyByUserID(user.Id)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, stored.RPID)
|
||||
assert.Equal(t, rpID, *stored.RPID)
|
||||
}
|
||||
|
||||
func TestPasskeyDomainEndpointRequiresRoot(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
setupPasskeyDomainOptions(t)
|
||||
bundle, err := service.CreateLoginSession(user.Id, "password", "127.0.0.1", "domain-settings-test")
|
||||
require.NoError(t, err)
|
||||
router := gin.New()
|
||||
router.PUT("/api/option/passkey/domains", middleware.RootAuth(), UpdatePasskeyDomains)
|
||||
for _, bearer := range []string{"", bundle.AccessToken} {
|
||||
request := httptest.NewRequest(http.MethodPut, "/api/option/passkey/domains", strings.NewReader(`{"rp_id":"example.com","legacy_rp_ids":"","origins":"https://example.com","preview":true}`))
|
||||
if bearer != "" {
|
||||
request.Header.Set("Authorization", "Bearer "+bearer)
|
||||
}
|
||||
response := httptest.NewRecorder()
|
||||
router.ServeHTTP(response, request)
|
||||
assert.Contains(t, []int{http.StatusUnauthorized, http.StatusForbidden}, response.Code)
|
||||
}
|
||||
options, err := model.AllOption()
|
||||
require.NoError(t, err)
|
||||
assert.Empty(t, options)
|
||||
}
|
||||
|
||||
func TestPasskeyDefaultsDoNotRestrictUnrelatedServerAddressSetup(t *testing.T) {
|
||||
setupSecurityEnrollmentTest(t)
|
||||
setupPasskeyDomainOptions(t)
|
||||
settings := system_setting.GetPasskeySettings()
|
||||
settings.Enabled, settings.RPID, settings.LegacyRPIDs, settings.Origins = false, "", "", ""
|
||||
system_setting.ServerAddress = ""
|
||||
require.NoError(t, model.UpdateOption("ServerAddress", "http://127.0.0.1:3000"))
|
||||
assert.Equal(t, "http://127.0.0.1:3000", system_setting.ServerAddress)
|
||||
assert.Empty(t, settings.RPID)
|
||||
assert.Empty(t, settings.LegacyRPIDs)
|
||||
}
|
||||
|
||||
func TestPasskeyRegistrationRejectsUnavailableRequestOrigin(t *testing.T) {
|
||||
_, identity := setupSecurityEnrollmentTest(t)
|
||||
settings := system_setting.GetPasskeySettings()
|
||||
settings.RPID = "www.example.com"
|
||||
settings.LegacyRPIDs = "example.com"
|
||||
settings.Origins = "https://example.com,https://www.example.com"
|
||||
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: "passkey.register"}, "password")
|
||||
response := securityEnrollmentRequest(http.MethodPost, "/api/user/passkey/register/begin", "", proof, identity, func(c *gin.Context) {
|
||||
c.Request.Header.Set("Origin", "https://example.com")
|
||||
PasskeyRegisterBegin(c)
|
||||
})
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.False(t, result.Success)
|
||||
assert.Equal(t, "PASSKEY_RP_ID_UNAVAILABLE", result.Code)
|
||||
var count int64
|
||||
require.NoError(t, model.DB.Model(&model.AuthFlow{}).Where("purpose = ?", model.AuthFlowPurposePasskeyRegister).Count(&count).Error)
|
||||
assert.Zero(t, count)
|
||||
}
|
||||
|
||||
func TestPasskeyRPIDRotationRetainsDefaultsAndRollsBackOnFailure(t *testing.T) {
|
||||
setupSecurityEnrollmentTest(t)
|
||||
setupPasskeyDomainOptions(t)
|
||||
settings := system_setting.GetPasskeySettings()
|
||||
settings.RPID, settings.LegacyRPIDs = "", ""
|
||||
settings.Origins = "https://www.example.com,https://example.com,https://login.example.com"
|
||||
system_setting.ServerAddress = "https://www.example.com"
|
||||
require.NoError(t, model.DB.Create(&model.Option{Key: "passkey.rp_id", Value: ""}).Error)
|
||||
require.NoError(t, model.UpdateOption("passkey.rp_id", "example.com"))
|
||||
assert.Equal(t, []string{"example.com", "www.example.com"}, system_setting.PasskeySettingsSnapshot().RelyingPartyIDs())
|
||||
// Simulate a second node with an older local cache. Persisted domains win.
|
||||
system_setting.GetPasskeySettings().RPID = "www.example.com"
|
||||
require.NoError(t, model.UpdateOption("passkey.rp_id", "login.example.com"))
|
||||
assert.Equal(t, []string{"login.example.com", "www.example.com", "example.com"}, system_setting.PasskeySettingsSnapshot().RelyingPartyIDs())
|
||||
before, err := model.AllOption()
|
||||
require.NoError(t, err)
|
||||
snapshot := system_setting.PasskeySettingsSnapshot()
|
||||
failure := errors.New("synthetic option write failure")
|
||||
require.NoError(t, model.DB.Callback().Update().Before("gorm:update").Register("passkey_option_failure", func(tx *gorm.DB) {
|
||||
if option, ok := tx.Statement.Dest.(*model.Option); ok && option.Key == "passkey.legacy_rp_ids" {
|
||||
tx.AddError(failure)
|
||||
}
|
||||
}))
|
||||
assert.ErrorIs(t, model.UpdateOptionsBulk(map[string]string{"passkey.rp_id": "next.example.com", "passkey.legacy_rp_ids": "www.example.com"}), failure)
|
||||
require.NoError(t, model.DB.Callback().Update().Remove("passkey_option_failure"))
|
||||
after, err := model.AllOption()
|
||||
require.NoError(t, err)
|
||||
assert.ElementsMatch(t, before, after)
|
||||
assert.Equal(t, snapshot, system_setting.PasskeySettingsSnapshot())
|
||||
for _, invalid := range []string{"*.example.com", "https://example.com", "example.com:443", "example.com/path", "com"} {
|
||||
assert.ErrorIs(t, model.UpdateOption("passkey.rp_id", invalid), system_setting.ErrPasskeyRPIDInvalid)
|
||||
assert.ErrorIs(t, model.UpdateOption("passkey.legacy_rp_ids", invalid), system_setting.ErrPasskeyRPIDInvalid)
|
||||
}
|
||||
assert.Equal(t, snapshot, system_setting.PasskeySettingsSnapshot())
|
||||
// Older settings accepted mixed case. Rotation must preserve its hash input.
|
||||
require.NoError(t, model.DB.Model(&model.Option{}).Where(&model.Option{Key: "passkey.rp_id"}).Update("value", "WWW.example.com").Error)
|
||||
require.NoError(t, model.UpdateOption("passkey.rp_id", "next.example.com"))
|
||||
assert.Contains(t, system_setting.PasskeySettingsSnapshot().RelyingPartyIDs(), "WWW.example.com")
|
||||
}
|
||||
|
||||
func TestPasskeyRPIDRotationKeepsInFlightRegistration(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
setupPasskeyDomainOptions(t)
|
||||
system_setting.GetPasskeySettings().Origins = "https://example.com,https://www.example.com"
|
||||
proof := issueSecurityEnrollmentProof(t, identity, service.VerificationOperation{Scope: "passkey.register"}, "password")
|
||||
begin := decodePasskeyDomainBegin(t, securityEnrollmentRequest(http.MethodPost, "/api/user/passkey/register/begin", "", proof, identity, PasskeyRegisterBegin))
|
||||
assert.Equal(t, "example.com", begin.Options.PublicKey.RP.ID)
|
||||
require.NoError(t, model.UpdateOption("passkey.rp_id", "www.example.com"))
|
||||
key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
|
||||
require.NoError(t, err)
|
||||
response := passkeyDomainRequest(t, "/api/user/passkey/register/finish", passkeyFinishRequest{FlowToken: begin.FlowToken, Credential: securityPasskeyResponse(t, key, begin.Options.PublicKey.Challenge, true, 0)}, identity, "https://example.com", PasskeyRegisterFinish)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
require.True(t, result.Success, response.Body.String())
|
||||
credential, err := model.GetPasskeyByUserID(user.Id)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, credential.RPID)
|
||||
assert.Equal(t, "example.com", *credential.RPID)
|
||||
// A later registration uses the current primary, not the retained domain.
|
||||
wa, err := passkeysvc.BuildWebAuthn(httptest.NewRequest(http.MethodPost, "https://www.example.com/api/user/passkey/register/begin", nil))
|
||||
require.NoError(t, err)
|
||||
options, _, err := wa.BeginRegistration(passkeysvc.NewWebAuthnUser(user, credential))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "www.example.com", options.Response.RelyingParty.ID)
|
||||
}
|
||||
|
||||
func TestPasskeyDomainFailuresDoNotLogCredentials(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
var output bytes.Buffer
|
||||
common.LogWriterMu.Lock()
|
||||
previousWriter := gin.DefaultErrorWriter
|
||||
gin.DefaultErrorWriter = &output
|
||||
common.LogWriterMu.Unlock()
|
||||
t.Cleanup(func() {
|
||||
common.LogWriterMu.Lock()
|
||||
gin.DefaultErrorWriter = previousWriter
|
||||
common.LogWriterMu.Unlock()
|
||||
})
|
||||
response := securityEnrollmentRequest(http.MethodPost, "/api/user/passkey/login/finish", "", "", service.AuthIdentity{}, func(c *gin.Context) {
|
||||
c.Set("passkey_rp_id", "www.example.com")
|
||||
writeSecurityOperationError(c, protocol.ErrVerification.WithDetails("private-flow-token").WithInfo("private-challenge"))
|
||||
})
|
||||
assert.Contains(t, response.Body.String(), "SECURITY_VERIFICATION_FAILED")
|
||||
assert.Contains(t, output.String(), "verification_error")
|
||||
assert.Contains(t, output.String(), "www.example.com")
|
||||
for _, secret := range []string{"private-flow-token", "private-challenge"} {
|
||||
assert.NotContains(t, response.Body.String(), secret)
|
||||
assert.NotContains(t, output.String(), secret)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPasskeyDomainErrorsRespectRequestLanguage(t *testing.T) {
|
||||
user, _ := setupSecurityEnrollmentTest(t)
|
||||
setupPasskeyDomainOptions(t)
|
||||
system_setting.GetPasskeySettings().LegacyRPIDs = "www.example.com"
|
||||
require.NoError(t, model.DB.Create(&model.PasskeyCredential{UserID: user.Id, CredentialID: "unknown-domain", PublicKey: "key"}).Error)
|
||||
for _, locale := range []struct {
|
||||
language, invalid, unavailable, removal string
|
||||
}{
|
||||
{"zh-CN", "通行密钥域名无效。请填写域名,不包含协议、端口、路径或通配符。", "此通行密钥域名无法在当前网站使用。请前往原网站或选择其他验证方式。", "请核对受影响的通行密钥并确认删除域名。配置或影响范围发生变化后,需要重新确认。"},
|
||||
{"zh-TW", "通行金鑰網域無效。請填寫網域,不包含通訊協定、連接埠、路徑或萬用字元。", "此通行金鑰網域無法在目前網站使用。請前往原網站或選擇其他驗證方式。", "請核對受影響的通行金鑰並確認刪除網域。設定或影響範圍變更後,需要重新確認。"},
|
||||
{"en", "Invalid Passkey domain. Enter a domain without a scheme, port, path or wildcard.", "This Passkey domain is not available on this website. Use its original website or another verification method.", "Review the affected Passkeys and confirm the domain removal. If the settings or impact have changed, confirmation is required again."},
|
||||
} {
|
||||
t.Run(locale.language, func(t *testing.T) {
|
||||
for _, request := range []struct {
|
||||
path, body, code, message string
|
||||
handler gin.HandlerFunc
|
||||
}{
|
||||
{"/api/option/", `{"key":"passkey.rp_id","value":"localhost:3000"}`, "PASSKEY_RP_ID_INVALID", locale.invalid, UpdateOption},
|
||||
{"/api/option/", `{"key":"passkey.legacy_rp_ids","value":"localhost:3001"}`, "PASSKEY_RP_ID_INVALID", locale.invalid, UpdateOption},
|
||||
{"/api/option/", `{"key":"passkey.legacy_rp_ids","value":""}`, "PASSKEY_RP_ID_REMOVAL_CONFIRMATION_REQUIRED", locale.removal, UpdateOption},
|
||||
{"/api/user/passkey/login/begin", `{"rp_id":"unconfigured.example.com"}`, "PASSKEY_RP_ID_UNAVAILABLE", locale.unavailable, PasskeyLoginBegin},
|
||||
} {
|
||||
response := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(response)
|
||||
method := http.MethodPost
|
||||
if request.path == "/api/option/" {
|
||||
method = http.MethodPut
|
||||
}
|
||||
c.Request = httptest.NewRequest(method, request.path, strings.NewReader(request.body))
|
||||
c.Request.Header.Set("Accept-Language", locale.language)
|
||||
c.Request.Header.Set("Origin", "https://example.com")
|
||||
c.Set("id", user.Id)
|
||||
c.Set("role", common.RoleRootUser)
|
||||
request.handler(c)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
expectedStatus := http.StatusOK
|
||||
if request.code == "PASSKEY_RP_ID_REMOVAL_CONFIRMATION_REQUIRED" {
|
||||
expectedStatus = http.StatusConflict
|
||||
}
|
||||
assert.Equal(t, expectedStatus, response.Code)
|
||||
assert.False(t, result.Success)
|
||||
assert.Equal(t, request.code, result.Code)
|
||||
assert.Equal(t, request.message, result.Message)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPasskeyRPIDMigrationPreservesExistingCredentials(t *testing.T) {
|
||||
for _, upgrade := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("upgrade=%t", upgrade), func(t *testing.T) {
|
||||
user, identity := setupSecurityEnrollmentTest(t)
|
||||
key := newSecurityLoginPasskey(t, user.Id)
|
||||
before, err := model.GetPasskeyByUserID(user.Id)
|
||||
require.NoError(t, err)
|
||||
if upgrade {
|
||||
// This schema is copied from the latest released model at
|
||||
// v1.0.0-rc.36, with the same GORM and database driver versions.
|
||||
require.NoError(t, model.DB.Migrator().DropTable(&model.PasskeyCredential{}))
|
||||
require.NoError(t, model.DB.AutoMigrate(&passkeyCredentialBeforeRPID{}))
|
||||
encoded, err := common.Marshal(before)
|
||||
require.NoError(t, err)
|
||||
var legacy passkeyCredentialBeforeRPID
|
||||
require.NoError(t, common.Unmarshal(encoded, &legacy))
|
||||
require.NoError(t, model.DB.Create(&legacy).Error)
|
||||
}
|
||||
pool, err := model.DB.DB()
|
||||
require.NoError(t, err)
|
||||
for range 2 {
|
||||
// Each migration represents a new application startup. Discard
|
||||
// PostgreSQL's old SELECT * prepared plans along with its old pool.
|
||||
pool.SetMaxIdleConns(0)
|
||||
pool.SetMaxIdleConns(2)
|
||||
require.NoError(t, model.DB.AutoMigrate(&model.PasskeyCredential{}))
|
||||
}
|
||||
after, err := model.GetPasskeyByUserID(user.Id)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, before, after)
|
||||
assert.True(t, model.DB.Migrator().HasIndex(&model.PasskeyCredential{}, "idx_passkey_credentials_deleted_at"))
|
||||
assert.Error(t, model.DB.Create(&model.PasskeyCredential{UserID: user.Id, CredentialID: "another-key", PublicKey: "key"}).Error)
|
||||
assert.Error(t, model.DB.Create(&model.PasskeyCredential{UserID: user.Id + 1, CredentialID: before.CredentialID, PublicKey: "key"}).Error)
|
||||
system_setting.GetPasskeySettings().LegacyRPIDs = "www.example.com"
|
||||
system_setting.GetPasskeySettings().Origins = "https://example.com,https://www.example.com"
|
||||
begin := decodePasskeyDomainBegin(t, passkeyDomainRequest(t, "/api/user/passkey/login/begin", map[string]string{"rp_id": "www.example.com"}, identity, "https://www.example.com", PasskeyLoginBegin))
|
||||
response := passkeyDomainRequest(t, "/api/user/passkey/login/finish", map[string]any{"flow_token": begin.FlowToken, "credential": passkeyDomainAssertion(t, key, begin.Options.PublicKey.Challenge, "www.example.com", "https://www.example.com", user.Id, true)}, identity, "https://www.example.com", PasskeyLoginFinish)
|
||||
var result securityEnrollmentResponse
|
||||
require.NoError(t, common.Unmarshal(response.Body.Bytes(), &result))
|
||||
assert.True(t, result.Success, response.Body.String())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (*passkeyTestBody) Close() error { return nil }
|
||||
|
||||
func TestParsePasskeyFinishRequestDoesNotRewriteRequestBody(t *testing.T) {
|
||||
@@ -70,3 +801,26 @@ func TestPasskeyRegisterFinishRejectsUnapprovedFlowWithoutConsumingIt(t *testing
|
||||
require.NoError(t, err)
|
||||
assert.Nil(t, flow.ConsumedAt)
|
||||
}
|
||||
|
||||
type passkeyCredentialBeforeRPID struct {
|
||||
ID int `json:"id" gorm:"primaryKey"`
|
||||
UserID int `json:"user_id" gorm:"uniqueIndex;not null"`
|
||||
CredentialID string `json:"credential_id" gorm:"type:varchar(512);uniqueIndex;not null"` // base64 encoded
|
||||
PublicKey string `json:"public_key" gorm:"type:text;not null"` // base64 encoded
|
||||
AttestationType string `json:"attestation_type" gorm:"type:varchar(255)"`
|
||||
AAGUID string `json:"aaguid" gorm:"type:varchar(512)"` // base64 encoded
|
||||
SignCount uint32 `json:"sign_count" gorm:"default:0"`
|
||||
CloneWarning bool `json:"clone_warning"`
|
||||
UserPresent bool `json:"user_present"`
|
||||
UserVerified bool `json:"user_verified"`
|
||||
BackupEligible bool `json:"backup_eligible"`
|
||||
BackupState bool `json:"backup_state"`
|
||||
Transports string `json:"transports" gorm:"type:text"`
|
||||
Attachment string `json:"attachment" gorm:"type:varchar(32)"`
|
||||
LastUsedAt *time.Time `json:"last_used_at"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
|
||||
}
|
||||
|
||||
func (passkeyCredentialBeforeRPID) TableName() string { return "passkey_credentials" }
|
||||
@@ -0,0 +1,96 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/tidwall/gjson"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var errDetectionServiceUnavailable = errors.New("detection service unreachable")
|
||||
|
||||
const localPelicanID = "local-pelican"
|
||||
const pelicanPrompt = "Create a complete self-contained HTML page showing a pelican riding a bicycle. Draw the pelican and bicycle with inline SVG and animate pedaling and rotating wheels using CSS. Make the scene visually polished. Do not load external resources or use JavaScript. Return only the complete HTML document."
|
||||
|
||||
func detectionHTML(raw []byte, contentType string) string {
|
||||
result := gjson.ParseBytes(raw)
|
||||
if strings.Contains(contentType, "text/event-stream") {
|
||||
result = gjson.Result{}
|
||||
for line := range strings.SplitSeq(string(raw), "\n") {
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
event := gjson.Parse(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
|
||||
if event.Get("type").String() == "response.completed" {
|
||||
result = event.Get("response")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if result.Get("status").String() != "completed" {
|
||||
return ""
|
||||
}
|
||||
var output strings.Builder
|
||||
for _, item := range result.Get("output").Array() {
|
||||
if item.Get("type").String() != "message" {
|
||||
continue
|
||||
}
|
||||
for _, part := range item.Get("content").Array() {
|
||||
if part.Get("type").String() == "output_text" {
|
||||
output.WriteString(part.Get("text").String())
|
||||
}
|
||||
}
|
||||
}
|
||||
html := strings.TrimSpace(output.String())
|
||||
lower := strings.ToLower(html)
|
||||
start := strings.Index(lower, "<!doctype html")
|
||||
if start < 0 {
|
||||
start = strings.Index(lower, "<html")
|
||||
}
|
||||
end := strings.LastIndex(lower, "</html>")
|
||||
if start < 0 || end < start || len(html) > 2<<20 {
|
||||
return ""
|
||||
}
|
||||
return html[start : end+7]
|
||||
}
|
||||
|
||||
func processLocalPelican(ctx context.Context, run model.DetectionRun) {
|
||||
if ctx.Err() != nil {
|
||||
return
|
||||
}
|
||||
if model.DB.Model(&model.DetectionRun{}).Where("id = ?", run.ID).Updates(map[string]any{"remote_id": localPelicanID, "quality": "unknown", "result": ""}).Error != nil {
|
||||
return
|
||||
}
|
||||
var call model.DetectionCall
|
||||
if model.DB.First(&call, "id = ?", run.ID).Error != nil {
|
||||
return
|
||||
}
|
||||
if call.State == "running" {
|
||||
return
|
||||
}
|
||||
if call.State == "ready" {
|
||||
payload, _ := common.Marshal(map[string]any{"model": "gpt-6-astra", "input": pelicanPrompt})
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequestWithContext(ctx, http.MethodPost, "/api/detection-proxy/v1/responses", bytes.NewReader(payload))
|
||||
c.Request.Header.Set("Authorization", "Bearer "+pelicanCredential(run.ID))
|
||||
DetectionProxy(c)
|
||||
if model.DB.First(&call, "id = ?", run.ID).Error != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
if call.State != "finished" {
|
||||
return
|
||||
}
|
||||
if model.DB.First(&run, "id = ?", run.ID).Error != nil {
|
||||
return
|
||||
}
|
||||
if run.HTML != "" {
|
||||
finishDetection(run, "succeeded", "")
|
||||
} else {
|
||||
finishDetection(run, "failed", "local_generation_failed")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"github.com/stretchr/testify/assert"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestDetectionHTMLRequiresCompletedDocument(t *testing.T) {
|
||||
raw := []byte(`{"status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"<html><body>pelican</body></html>"}]}]}`)
|
||||
assert.Equal(t, "<html><body>pelican</body></html>", detectionHTML(raw, "application/json"))
|
||||
assert.Empty(t, detectionHTML([]byte(`{"status":"incomplete","output":[]}`), "application/json"))
|
||||
sse := []byte("data: {\"type\":\"response.completed\",\"response\":" + string(raw) + "}\n\n")
|
||||
assert.Equal(t, detectionHTML(raw, "application/json"), detectionHTML(sse, "text/event-stream"))
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/billingexpr"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestLocalPelicanGeneratesOnceWithUnknownVerdict(t *testing.T) {
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
conn, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
conn.SetMaxOpenConns(1)
|
||||
old := model.DB
|
||||
model.DB = db
|
||||
t.Cleanup(func() { model.DB = old; _ = conn.Close() })
|
||||
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.DetectionSchedule{}))
|
||||
calls := 0
|
||||
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls++
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"status":"completed","usage":{"input_tokens":10,"output_tokens":5},"output":[{"type":"message","content":[{"type":"output_text","text":"<html><body>pelican</body></html>"}]}]}`))
|
||||
}))
|
||||
defer upstream.Close()
|
||||
url := upstream.URL
|
||||
channel := model.Channel{Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: "gpt-6-astra"}
|
||||
require.NoError(t, db.Create(&channel).Error)
|
||||
price, _ := common.Marshal(service.DetectionPrice{Snapshot: billingexpr.BillingSnapshot{ExprString: "p+c", ExprHash: billingexpr.ExprHashString("p+c"), GroupRatio: 1, QuotaPerUnit: 500000, ExprVersion: 1}})
|
||||
run := model.DetectionRun{ID: "local-test", RequestKey: "local-test", Kind: "pelican", ChannelID: channel.Id, Status: "running", Deadline: time.Now().Unix() + 60, Reserved: 100, Price: string(price)}
|
||||
require.NoError(t, db.Create(&run).Error)
|
||||
require.NoError(t, db.Create(&model.DetectionCall{ID: run.ID, RunID: run.ID, SecretHash: model.DetectionSecretHash(pelicanCredential(run.ID)), State: "ready", ExpiresAt: run.Deadline}).Error)
|
||||
processLocalPelican(context.Background(), run)
|
||||
require.NoError(t, db.First(&run, "id = ?", run.ID).Error)
|
||||
require.Equal(t, "succeeded", run.Status)
|
||||
require.Equal(t, "unknown", run.Quality)
|
||||
require.Contains(t, run.HTML, "pelican")
|
||||
require.True(t, run.Settled)
|
||||
processLocalPelican(context.Background(), run)
|
||||
require.Equal(t, 1, calls)
|
||||
}
|
||||
@@ -53,9 +53,10 @@ func GetPerfMetrics(c *gin.Context) {
|
||||
}
|
||||
|
||||
result, err := perfmetrics.Query(perfmetrics.QueryParams{
|
||||
Model: modelName,
|
||||
Group: c.Query("group"),
|
||||
Hours: hours,
|
||||
Model: modelName,
|
||||
Group: c.Query("group"),
|
||||
Hours: hours,
|
||||
AllowedGroups: append(lo.Keys(ratio_setting.GetGroupRatioCopy()), "auto"),
|
||||
})
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{
|
||||
@@ -65,18 +66,8 @@ func GetPerfMetrics(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
result.Groups = filterActiveGroups(result.Groups)
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"data": result,
|
||||
})
|
||||
}
|
||||
|
||||
func filterActiveGroups(groups []perfmetrics.GroupResult) []perfmetrics.GroupResult {
|
||||
activeRatios := ratio_setting.GetGroupRatioCopy()
|
||||
return lo.Filter(groups, func(g perfmetrics.GroupResult, _ int) bool {
|
||||
_, ok := activeRatios[g.Group]
|
||||
return ok || g.Group == "auto"
|
||||
})
|
||||
}
|
||||
+190
-198
@@ -25,12 +25,8 @@ import (
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
|
||||
"github.com/bytedance/gopkg/util/gopool"
|
||||
"github.com/samber/lo"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
@@ -77,6 +73,13 @@ func publicUpstreamErrorContent(statusCode int) string {
|
||||
}
|
||||
|
||||
func clientFacingOpenAIError(err *types.NewAPIError, upstreamFailure bool, requestId string) types.OpenAIError {
|
||||
if err.GetErrorCode() == relaycommon.ResponseGuardErrorCode {
|
||||
return types.OpenAIError{
|
||||
Message: common.MessageWithRequestId(err.Error(), requestId),
|
||||
Type: "api_error",
|
||||
Code: relaycommon.ResponseGuardErrorCode,
|
||||
}
|
||||
}
|
||||
if upstreamFailure {
|
||||
return types.OpenAIError{
|
||||
Message: common.MessageWithRequestId(publicUpstreamErrorMessage, requestId),
|
||||
@@ -122,7 +125,11 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
|
||||
|
||||
defer func() {
|
||||
if newAPIError != nil {
|
||||
service.RecordRequestPolicyTermination(c, newAPIError)
|
||||
logger.LogError(c, fmt.Sprintf("relay error: %s", common.LocalLogPreview(newAPIError.Error())))
|
||||
if c.GetBool("response_guard_error_sent") {
|
||||
return
|
||||
}
|
||||
switch relayFormat {
|
||||
case types.RelayFormatOpenAIRealtime:
|
||||
helper.WssError(c, ws, clientFacingOpenAIError(newAPIError, upstreamFailure, requestId))
|
||||
@@ -156,59 +163,25 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
|
||||
return
|
||||
}
|
||||
|
||||
needSensitiveCheck := setting.ShouldCheckPromptSensitive()
|
||||
needCountToken := constant.CountToken
|
||||
// Avoid building huge CombineText (strings.Join) when token counting and sensitive check are both disabled.
|
||||
var meta *types.TokenCountMeta
|
||||
if needSensitiveCheck || needCountToken {
|
||||
meta = request.GetTokenCountMeta()
|
||||
} else {
|
||||
meta = fastTokenCountMetaForPricing(request)
|
||||
}
|
||||
|
||||
if needSensitiveCheck && meta != nil {
|
||||
contains, words := service.CheckSensitiveText(meta.CombineText)
|
||||
if contains {
|
||||
logger.LogWarn(c, fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", ")))
|
||||
newAPIError = types.NewError(err, types.ErrorCodeSensitiveWordsDetected)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
tokens, err := service.EstimateRequestToken(c, meta, relayInfo)
|
||||
if err != nil {
|
||||
newAPIError = types.NewError(err, types.ErrorCodeCountTokenFailed)
|
||||
return
|
||||
}
|
||||
|
||||
relayInfo.SetEstimatePromptTokens(tokens)
|
||||
|
||||
priceData, err := helper.ModelPriceHelper(c, relayInfo, tokens, meta)
|
||||
if err != nil {
|
||||
newAPIError = types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest))
|
||||
return
|
||||
}
|
||||
|
||||
// common.SetContextKey(c, constant.ContextKeyTokenCountMeta, meta)
|
||||
|
||||
if priceData.FreeModel {
|
||||
logger.LogInfo(c, fmt.Sprintf("模型 %s 免费,跳过预扣费", relayInfo.OriginModelName))
|
||||
} else {
|
||||
newAPIError = service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo)
|
||||
if newAPIError != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
defer func() {
|
||||
// Only return quota if downstream failed and quota was actually pre-consumed
|
||||
if newAPIError != nil {
|
||||
newAPIError = service.NormalizeViolationFeeError(newAPIError)
|
||||
if relayInfo.Billing != nil {
|
||||
relayInfo.Billing.Refund(c)
|
||||
}
|
||||
service.ChargeViolationFeeIfNeeded(c, relayInfo, newAPIError)
|
||||
recovered := recover()
|
||||
resultErr := newAPIError
|
||||
if recovered != nil {
|
||||
resultErr = types.NewError(fmt.Errorf("relay panic: %v", recovered), types.ErrorCodeBadResponse)
|
||||
}
|
||||
if relayFormat != types.RelayFormatOpenAIRealtime {
|
||||
perfmetrics.RecordRelayResult(c.Request.Context(), relayInfo, resultErr)
|
||||
}
|
||||
if recovered != nil {
|
||||
panic(recovered)
|
||||
}
|
||||
}()
|
||||
|
||||
if newAPIError = relay.PrepareRequestBilling(c, relayInfo); newAPIError != nil {
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
newAPIError = relay.RefundFailedRequestBilling(c, relayInfo, newAPIError)
|
||||
}()
|
||||
|
||||
retryParam := &service.RetryParam{
|
||||
@@ -220,8 +193,30 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
|
||||
}
|
||||
relayInfo.RetryIndex = 0
|
||||
relayInfo.LastError = nil
|
||||
maxRetries := common.RetryTimes
|
||||
ordinaryRetries := 0
|
||||
if relayInfo.DurableAttemptBilling {
|
||||
maxRetries = 2
|
||||
timeout := constant.StreamingTimeout
|
||||
if timeout <= 0 {
|
||||
timeout = 300
|
||||
}
|
||||
relayInfo.GuardDeadline = time.Now().Add(time.Duration(timeout) * time.Second)
|
||||
defer func() { service.EndFailedRelayAttempt(c, relayInfo) }()
|
||||
}
|
||||
|
||||
for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
|
||||
for ; retryParam.GetRetry() <= maxRetries; retryParam.IncreaseRetry() {
|
||||
if relayInfo.DurableAttemptBilling && time.Now().After(relayInfo.GuardDeadline) {
|
||||
newAPIError = types.NewErrorWithStatusCode(errors.New("response model confirmation budget exhausted"), relaycommon.ResponseGuardErrorCode, http.StatusGatewayTimeout, types.ErrOptionWithSkipRetry())
|
||||
break
|
||||
}
|
||||
if relayInfo.ResponseGuard != nil && relayInfo.UpstreamRequests >= relaycommon.ResponseGuardMaxAttempts {
|
||||
newAPIError = types.NewErrorWithStatusCode(errors.New("response model attempt budget exhausted"), relaycommon.ResponseGuardErrorCode, http.StatusBadGateway, types.ErrOptionWithSkipRetry())
|
||||
break
|
||||
}
|
||||
relayInfo.StreamStatus = nil
|
||||
relayInfo.PerformanceBusinessRejection = false
|
||||
relayInfo.PerformanceOutputTokens = 0
|
||||
relayInfo.RetryIndex = retryParam.GetRetry()
|
||||
channel, channelErr := getChannel(c, relayInfo, retryParam)
|
||||
if channelErr != nil {
|
||||
@@ -229,11 +224,15 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
|
||||
newAPIError = channelErr
|
||||
break
|
||||
}
|
||||
addUsedChannel(c, channel.Id)
|
||||
service.AppendUsedChannel(c, channel.Id)
|
||||
if billingErr := service.PrepareTieredBillingForSelectedGroup(c, relayInfo); billingErr != nil {
|
||||
newAPIError = billingErr
|
||||
break
|
||||
}
|
||||
if billingErr := service.BeginRelayAttempt(c, relayInfo); billingErr != nil {
|
||||
newAPIError = billingErr
|
||||
break
|
||||
}
|
||||
|
||||
bodyStorage, bodyErr := common.GetBodyStorage(c)
|
||||
if bodyErr != nil {
|
||||
@@ -259,19 +258,40 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
|
||||
}
|
||||
|
||||
if newAPIError == nil {
|
||||
service.MarkRequestPolicySuccess(c, relayInfo.StreamStatus)
|
||||
relayInfo.LastError = nil
|
||||
service.RecordGroupOutcome(currentRequestGroup(c, channel), true)
|
||||
return
|
||||
}
|
||||
upstreamFailure = true
|
||||
if isGroupAvailabilityFailure(newAPIError) {
|
||||
service.RecordGroupOutcome(currentRequestGroup(c, channel), false)
|
||||
}
|
||||
|
||||
newAPIError = service.NormalizeViolationFeeError(newAPIError)
|
||||
relayInfo.LastError = newAPIError
|
||||
|
||||
service.EndFailedRelayAttempt(c, relayInfo)
|
||||
remaining := common.RetryTimes - retryParam.GetRetry()
|
||||
if relayInfo.DurableAttemptBilling {
|
||||
remaining = common.RetryTimes - ordinaryRetries
|
||||
}
|
||||
decision := service.DecideRelayRetry(c, newAPIError, remaining)
|
||||
if canSwitchRejectedAstra(c, relayInfo, newAPIError) {
|
||||
decision = service.PolicyDecision{Action: "retry", Reason: "model_mismatch_switch", Source: "model_guard"}
|
||||
service.ClearPoolAffinity(c)
|
||||
} else if decision.Action == "retry" {
|
||||
ordinaryRetries++
|
||||
}
|
||||
service.RecordPolicyFailure(c, channel.Id, newAPIError, decision)
|
||||
processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError, relayInfo)
|
||||
|
||||
if !shouldRetry(c, newAPIError, common.RetryTimes-retryParam.GetRetry()) {
|
||||
if decision.Action != "retry" {
|
||||
break
|
||||
}
|
||||
if service.RequestRoutingPolicy(c) == nil {
|
||||
service.MarkRequestGroupFailed(c, currentRequestGroup(c, channel))
|
||||
}
|
||||
}
|
||||
|
||||
useChannel := c.GetStringSlice("use_channel")
|
||||
@@ -279,11 +299,24 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
|
||||
retryLogStr := fmt.Sprintf("重试:%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
|
||||
logger.LogInfo(c, retryLogStr)
|
||||
}
|
||||
if newAPIError != nil {
|
||||
gopool.Go(func() {
|
||||
perfmetrics.RecordRelaySample(relayInfo, false, 0)
|
||||
})
|
||||
}
|
||||
|
||||
func canSwitchRejectedAstra(c *gin.Context, info *relaycommon.RelayInfo, err *types.NewAPIError) bool {
|
||||
if !info.AstraAutoSwitch || info.ResponseGuard == nil {
|
||||
return false
|
||||
}
|
||||
if !info.DurableAttemptBilling || info.UpstreamRequests >= 3 || err.GetErrorCode() != relaycommon.ResponseGuardErrorCode || c.Writer.Written() || c.Request.Context().Err() != nil || service.RequestRoutingPolicy(c) == nil {
|
||||
return false
|
||||
}
|
||||
if _, pinned, _ := service.GetChannelConstraints(c).ResolvedPin(); pinned {
|
||||
return false
|
||||
}
|
||||
attempt := info.ResponseGuard.Last()
|
||||
if attempt == nil || attempt.Reason != "model_mismatch" {
|
||||
return false
|
||||
}
|
||||
request, ok := info.Request.(*dto.OpenAIResponsesRequest)
|
||||
return ok && request.PreviousResponseID == "" && (len(request.Conversation) == 0 || string(request.Conversation) == "null")
|
||||
}
|
||||
|
||||
// CountClaudeTokens implements Anthropic's token-counting utility endpoint.
|
||||
@@ -319,47 +352,12 @@ func CountClaudeTokens(c *gin.Context) {
|
||||
}
|
||||
|
||||
var upgrader = websocket.Upgrader{
|
||||
Subprotocols: []string{"realtime"}, // WS 握手支持的协议,如果有使用 Sec-WebSocket-Protocol,则必须在此声明对应的 Protocol TODO add other protocol
|
||||
Subprotocols: []string{"realtime", "responses"}, // WS 握手支持的协议,如果有使用 Sec-WebSocket-Protocol,则必须在此声明对应的 Protocol
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
return true // 允许跨域
|
||||
},
|
||||
}
|
||||
|
||||
func addUsedChannel(c *gin.Context, channelId int) {
|
||||
useChannel := c.GetStringSlice("use_channel")
|
||||
useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
|
||||
c.Set("use_channel", useChannel)
|
||||
}
|
||||
|
||||
func fastTokenCountMetaForPricing(request dto.Request) *types.TokenCountMeta {
|
||||
if request == nil {
|
||||
return &types.TokenCountMeta{}
|
||||
}
|
||||
meta := &types.TokenCountMeta{
|
||||
TokenType: types.TokenTypeTokenizer,
|
||||
}
|
||||
switch r := request.(type) {
|
||||
case *dto.GeneralOpenAIRequest:
|
||||
maxCompletionTokens := lo.FromPtrOr(r.MaxCompletionTokens, uint(0))
|
||||
maxTokens := lo.FromPtrOr(r.MaxTokens, uint(0))
|
||||
if maxCompletionTokens > maxTokens {
|
||||
meta.MaxTokens = int(maxCompletionTokens)
|
||||
} else {
|
||||
meta.MaxTokens = int(maxTokens)
|
||||
}
|
||||
case *dto.OpenAIResponsesRequest:
|
||||
meta.MaxTokens = int(lo.FromPtrOr(r.MaxOutputTokens, uint(0)))
|
||||
case *dto.ClaudeRequest:
|
||||
meta.MaxTokens = int(lo.FromPtr(r.MaxTokens))
|
||||
case *dto.ImageRequest:
|
||||
// Pricing for image requests depends on ImagePriceRatio; safe to compute even when CountToken is disabled.
|
||||
return r.GetTokenCountMeta()
|
||||
default:
|
||||
// Best-effort: leave CombineText empty to avoid large allocations.
|
||||
}
|
||||
return meta
|
||||
}
|
||||
|
||||
func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service.RetryParam) (*model.Channel, *types.NewAPIError) {
|
||||
if info.ChannelMeta == nil {
|
||||
autoBan := c.GetBool("auto_ban")
|
||||
@@ -367,12 +365,14 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service
|
||||
if !autoBan {
|
||||
autoBanInt = 0
|
||||
}
|
||||
return &model.Channel{
|
||||
channel := &model.Channel{
|
||||
Id: c.GetInt("channel_id"),
|
||||
Type: c.GetInt("channel_type"),
|
||||
Name: c.GetString("channel_name"),
|
||||
AutoBan: &autoBanInt,
|
||||
}, nil
|
||||
}
|
||||
service.RequestPolicy(c).BeginAttempt(channel, info.UsingGroup)
|
||||
return channel, nil
|
||||
}
|
||||
channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam)
|
||||
if err != nil {
|
||||
@@ -384,6 +384,7 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service
|
||||
|
||||
info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info)
|
||||
|
||||
service.RequestPolicy(c).BeginAttempt(channel, selectGroup)
|
||||
newAPIError := middleware.SetupContextForSelectedChannel(c, channel, info.OriginModelName)
|
||||
if newAPIError != nil {
|
||||
return nil, newAPIError
|
||||
@@ -391,76 +392,45 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service
|
||||
return channel, nil
|
||||
}
|
||||
|
||||
func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool {
|
||||
if openaiErr == nil {
|
||||
// currentRequestGroup returns the group that handled the current attempt,
|
||||
// falling back to the channel's own group when the selection did not record one.
|
||||
func currentRequestGroup(c *gin.Context, channel *model.Channel) string {
|
||||
if group := common.GetContextKeyString(c, constant.ContextKeyAutoGroup); group != "" {
|
||||
return group
|
||||
}
|
||||
if channel != nil {
|
||||
return channel.Group
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// isGroupAvailabilityFailure reports whether an error says something about the
|
||||
// selected group's health rather than about the client request itself.
|
||||
func isGroupAvailabilityFailure(err *types.NewAPIError) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
|
||||
return false
|
||||
}
|
||||
if types.IsChannelError(openaiErr) {
|
||||
if types.IsChannelError(err) {
|
||||
return true
|
||||
}
|
||||
if types.IsSkipRetryError(openaiErr) {
|
||||
return false
|
||||
}
|
||||
if retryTimes <= 0 {
|
||||
return false
|
||||
}
|
||||
if service.GetChannelConstraints(c).SuppressesRetry() {
|
||||
return false
|
||||
}
|
||||
code := openaiErr.StatusCode
|
||||
if code >= 200 && code < 300 {
|
||||
return false
|
||||
}
|
||||
if code < 100 || code > 599 {
|
||||
return true
|
||||
}
|
||||
if operation_setting.IsAlwaysSkipRetryCode(openaiErr.GetErrorCode()) {
|
||||
return false
|
||||
}
|
||||
return operation_setting.ShouldRetryByStatusCode(code)
|
||||
return err.StatusCode >= 500
|
||||
}
|
||||
|
||||
func shouldRetry(c *gin.Context, err *types.NewAPIError, retryTimes int) bool {
|
||||
return service.ShouldRetryRelayError(c, err, retryTimes)
|
||||
}
|
||||
|
||||
func processChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError, relayInfo *relaycommon.RelayInfo) {
|
||||
logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, common.LocalLogPreview(err.Error())))
|
||||
// 不要使用context获取渠道信息,异步处理时可能会出现渠道信息不一致的情况
|
||||
// do not use context to get channel info, there may be inconsistent channel info when processing asynchronously
|
||||
if service.ShouldDisableChannel(err) && channelError.AutoBan {
|
||||
gopool.Go(func() {
|
||||
service.DisableChannel(channelError, err.ErrorWithStatusCode())
|
||||
})
|
||||
}
|
||||
|
||||
if constant.ErrorLogEnabled && types.IsRecordErrorLog(err) {
|
||||
// 保存错误日志到mysql中
|
||||
userId := c.GetInt("id")
|
||||
tokenName := c.GetString("token_name")
|
||||
modelName := c.GetString("original_model")
|
||||
tokenId := c.GetInt("token_id")
|
||||
userGroup := c.GetString("group")
|
||||
other := model.NewLogOther()
|
||||
if c.Request != nil && c.Request.URL != nil {
|
||||
other.SetPublic("request_path", c.Request.URL.Path)
|
||||
}
|
||||
other.SetPublic("error_type", err.GetErrorType())
|
||||
other.SetPublic("error_code", err.GetErrorCode())
|
||||
other.SetPublic("status_code", err.StatusCode)
|
||||
service.AppendRelayLogAdminInfo(c, relayInfo, other)
|
||||
service.AppendTaskPluginContextAuditInfo(c, other)
|
||||
startTime := common.GetContextKeyTime(c, constant.ContextKeyRequestStartTime)
|
||||
if startTime.IsZero() {
|
||||
startTime = time.Now()
|
||||
}
|
||||
useTimeSeconds := int(time.Since(startTime).Seconds())
|
||||
other.SetRoot("upstream_error", err.MaskSensitiveErrorWithStatusCode())
|
||||
model.RecordErrorLog(c, userId, channelError.ChannelId, modelName, tokenName, publicUpstreamErrorContent(err.StatusCode), tokenId, useTimeSeconds, common.GetContextKeyBool(c, constant.ContextKeyIsStream), userGroup, other)
|
||||
}
|
||||
|
||||
service.ProcessChannelError(c, channelError, err, relayInfo)
|
||||
}
|
||||
|
||||
func RelayMidjourney(c *gin.Context) {
|
||||
policy := service.RequestPolicy(c)
|
||||
defer func() {
|
||||
if policy.Attempts > 0 && !policy.Successful {
|
||||
service.RecordRequestPolicyTermination(c, types.NewErrorWithStatusCode(errors.New("Midjourney submission failed"), types.ErrorCodeBadResponseStatusCode, http.StatusBadGateway, types.ErrOptionWithSkipRetry()))
|
||||
}
|
||||
}()
|
||||
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatMjProxy, nil, nil)
|
||||
|
||||
if err != nil {
|
||||
@@ -488,6 +458,7 @@ func RelayMidjourney(c *gin.Context) {
|
||||
//err = relayMidjourneySubmit(c, relayMode)
|
||||
log.Println(mjErr)
|
||||
if mjErr != nil {
|
||||
policy.Successful = false
|
||||
statusCode := http.StatusBadRequest
|
||||
if mjErr.Code == 30 {
|
||||
mjErr.Result = "当前分组负载已饱和,请稍后再试,或升级账户以提升服务质量。"
|
||||
@@ -627,6 +598,7 @@ func executeTaskSubmissionWith(
|
||||
relayInfo *relaycommon.RelayInfo,
|
||||
submit taskSubmitAttempt,
|
||||
) (*taskSubmissionOutcome, *taskdto.TaskError) {
|
||||
policy := service.RequestPolicy(c)
|
||||
diagnostics := newTaskPluginSubmitDiagnostics(c)
|
||||
diagnostics.start(relayInfo)
|
||||
var result *relay.TaskSubmitResult
|
||||
@@ -664,6 +636,7 @@ func executeTaskSubmissionWith(
|
||||
|
||||
if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil {
|
||||
channel = lockedCh
|
||||
policy.BeginAttempt(channel, relayInfo.UsingGroup)
|
||||
if retryParam.GetRetry() > 0 {
|
||||
if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil {
|
||||
taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError)
|
||||
@@ -675,13 +648,13 @@ func executeTaskSubmissionWith(
|
||||
channel, channelErr = getChannel(c, relayInfo, retryParam)
|
||||
if channelErr != nil {
|
||||
logger.LogError(c, channelErr.Error())
|
||||
taskErr = service.TaskErrorWrapperLocal(channelErr.Err, "get_channel_failed", http.StatusInternalServerError)
|
||||
taskErr = service.TaskErrorWrapperLocal(channelErr.Err, "get_channel_failed", channelErr.StatusCode)
|
||||
break
|
||||
}
|
||||
}
|
||||
diagnostics.attempt(retryParam.GetRetry()+1, channel, relayInfo.LockedChannel != nil)
|
||||
|
||||
addUsedChannel(c, channel.Id)
|
||||
service.AppendUsedChannel(c, channel.Id)
|
||||
bodyStorage, bodyErr := common.GetBodyStorage(c)
|
||||
if bodyErr != nil {
|
||||
stage = "read_body"
|
||||
@@ -706,15 +679,19 @@ func executeTaskSubmissionWith(
|
||||
break
|
||||
}
|
||||
|
||||
taskAPIError := taskSubmissionAPIError(taskErr)
|
||||
relayInfo.LastError = taskAPIError
|
||||
decision := decideTaskRetry(c, taskErr, common.RetryTimes-retryParam.GetRetry())
|
||||
service.RecordPolicyFailure(c, channel.Id, taskAPIError, decision)
|
||||
if !taskErr.LocalError {
|
||||
processChannelError(c,
|
||||
*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey,
|
||||
common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()),
|
||||
types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode),
|
||||
taskAPIError,
|
||||
relayInfo)
|
||||
}
|
||||
|
||||
willRetry := shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry())
|
||||
willRetry := decision.Action == "retry"
|
||||
diagnostics.attemptFailed(retryParam.GetRetry()+1, channel, taskErr, willRetry)
|
||||
if !willRetry {
|
||||
break
|
||||
@@ -816,6 +793,11 @@ func executeTaskSubmissionWith(
|
||||
diagnostics.failed("settle", "billing_error", taskErr, true)
|
||||
return nil, taskErr
|
||||
}
|
||||
if task.Status != model.TaskStatusFailure {
|
||||
service.MarkRequestPolicySuccess(c, nil)
|
||||
} else {
|
||||
policy.AddEvent(service.PolicyEvent{Decision: service.PolicyDecision{Action: "stop", Reason: "task_failed", Source: "upstream"}})
|
||||
}
|
||||
service.LogTaskConsumption(c, relayInfo, task)
|
||||
diagnostics.complete(task, result.Quota)
|
||||
|
||||
@@ -876,6 +858,7 @@ func presentTaskSubmission(c *gin.Context, outcome *taskSubmissionOutcome) {
|
||||
}
|
||||
|
||||
func respondTaskSubmissionError(c *gin.Context, taskErr *taskdto.TaskError) {
|
||||
service.RecordRequestPolicyTermination(c, taskSubmissionAPIError(taskErr))
|
||||
newTaskPluginSubmitDiagnostics(c).presentError(taskErr)
|
||||
if middleware.RespondTaskPluginError(c, taskErr) {
|
||||
return
|
||||
@@ -891,44 +874,53 @@ func respondTaskError(c *gin.Context, taskErr *taskdto.TaskError) {
|
||||
c.JSON(taskErr.StatusCode, taskErr)
|
||||
}
|
||||
|
||||
func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *taskdto.TaskError, retryTimes int) bool {
|
||||
if taskErr == nil {
|
||||
return false
|
||||
// taskSubmissionAPIError adapts a task error for the shared relay error paths.
|
||||
// TaskError.Error is nil for many local rejections, so fall back to the message.
|
||||
func taskSubmissionAPIError(taskErr *taskdto.TaskError) *types.NewAPIError {
|
||||
err := taskErr.Error
|
||||
if err == nil {
|
||||
err = errors.New(taskErr.Message)
|
||||
}
|
||||
if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
|
||||
return false
|
||||
}
|
||||
if retryTimes <= 0 {
|
||||
return false
|
||||
}
|
||||
if service.GetChannelConstraints(c).SuppressesRetry() {
|
||||
return false
|
||||
}
|
||||
if taskErr.StatusCode == http.StatusTooManyRequests {
|
||||
return true
|
||||
}
|
||||
if taskErr.StatusCode == 307 {
|
||||
return true
|
||||
}
|
||||
if taskErr.StatusCode/100 == 5 {
|
||||
return types.NewOpenAIError(err, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode)
|
||||
}
|
||||
|
||||
// decideTaskRetry is the single retry decision for task submissions. The
|
||||
// reason is recorded in the request policy decision events of the log details.
|
||||
func decideTaskRetry(c *gin.Context, taskErr *taskdto.TaskError, retryTimes int) service.PolicyDecision {
|
||||
stop := service.PolicyDecision{Action: "stop", Source: "system"}
|
||||
retry := service.PolicyDecision{Action: "retry", Reason: "retry_status_matched", Source: "system"}
|
||||
switch {
|
||||
case taskErr == nil:
|
||||
stop.Reason = "request_completed"
|
||||
case taskErr.NoRetry:
|
||||
stop.Reason = "task_accepted"
|
||||
case service.ShouldSkipRetryAfterChannelAffinityFailure(c):
|
||||
stop.Reason, stop.Source = "strict_session", "session_rule"
|
||||
if source := service.RequestPolicy(c).SessionModeSource; source != "" {
|
||||
stop.Source = source
|
||||
}
|
||||
case retryTimes <= 0:
|
||||
stop.Reason, stop.Source = "attempt_budget_exhausted", "global"
|
||||
case service.GetChannelConstraints(c).SuppressesRetry():
|
||||
stop.Reason, stop.Source = "pinned_channel", "channel_constraint"
|
||||
case taskErr.StatusCode == http.StatusTooManyRequests, taskErr.StatusCode == 307:
|
||||
return retry
|
||||
case taskErr.StatusCode/100 == 5:
|
||||
// 超时不重试
|
||||
if operation_setting.IsAlwaysSkipRetryStatusCode(taskErr.StatusCode) {
|
||||
return false
|
||||
stop.Reason = "system_retry_exclusion"
|
||||
break
|
||||
}
|
||||
return true
|
||||
}
|
||||
if taskErr.StatusCode == http.StatusBadRequest {
|
||||
return false
|
||||
}
|
||||
if taskErr.StatusCode == 408 {
|
||||
return retry
|
||||
case taskErr.StatusCode == http.StatusBadRequest, taskErr.StatusCode == 408:
|
||||
// azure处理超时不重试
|
||||
return false
|
||||
stop.Reason = "status_not_retryable"
|
||||
case taskErr.LocalError:
|
||||
stop.Reason = "local_rejection"
|
||||
case taskErr.StatusCode/100 == 2:
|
||||
stop.Reason = "system_retry_exclusion"
|
||||
default:
|
||||
return retry
|
||||
}
|
||||
if taskErr.LocalError {
|
||||
return false
|
||||
}
|
||||
if taskErr.StatusCode/100 == 2 {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
return stop
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func relayAttemptGuardResult(a model.RelayAttempt) string {
|
||||
if a.GuardResult == "model_mismatch" {
|
||||
return "mismatch_rejected"
|
||||
}
|
||||
if a.GuardResult == "unchecked" && a.ReturnedModel != "" && a.ReturnedModel == a.RequestedModel {
|
||||
return "accepted"
|
||||
}
|
||||
if a.GuardResult == "accepted" && a.Number > 1 {
|
||||
return "switched_success"
|
||||
}
|
||||
return a.GuardResult
|
||||
}
|
||||
|
||||
// Attempt history is owner-scoped; channel identifiers are administrator-only.
|
||||
func GetRelayAttempts(c *gin.Context) {
|
||||
requestID := c.Param("request_id")
|
||||
if requestID == "" || len(requestID) > 64 {
|
||||
c.AbortWithStatus(400)
|
||||
return
|
||||
}
|
||||
query := model.DB.Where("request_id = ?", requestID)
|
||||
admin := c.GetInt("role") >= common.RoleAdminUser
|
||||
if !admin {
|
||||
query = query.Where("user_id = ?", c.GetInt("id"))
|
||||
}
|
||||
var attempts []model.RelayAttempt
|
||||
if err := query.Order("number asc, id asc").Limit(100).Find(&attempts).Error; err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
rows := make([]gin.H, 0, len(attempts))
|
||||
total := 0
|
||||
pending := false
|
||||
for _, a := range attempts {
|
||||
var usage dto.Usage
|
||||
_ = common.UnmarshalJsonStr(a.UsageJSON, &usage)
|
||||
row := gin.H{"number": a.Number, "group": a.GroupName, "requested_model": a.RequestedModel, "returned_model": a.ReturnedModel, "guard_result": relayAttemptGuardResult(a), "usage_source": a.UsageSource, "state": a.State, "quota": a.Quota, "billing_source": a.Funding, "usage": usage}
|
||||
row["charged"] = a.State == "settled" && a.Quota > 0
|
||||
row["charge_explanation"] = "Each consumed attempt is charged independently, including rejected results."
|
||||
var log model.Log
|
||||
var other map[string]json.RawMessage
|
||||
price := map[string]json.RawMessage{}
|
||||
if common.UnmarshalJsonStr(a.LogJSON, &log) == nil && common.UnmarshalJsonStr(log.Other, &other) == nil {
|
||||
for _, key := range []string{"group_ratio", "model_ratio", "completion_ratio", "cache_ratio", "model_price", "billing_mode", "expr_b64", "matched_tier", "billing_tokens", "billing_unit", "fixed_price"} {
|
||||
if value, ok := other[key]; ok {
|
||||
price[key] = value
|
||||
}
|
||||
}
|
||||
}
|
||||
row["price"] = price
|
||||
if admin {
|
||||
row["channel_id"] = a.ChannelID
|
||||
}
|
||||
rows = append(rows, row)
|
||||
total += a.Quota
|
||||
pending = pending || a.State != "settled"
|
||||
}
|
||||
common.ApiSuccess(c, gin.H{"attempts": rows, "total_quota": total, "pending": pending})
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
hostdto "github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestRelayAttemptSwitchSafety(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name, reason, previous, conversation string
|
||||
count int
|
||||
pin, cancel, written, disabled, want bool
|
||||
}{
|
||||
{name: "explicit mismatch", reason: "model_mismatch", count: 1, want: true},
|
||||
{name: "token disabled", reason: "model_mismatch", count: 1, disabled: true},
|
||||
{name: "third channel allowed", reason: "model_mismatch", count: 2, want: true},
|
||||
{name: "third channel final", reason: "model_mismatch", count: 3},
|
||||
{name: "missing model", reason: "missing_model", count: 1},
|
||||
{name: "confirmation timeout", reason: "confirmation_timeout", count: 1},
|
||||
{name: "pinned", reason: "model_mismatch", count: 1, pin: true},
|
||||
{name: "cancelled", reason: "model_mismatch", count: 1, cancel: true},
|
||||
{name: "already delivered", reason: "model_mismatch", count: 1, written: true},
|
||||
{name: "upstream response id", reason: "model_mismatch", count: 1, previous: "resp_private"},
|
||||
{name: "upstream conversation", reason: "model_mismatch", count: 1, conversation: `"conv_private"`},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
c.Request = httptest.NewRequest("POST", "/v1/responses", nil).WithContext(ctx)
|
||||
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "a"}}})
|
||||
if tc.pin {
|
||||
service.GetChannelConstraints(c).AddPin(hostdto.ChannelPin{ChannelId: 1, Source: hostdto.PinSourceToken, Rank: hostdto.PinRankToken, RetryMode: hostdto.PinRetrySingleAttempt})
|
||||
}
|
||||
if tc.cancel {
|
||||
cancel()
|
||||
}
|
||||
if tc.written {
|
||||
_, _ = c.Writer.WriteString("already sent")
|
||||
}
|
||||
info := &relaycommon.RelayInfo{DurableAttemptBilling: true, UpstreamRequests: tc.count, Request: &dto.OpenAIResponsesRequest{PreviousResponseID: tc.previous, Conversation: []byte(tc.conversation)}, ResponseGuard: &relaycommon.ResponseModelGuard{Attempts: []*relaycommon.ResponseGuardAttempt{{Reason: tc.reason}}}}
|
||||
err := types.NewError(errors.New("guard rejected"), relaycommon.ResponseGuardErrorCode)
|
||||
info.AstraAutoSwitch = !tc.disabled
|
||||
require.Equal(t, tc.want, canSwitchRejectedAstra(c, info, err))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRelayAttemptGuardResult(t *testing.T) {
|
||||
require.Equal(t, "mismatch_rejected", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "model_mismatch", Number: 1}))
|
||||
require.Equal(t, "switched_success", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "accepted", Number: 2}))
|
||||
require.Equal(t, "accepted", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "accepted", Number: 1}))
|
||||
require.Equal(t, "accepted", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "unchecked", Number: 1, RequestedModel: "gpt-6-astra", ReturnedModel: "gpt-6-astra"}))
|
||||
require.Equal(t, "unchecked", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "unchecked", Number: 1, RequestedModel: "gpt-6-astra"}))
|
||||
require.Equal(t, "transport_error", relayAttemptGuardResult(model.RelayAttempt{GuardResult: "transport_error", Number: 1}))
|
||||
}
|
||||
@@ -2,6 +2,7 @@ package controller
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
@@ -43,3 +44,20 @@ func TestClientFacingErrorKeepsLocalValidationMessage(t *testing.T) {
|
||||
t.Fatalf("local validation message was hidden: %q", got.Message)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientFacingResponseGuardErrorNamesDeclaredModel(t *testing.T) {
|
||||
err := types.NewErrorWithStatusCode(
|
||||
errors.New("Requested model gpt-6-astra, but the upstream response declared gpt-5.6-luna."),
|
||||
"response_model_rejected",
|
||||
http.StatusConflict,
|
||||
)
|
||||
got := clientFacingOpenAIError(err, true, "local-request-id")
|
||||
if fmt.Sprint(got.Code) != "response_model_rejected" {
|
||||
t.Fatalf("unexpected public code: %v", got.Code)
|
||||
}
|
||||
for _, expected := range []string{"gpt-6-astra", "gpt-5.6-luna", "local-request-id"} {
|
||||
if !strings.Contains(got.Message, expected) {
|
||||
t.Fatalf("public mismatch message %q does not contain %q", got.Message, expected)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -3,8 +3,10 @@ package controller
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -22,14 +24,18 @@ import (
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/mysql"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/schema"
|
||||
)
|
||||
|
||||
type taskSubmissionTestBilling struct {
|
||||
events *[]string
|
||||
settleErr error
|
||||
onSettle func()
|
||||
refunds int
|
||||
events *[]string
|
||||
settleErr error
|
||||
reserveErr error
|
||||
onSettle func()
|
||||
refunds int
|
||||
}
|
||||
|
||||
func (b *taskSubmissionTestBilling) Settle(int) error {
|
||||
@@ -49,7 +55,7 @@ func (b *taskSubmissionTestBilling) NeedsRefund() bool { return b.refunds
|
||||
func (b *taskSubmissionTestBilling) GetPreConsumedQuota() int { return 0 }
|
||||
func (b *taskSubmissionTestBilling) Reserve(int) error {
|
||||
*b.events = append(*b.events, "reserve")
|
||||
return nil
|
||||
return b.reserveErr
|
||||
}
|
||||
|
||||
func TestPresentTaskSubmissionUsesNativePresenterAfterPersistence(t *testing.T) {
|
||||
@@ -387,3 +393,185 @@ func taskSubmissionRelayInfo(billing relaycommon.BillingSettler) *relaycommon.Re
|
||||
ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1, ChannelType: constant.ChannelTypeTaskPlugin},
|
||||
}
|
||||
}
|
||||
|
||||
// Uses the real submit adaptor, expression evaluator, BillingSession, task row
|
||||
// and consume log. Set TEST_TASK_DB_DIALECT plus TEST_MYSQL_DSN or
|
||||
// TEST_POSTGRES_DSN to exercise the same contract on an external test database.
|
||||
// Unique table prefixes keep the fixture isolated from all existing tables.
|
||||
func TestImmediateTaskSettlementDatabase(t *testing.T) {
|
||||
dialect := common.DatabaseType(os.Getenv("TEST_TASK_DB_DIALECT"))
|
||||
var driver gorm.Dialector
|
||||
switch dialect {
|
||||
case "", common.DatabaseTypeSQLite:
|
||||
dialect = common.DatabaseTypeSQLite
|
||||
driver = sqlite.Open(":memory:")
|
||||
case common.DatabaseTypeMySQL:
|
||||
require.NotEmpty(t, os.Getenv("TEST_MYSQL_DSN"))
|
||||
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
|
||||
case common.DatabaseTypePostgreSQL:
|
||||
require.NotEmpty(t, os.Getenv("TEST_POSTGRES_DSN"))
|
||||
driver = postgres.New(postgres.Config{DSN: os.Getenv("TEST_POSTGRES_DSN"), PreferSimpleProtocol: true})
|
||||
default:
|
||||
t.Fatalf("unsupported test dialect %q", dialect)
|
||||
}
|
||||
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("tsubmit_%d_", time.Now().UnixNano())}})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
sqlDB.SetMaxOpenConns(1)
|
||||
t.Cleanup(func() { _ = sqlDB.Close() })
|
||||
models := []any{&model.User{}, &model.Channel{}, &model.Task{}, &model.Log{}}
|
||||
require.NoError(t, db.AutoMigrate(models...))
|
||||
t.Cleanup(func() { require.NoError(t, db.Migrator().DropTable(models...)) })
|
||||
var version string
|
||||
if dialect == common.DatabaseTypeSQLite {
|
||||
require.NoError(t, db.Raw("SELECT sqlite_version()").Scan(&version).Error)
|
||||
} else {
|
||||
require.NoError(t, db.Raw("SELECT version()").Scan(&version).Error)
|
||||
}
|
||||
t.Logf("database: %s %s", dialect, version)
|
||||
oldDB, oldLogDB := model.DB, model.LOG_DB
|
||||
oldMain, oldLog := common.MainDatabaseType(), common.LogDatabaseType()
|
||||
oldRedis, oldMemory, oldBatch, oldConsume, oldExport := common.RedisEnabled, common.MemoryCacheEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled, common.DataExportEnabled
|
||||
model.DB, model.LOG_DB = db, db
|
||||
common.SetDatabaseTypes(dialect, dialect)
|
||||
common.RedisEnabled, common.MemoryCacheEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled, common.DataExportEnabled = false, false, false, true, false
|
||||
t.Cleanup(func() {
|
||||
model.DB, model.LOG_DB = oldDB, oldLogDB
|
||||
common.SetDatabaseTypes(oldMain, oldLog)
|
||||
common.RedisEnabled, common.MemoryCacheEnabled, common.BatchUpdateEnabled, common.LogConsumeEnabled, common.DataExportEnabled = oldRedis, oldMemory, oldBatch, oldConsume, oldExport
|
||||
})
|
||||
|
||||
const expression = `u("units") == 7 ? tier("missing", u("missing") * 1.0) : u("units") == 8 ? tier("invalid", -1.0) : u("units") > 3 ? tier("bulk", u("units") * 0.01) : tier("small", u("units") * 0.01)`
|
||||
withTieredBillingConfig(t, map[string]string{"document-model": "tiered_expr"}, map[string]string{"document-model": expression})
|
||||
const source = `
|
||||
export const meta={apiVersion:1,key:"generic-settlement",name:"Generic settlement",version:"1.0.0",author:{name:"Test"},models:["document-model"],fetchMode:"per_task",usageSchema:{units:{type:"number",unit:"count"}}};
|
||||
export function buildSubmitRequest(ctx){return {url:ctx.baseUrl+"/compile",body:ctx.requestBody};}
|
||||
export function parseSubmitResponse(ctx,resp){return {taskId:"vendor-job",taskData:resp.body,immediate:{status:resp.body.status,reason:"provider rejected job"}};}
|
||||
export function extractUsage(){return {units:4};}
|
||||
export function extractUsageOnComplete(ctx,result,body){return body.usage;}
|
||||
export function parseTaskResult(){throw new Error("completed submissions must not poll");}
|
||||
export function buildQueryRequest(){throw new Error("completed submissions must not poll");}
|
||||
`
|
||||
plugin, err := pluginruntime.CompilePlugin(source, pluginruntime.Options{})
|
||||
require.NoError(t, err)
|
||||
for index, tc := range []struct {
|
||||
name, status string
|
||||
actual any
|
||||
count float64
|
||||
}{
|
||||
{"partial", "SUCCESS", 2, 2}, {"zero", "SUCCESS", 0, 0}, {"larger", "SUCCESS", 6, 6},
|
||||
{"invalid usage", "SUCCESS", -1, 4}, {"expression failure", "SUCCESS", 7, 4}, {"negative result", "SUCCESS", 8, 4}, {"missing usage", "SUCCESS", nil, 4}, {"failed", "FAILURE", 9, 0},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
body := map[string]any{"status": tc.status}
|
||||
if tc.actual != nil {
|
||||
body["usage"] = map[string]any{"units": tc.actual}
|
||||
}
|
||||
encoded, err := common.Marshal(body)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
_, _ = w.Write(encoded)
|
||||
}))
|
||||
defer server.Close()
|
||||
initial := int(20 * common.QuotaPerUnit)
|
||||
user := model.User{Username: fmt.Sprintf("task_user_%d", index), AffCode: fmt.Sprintf("task_aff_%d", index), Quota: initial}
|
||||
require.NoError(t, db.Create(&user).Error)
|
||||
ch := model.Channel{Name: "test provider", Type: constant.ChannelTypeTaskPlugin}
|
||||
require.NoError(t, db.Create(&ch).Error)
|
||||
c := taskSubmissionTestContext()
|
||||
c.Set("group", "default")
|
||||
c.Set("username", user.Username)
|
||||
c.Set("task_request", map[string]any{"model": "document-model"})
|
||||
c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{Plugin: plugin})
|
||||
common.SetContextKey(c, constant.ContextKeyOriginalModel, "document-model")
|
||||
common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, server.URL)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelId, ch.Id)
|
||||
common.SetContextKey(c, constant.ContextKeyChannelType, ch.Type)
|
||||
info := taskSubmissionRelayInfo(nil)
|
||||
info.UserId = user.Id
|
||||
info.OriginModelName = "document-model"
|
||||
info.UserGroup = "default"
|
||||
info.UsingGroup = "default"
|
||||
info.IsPlayground = true
|
||||
info.UserSetting.BillingPreference = "wallet_only"
|
||||
info.PublicTaskID = model.GenerateTaskID()
|
||||
info.LockedChannel = &ch
|
||||
outcome, taskErr := executeTaskSubmission(c, info)
|
||||
require.Nil(t, taskErr)
|
||||
require.NotNil(t, outcome)
|
||||
want := common.QuotaRound(tc.count * 0.01 * common.QuotaPerUnit)
|
||||
assert.Equal(t, want, outcome.Result.Quota)
|
||||
assert.Equal(t, want, info.PriceData.Quota)
|
||||
var stored model.Task
|
||||
require.NoError(t, db.Where("task_id = ?", info.PublicTaskID).First(&stored).Error)
|
||||
assert.Equal(t, want, stored.Quota)
|
||||
assert.Equal(t, model.TaskStatus(tc.status), stored.Status)
|
||||
assert.Positive(t, stored.FinishTime)
|
||||
assert.Equal(t, float64(4), info.TieredBillingSnapshot.EstimatedQuotaBeforeGroup/(0.01*common.QuotaPerUnit))
|
||||
if tc.status == "SUCCESS" {
|
||||
assert.Equal(t, tc.count, stored.PrivateData.BillingContext.TieredSnapshot.UsageFacts["units"])
|
||||
}
|
||||
var updated model.User
|
||||
require.NoError(t, db.First(&updated, user.Id).Error)
|
||||
assert.Equal(t, initial-want, updated.Quota)
|
||||
assert.Equal(t, want, updated.UsedQuota)
|
||||
var logs []model.Log
|
||||
require.NoError(t, db.Where("user_id = ?", user.Id).Find(&logs).Error)
|
||||
require.Len(t, logs, 1)
|
||||
assert.Equal(t, want, logs[0].Quota)
|
||||
var other map[string]any
|
||||
require.NoError(t, common.UnmarshalJsonStr(logs[0].Other, &other))
|
||||
if tc.status == "SUCCESS" {
|
||||
assert.Equal(t, tc.count, other["usage_facts"].(map[string]any)["units"])
|
||||
}
|
||||
assert.False(t, c.Writer.Written(), "presentation must follow persistence and settlement")
|
||||
require.NoError(t, info.Billing.Settle(want))
|
||||
info.Billing.Refund(c)
|
||||
require.NoError(t, db.First(&updated, user.Id).Error)
|
||||
assert.Equal(t, initial-want, updated.Quota, "terminal settlement is idempotent")
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcceptedSubmitStreamNeverRetries(t *testing.T) {
|
||||
c := taskSubmissionTestContext()
|
||||
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "task_accepted", Source: "system"}, decideTaskRetry(c, &dto.TaskError{StatusCode: 502, LocalError: true, NoRetry: true}, 3))
|
||||
}
|
||||
|
||||
// Local task rejections carry a message but no cause; the response and the
|
||||
// decision record must still be produced.
|
||||
func TestRespondTaskSubmissionErrorWithoutCause(t *testing.T) {
|
||||
previousErrorLog := constant.ErrorLogEnabled
|
||||
constant.ErrorLogEnabled = false
|
||||
t.Cleanup(func() { constant.ErrorLogEnabled = previousErrorLog })
|
||||
c := taskSubmissionTestContext()
|
||||
taskErr := &dto.TaskError{Code: "get_channel_failed", Message: "no channel", StatusCode: http.StatusServiceUnavailable, LocalError: true}
|
||||
require.NotPanics(t, func() { respondTaskSubmissionError(c, taskErr) })
|
||||
assert.Equal(t, http.StatusServiceUnavailable, c.Writer.Status())
|
||||
events := service.RequestPolicy(c).Events()
|
||||
require.Len(t, events, 1)
|
||||
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "request_failed", Source: "system"}, events[0].Decision)
|
||||
assert.Equal(t, http.StatusServiceUnavailable, events[0].Status)
|
||||
assert.Equal(t, service.PolicyDecision{Action: "stop", Reason: "local_rejection", Source: "system"}, decideTaskRetry(c, &dto.TaskError{StatusCode: http.StatusForbidden, LocalError: true, Message: "billing"}, 2), "local errors stop once the status rules do not force a retry")
|
||||
}
|
||||
|
||||
func TestExecuteTaskSubmissionRefundsWhenFinalReserveFails(t *testing.T) {
|
||||
events := []string{}
|
||||
setupTaskSubmissionDatabase(t, true, &events)
|
||||
billing := &taskSubmissionTestBilling{events: &events, reserveErr: errors.New("insufficient funds")}
|
||||
c := taskSubmissionTestContext()
|
||||
info := taskSubmissionRelayInfo(billing)
|
||||
outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) {
|
||||
return &relay.TaskSubmitResult{Platform: "plugin", Quota: 600, Immediate: &relaycommon.TaskInfo{Status: "SUCCESS"}}, nil
|
||||
})
|
||||
require.Nil(t, outcome)
|
||||
require.NotNil(t, taskErr)
|
||||
assert.Equal(t, http.StatusForbidden, taskErr.StatusCode)
|
||||
assert.Equal(t, []string{"reserve", "refund"}, events)
|
||||
assert.Equal(t, 1, billing.refunds)
|
||||
assert.False(t, c.Writer.Written())
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetRequestPolicy(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "data": gin.H{"options": model.CurrentRequestPolicy().Options}})
|
||||
}
|
||||
|
||||
func UpdateRequestPolicy(c *gin.Context) {
|
||||
var request struct {
|
||||
Options map[string]string `json:"options"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&request); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
|
||||
return
|
||||
}
|
||||
if err := model.UpdateRequestPolicyOptions(request.Options); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": err.Error()})
|
||||
return
|
||||
}
|
||||
GetRequestPolicy(c)
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/relay"
|
||||
"github.com/QuantumNous/new-api/relay/helper"
|
||||
"github.com/QuantumNous/new-api/relaykit/types"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
type responsesWSRequestContextKey struct{}
|
||||
|
||||
type responsesWSRequestState struct {
|
||||
requestID string
|
||||
handle func(*gin.Context) *types.NewAPIError
|
||||
apiError *types.NewAPIError
|
||||
}
|
||||
|
||||
// Each response.create runs the ordinary request middleware to completion,
|
||||
// without routing another HTTP request or retaining a pooled Gin context.
|
||||
var responsesWSRequestEngine = sync.OnceValue(func() *gin.Engine {
|
||||
engine := gin.New()
|
||||
engine.ForwardedByClientIP = false
|
||||
_ = engine.SetTrustedProxies(nil)
|
||||
engine.POST("/v1/responses", func(c *gin.Context) {
|
||||
state := c.Request.Context().Value(responsesWSRequestContextKey{}).(*responsesWSRequestState)
|
||||
c.Set(common.RequestIdKey, state.requestID)
|
||||
common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now())
|
||||
c.Next()
|
||||
}, middleware.BodyStorageCleanup(), middleware.TokenAuth(), middleware.ModelRequestRateLimit(), func(c *gin.Context) {
|
||||
state := c.Request.Context().Value(responsesWSRequestContextKey{}).(*responsesWSRequestState)
|
||||
state.apiError = state.handle(c)
|
||||
if state.apiError != nil {
|
||||
status := state.apiError.StatusCode
|
||||
if status < http.StatusBadRequest {
|
||||
status = http.StatusInternalServerError
|
||||
}
|
||||
c.Status(status)
|
||||
}
|
||||
})
|
||||
return engine
|
||||
})
|
||||
|
||||
type responsesWSResponseWriter struct {
|
||||
header http.Header
|
||||
status int
|
||||
body bytes.Buffer
|
||||
}
|
||||
|
||||
func (w *responsesWSResponseWriter) Header() http.Header {
|
||||
return w.header
|
||||
}
|
||||
|
||||
func (w *responsesWSResponseWriter) WriteHeader(status int) {
|
||||
if w.status == 0 {
|
||||
w.status = status
|
||||
}
|
||||
}
|
||||
|
||||
func (w *responsesWSResponseWriter) Write(data []byte) (int, error) {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
return w.body.Write(data)
|
||||
}
|
||||
|
||||
func newResponsesWSRequestRunner(c *gin.Context) relay.ResponsesWSRequestRunner {
|
||||
// Capture credentials before channel selection or header overrides. Resolve
|
||||
// the peer once with the public router's trusted-proxy configuration.
|
||||
headers := c.Request.Header.Clone()
|
||||
for _, name := range []string{"Connection", "Upgrade", "Sec-WebSocket-Key", "Sec-WebSocket-Version", "Sec-WebSocket-Extensions", "Sec-WebSocket-Protocol", "Content-Length", "Content-Encoding"} {
|
||||
headers.Del(name)
|
||||
}
|
||||
remoteAddr := net.JoinHostPort(c.ClientIP(), "0")
|
||||
return func(request *http.Request, requestID string, handle func(*gin.Context) *types.NewAPIError) *types.NewAPIError {
|
||||
state := &responsesWSRequestState{requestID: requestID, handle: handle}
|
||||
ctx := context.WithValue(request.Context(), responsesWSRequestContextKey{}, state)
|
||||
ctx = context.WithValue(ctx, common.RequestIdKey, requestID)
|
||||
request = request.Clone(ctx)
|
||||
request.Method = http.MethodPost
|
||||
request.URL.Path = "/v1/responses"
|
||||
request.URL.RawPath = ""
|
||||
request.Header = headers.Clone()
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
request.RemoteAddr = remoteAddr
|
||||
response := &responsesWSResponseWriter{header: make(http.Header)}
|
||||
responsesWSRequestEngine().ServeHTTP(response, request)
|
||||
if state.apiError != nil {
|
||||
return state.apiError
|
||||
}
|
||||
if response.status < http.StatusBadRequest {
|
||||
return nil
|
||||
}
|
||||
var body struct {
|
||||
Error *types.OpenAIError `json:"error"`
|
||||
}
|
||||
if common.Unmarshal(response.body.Bytes(), &body) == nil && body.Error != nil {
|
||||
return types.WithOpenAIError(*body.Error, response.status, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
|
||||
}
|
||||
// The existing in-memory rate limiter returns a bare 429 response.
|
||||
return types.NewErrorWithStatusCode(errors.New(http.StatusText(response.status)), types.ErrorCodeInvalidRequest, response.status, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
|
||||
}
|
||||
}
|
||||
|
||||
func ResponsesWebSocket(c *gin.Context) {
|
||||
requestID := c.GetString(common.RequestIdKey)
|
||||
runner := newResponsesWSRequestRunner(c)
|
||||
ws, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer ws.Close()
|
||||
|
||||
if apiError := relay.ResponsesWebSocketHelper(c, ws, runner); apiError != nil {
|
||||
logger.LogError(c, fmt.Sprintf("responses websocket relay error: %s", common.LocalLogPreview(apiError.Error())))
|
||||
apiError.SetMessage(common.MessageWithRequestId(apiError.Error(), requestID))
|
||||
helper.WssError(c, ws, apiError.ToOpenAIError())
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -3,12 +3,18 @@ package controller
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/oauth"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
passkeysvc "github.com/QuantumNous/new-api/service/passkey"
|
||||
"github.com/QuantumNous/new-api/setting/system_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/go-webauthn/webauthn/protocol"
|
||||
)
|
||||
@@ -34,6 +40,10 @@ func writeSecurityOperationError(c *gin.Context, err error) {
|
||||
var code, message string
|
||||
var protocolError *protocol.Error
|
||||
switch {
|
||||
case errors.Is(err, passkeysvc.ErrRPIDUnavailable):
|
||||
code, message = "PASSKEY_RP_ID_UNAVAILABLE", i18n.T(c, i18n.MsgPasskeyRPIDUnavailable)
|
||||
case errors.Is(err, system_setting.ErrPasskeyRPIDInvalid):
|
||||
code, message = "PASSKEY_RP_ID_INVALID", i18n.T(c, i18n.MsgPasskeyRPIDInvalid)
|
||||
case errors.Is(err, service.ErrAccountEmailInvalid), errors.Is(err, service.ErrAccountEmailRestricted):
|
||||
code, message = "EMAIL_ADDRESS_REJECTED", err.Error()
|
||||
case errors.Is(err, model.ErrEmailAlreadyTaken):
|
||||
@@ -112,6 +122,15 @@ func writeSecurityOperationError(c *gin.Context, err error) {
|
||||
return
|
||||
}
|
||||
c.Set("security_error_code", code)
|
||||
if strings.Contains(c.Request.URL.Path, "/passkey/") {
|
||||
reason := "verification_failed"
|
||||
if errors.As(err, &protocolError) && slices.Contains([]string{"invalid_request", "challenge_mismatch", "parse_error", "auth_data", "verification_error", "invalid_signature", "invalid_key_type", "unsupported_key_algorithm"}, protocolError.Type) {
|
||||
reason = protocolError.Type
|
||||
}
|
||||
// Protocol details can contain challenges and client-controlled data.
|
||||
// Only fixed categories and the server-selected public RP ID are logged.
|
||||
logger.LogWarn(c.Request.Context(), "passkey verification rejected: code=%s reason=%s rp_id=%q", code, reason, c.GetString("passkey_rp_id"))
|
||||
}
|
||||
c.JSON(status, gin.H{"success": false, "code": code, "message": message})
|
||||
}
|
||||
|
||||
|
||||
@@ -64,7 +64,7 @@ func setupSecurityEnrollmentTest(t *testing.T) (*model.User, service.AuthIdentit
|
||||
var version string
|
||||
require.NoError(t, db.Raw(versionQuery).Scan(&version).Error)
|
||||
t.Logf("database: %s %s", dialect, version)
|
||||
require.NoError(t, db.AutoMigrate(&model.User{}, &model.UserSession{}, &model.TwoFA{}, &model.TwoFABackupCode{}, &model.PasskeyCredential{}, &model.AuthFlow{}, &model.UserOAuthBinding{}))
|
||||
require.NoError(t, db.AutoMigrate(&model.User{}, &model.UserSession{}, &model.TwoFA{}, &model.TwoFABackupCode{}, &model.PasskeyCredential{}, &model.AuthFlow{}, &model.UserOAuthBinding{}, &model.Option{}))
|
||||
require.NoError(t, logDB.AutoMigrate(&model.AuditLog{}))
|
||||
model.DB, model.LOG_DB = db, logDB
|
||||
dbType := common.DatabaseTypeSQLite
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func GetStatusPageAdmin(c *gin.Context) {
|
||||
results, err := currentGroupMonitors(c.Request.Context(), true)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "data": results})
|
||||
}
|
||||
|
||||
func UpdateStatusPageGroup(c *gin.Context) {
|
||||
var input struct {
|
||||
Group string `json:"group"`
|
||||
Alias string `json:"alias"`
|
||||
Hidden *bool `json:"hidden"`
|
||||
}
|
||||
if c.ShouldBindJSON(&input) != nil || strings.TrimSpace(input.Group) == "" || utf8.RuneCountInString(input.Group) > 64 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "Invalid group"})
|
||||
return
|
||||
}
|
||||
reset := strings.HasSuffix(c.FullPath(), "/reset")
|
||||
input.Alias = strings.TrimSpace(input.Alias)
|
||||
if !reset && (input.Hidden == nil || utf8.RuneCountInString(input.Alias) > 80 || strings.ContainsAny(input.Alias, "\r\n\t")) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "Invalid display settings"})
|
||||
return
|
||||
}
|
||||
channels, err := model.GetStatusChannels(c.Request.Context())
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
found := false
|
||||
for _, channel := range channels {
|
||||
for group := range strings.SplitSeq(channel.Group, ",") {
|
||||
if strings.TrimSpace(group) == input.Group {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
c.JSON(http.StatusNotFound, gin.H{"success": false, "message": "Group not found"})
|
||||
return
|
||||
}
|
||||
var display *model.StatusPageGroup
|
||||
if !reset {
|
||||
display = &model.StatusPageGroup{Alias: input.Alias, Hidden: *input.Hidden}
|
||||
}
|
||||
if err := model.SaveStatusPageGroup(c.Request.Context(), input.Group, display); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true})
|
||||
}
|
||||
|
||||
// Main-site group status is independent of the public uptime display settings.
|
||||
func GetSiteGroupStatus(c *gin.Context) {
|
||||
results, err := currentGroupMonitors(c.Request.Context(), true)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
for i := range results {
|
||||
monitors := make([]Monitor, 0, len(results[i].Monitors))
|
||||
for _, monitor := range results[i].Monitors {
|
||||
monitor.Name = monitor.Group
|
||||
monitor.Alias = ""
|
||||
monitor.Hidden = false
|
||||
monitors = append(monitors, monitor)
|
||||
}
|
||||
results[i].Monitors = monitors
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{"success": true, "data": results})
|
||||
}
|
||||
@@ -66,9 +66,17 @@ func GetCurrentSystemTask(c *gin.Context) {
|
||||
}
|
||||
|
||||
func ListSystemTasks(c *gin.Context) {
|
||||
limit, _ := strconv.Atoi(c.Query("limit"))
|
||||
var query struct {
|
||||
model.SystemTaskFilter
|
||||
Offset int `form:"offset" binding:"gte=0"`
|
||||
Limit int `form:"limit"`
|
||||
}
|
||||
if err := c.ShouldBindQuery(&query); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "invalid system task filters"})
|
||||
return
|
||||
}
|
||||
|
||||
tasks, err := model.ListSystemTasks(limit)
|
||||
tasks, total, err := model.ListSystemTasks(query.SystemTaskFilter, query.Offset, query.Limit)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
@@ -83,6 +91,25 @@ func ListSystemTasks(c *gin.Context) {
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": responses,
|
||||
"total": total,
|
||||
})
|
||||
}
|
||||
|
||||
func DeleteSystemTaskHistory(c *gin.Context) {
|
||||
var filter model.SystemTaskFilter
|
||||
if err := c.ShouldBindQuery(&filter); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"success": false, "message": "invalid system task filters"})
|
||||
return
|
||||
}
|
||||
deleted, err := model.DeleteSystemTaskHistory(filter)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"success": true,
|
||||
"message": "",
|
||||
"data": gin.H{"deleted_count": deleted},
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@ package controller
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
@@ -18,6 +19,9 @@ import (
|
||||
// instances and each run is recorded as one task row. Call this before
|
||||
// service.StartSystemTaskRunner.
|
||||
func RegisterScheduledSystemTasks() {
|
||||
service.RegisterSystemTaskHandler(detectionHandler{})
|
||||
service.RegisterSystemTaskHandler(astraSamplingHandler{})
|
||||
service.RegisterSystemTaskHandler(channelAvailabilityHandler{})
|
||||
service.RegisterSystemTaskHandler(channelTestHandler{})
|
||||
service.RegisterSystemTaskHandler(modelUpdateHandler{})
|
||||
service.RegisterSystemTaskHandler(midjourneyPollHandler{})
|
||||
@@ -29,6 +33,74 @@ func RegisterScheduledSystemTasks() {
|
||||
// into the system task runner.
|
||||
type channelTestHandler struct{}
|
||||
|
||||
type channelAvailabilityHandler struct{}
|
||||
|
||||
func (channelAvailabilityHandler) Type() string { return model.SystemTaskTypeChannelAvailability }
|
||||
|
||||
func (channelAvailabilityHandler) Enabled() bool { return service.AstraAvailabilityEnabled() }
|
||||
|
||||
func (channelAvailabilityHandler) Interval() time.Duration {
|
||||
return time.Duration(service.AstraAvailabilityIntervalMinutes()) * time.Minute
|
||||
}
|
||||
|
||||
func (channelAvailabilityHandler) ScheduleFromCreation() bool { return true }
|
||||
|
||||
func (channelAvailabilityHandler) NewPayload() any { return nil }
|
||||
|
||||
func (channelAvailabilityHandler) Run(ctx context.Context, task *model.SystemTask, runnerID string) {
|
||||
testUserID, err := resolveChannelTestUserID(nil)
|
||||
if err != nil {
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err)
|
||||
return
|
||||
}
|
||||
channels, err := model.GetAllChannels(0, 0, true, false)
|
||||
if err != nil {
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err)
|
||||
return
|
||||
}
|
||||
selected := make([]*model.Channel, 0, len(channels))
|
||||
for _, channel := range channels {
|
||||
if channel.Status == common.ChannelStatusEnabled {
|
||||
selected = append(selected, channel)
|
||||
}
|
||||
}
|
||||
settings := operation_setting.GetMonitorSetting()
|
||||
state := service.ChannelAvailabilityState{Rounds: []service.ChannelAvailabilityRound{}}
|
||||
var mu sync.Mutex
|
||||
var persistErr error
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
disableThreshold := int64(common.ChannelDisableThreshold * 1000)
|
||||
if disableThreshold == 0 {
|
||||
disableThreshold = 10000000
|
||||
}
|
||||
summary := runChannelTestWorkers(ctx, selected, settings.ChannelTestConcurrency, func(ctx context.Context, channel *model.Channel) channelTestSummary {
|
||||
result := testChannelForHealthCheck(ctx, channel, testUserID, false, disableThreshold)
|
||||
if result.Tested == 0 {
|
||||
return result
|
||||
}
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
if persistErr == nil {
|
||||
state.Rounds = append(state.Rounds, service.ChannelAvailabilityRound{ChannelID: channel.Id, At: time.Now().Unix(), Success: result.Succeeded == 1})
|
||||
persistErr = model.UpdateSystemTaskState(task.TaskID, runnerID, state)
|
||||
if persistErr != nil {
|
||||
cancel()
|
||||
}
|
||||
}
|
||||
return result
|
||||
}, nil)
|
||||
if persistErr != nil {
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, summary, persistErr)
|
||||
return
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, summary, ctx.Err())
|
||||
return
|
||||
}
|
||||
finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, summary, nil)
|
||||
}
|
||||
|
||||
func (channelTestHandler) Type() string { return model.SystemTaskTypeChannelTest }
|
||||
|
||||
func (channelTestHandler) Enabled() bool {
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestSystemTaskListFiltersAndPaginationResponse(t *testing.T) {
|
||||
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
|
||||
require.NoError(t, err)
|
||||
sqlDB, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
previousDB := model.DB
|
||||
model.DB = db
|
||||
t.Cleanup(func() {
|
||||
model.DB = previousDB
|
||||
require.NoError(t, sqlDB.Close())
|
||||
})
|
||||
require.NoError(t, db.AutoMigrate(&model.SystemTask{}))
|
||||
require.NoError(t, db.Create(&[]model.SystemTask{
|
||||
{TaskID: "older", Type: model.SystemTaskTypeModelUpdate, Status: model.SystemTaskStatusFailed},
|
||||
{TaskID: "newer", Type: model.SystemTaskTypeModelUpdate, Status: model.SystemTaskStatusFailed},
|
||||
{TaskID: "other-status", Type: model.SystemTaskTypeModelUpdate, Status: model.SystemTaskStatusSucceeded},
|
||||
{TaskID: "other-type", Type: model.SystemTaskTypeChannelTest, Status: model.SystemTaskStatusFailed},
|
||||
}).Error)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/api/system-task/list?scope=history&type=model_update&status=failed&offset=1&limit=1", nil)
|
||||
ListSystemTasks(c)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
var response struct {
|
||||
Success bool `json:"success"`
|
||||
Total int64 `json:"total"`
|
||||
Data []model.SystemTaskResponse `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
assert.True(t, response.Success)
|
||||
assert.EqualValues(t, 2, response.Total)
|
||||
require.Len(t, response.Data, 1)
|
||||
assert.Equal(t, "older", response.Data[0].TaskID)
|
||||
}
|
||||
|
||||
func TestSystemTaskInvalidFiltersAreRejected(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
query string
|
||||
handler gin.HandlerFunc
|
||||
}{
|
||||
{"?scope=invalid", ListSystemTasks},
|
||||
{"?status=unknown", ListSystemTasks},
|
||||
{"?offset=-1", ListSystemTasks},
|
||||
{"?limit=invalid", ListSystemTasks},
|
||||
{"?status=unknown", DeleteSystemTaskHistory},
|
||||
{"?scope=invalid", DeleteSystemTaskHistory},
|
||||
} {
|
||||
t.Run(tc.query, func(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/"+tc.query, nil)
|
||||
tc.handler(c)
|
||||
assert.Equal(t, http.StatusBadRequest, recorder.Code)
|
||||
assert.JSONEq(t, `{"success":false,"message":"invalid system task filters"}`, recorder.Body.String())
|
||||
})
|
||||
}
|
||||
}
|
||||
+27
-14
@@ -15,6 +15,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
@@ -26,6 +27,15 @@ import (
|
||||
|
||||
const maxTaskPluginSourceBytes = 1024 * 1024
|
||||
|
||||
func taskPluginCompileError(c *gin.Context, err error) {
|
||||
var unknownField *jsplugin.UnknownMetaFieldError
|
||||
if errors.As(err, &unknownField) {
|
||||
common.ApiErrorI18n(c, i18n.MsgTaskPluginUnknownMetaField, map[string]any{"Field": unknownField.Field})
|
||||
return
|
||||
}
|
||||
common.ApiError(c, err)
|
||||
}
|
||||
|
||||
type taskPluginUploadRequest struct {
|
||||
Source string `json:"source" binding:"required"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
@@ -57,7 +67,7 @@ func UploadTaskPlugin(c *gin.Context) {
|
||||
temporary := jsplugin.NewRegistry()
|
||||
loaded, err := temporary.Register(request.Source, jsplugin.Options{})
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
taskPluginCompileError(c, err)
|
||||
return
|
||||
}
|
||||
if err = jsplugin.ValidateV1Meta(loaded.Meta); err != nil {
|
||||
@@ -335,7 +345,7 @@ func GetTaskPlugin(c *gin.Context) {
|
||||
if err == nil {
|
||||
loaded, compileErr := jsplugin.NewRegistry().Register(plugin.Source, jsplugin.Options{Key: plugin.Key, Version: plugin.Version})
|
||||
if compileErr != nil {
|
||||
common.ApiErrorMsg(c, compileErr.Error())
|
||||
taskPluginCompileError(c, compileErr)
|
||||
return
|
||||
}
|
||||
common.ApiSuccess(c, taskPluginDetail{Plugin: plugin, Meta: loaded.Meta, Source: plugin.Source, Layer: "override", HasIcon: plugin.HasIcon()})
|
||||
@@ -352,7 +362,7 @@ func GetTaskPlugin(c *gin.Context) {
|
||||
}
|
||||
loaded, err := jsplugin.NewRegistry().RegisterFactory(source, jsplugin.Options{Key: key})
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
taskPluginCompileError(c, err)
|
||||
return
|
||||
}
|
||||
_, _, hasIcon := plugins.Icon(key)
|
||||
@@ -384,7 +394,7 @@ func DryRunTaskPlugin(c *gin.Context) {
|
||||
}
|
||||
loaded, err := jsplugin.NewRegistry().Register(detailSource, jsplugin.Options{Key: c.Param("key")})
|
||||
if err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
taskPluginCompileError(c, err)
|
||||
return
|
||||
}
|
||||
args := make([]any, len(request.Args))
|
||||
@@ -473,7 +483,7 @@ func ActivateTaskPlugin(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
if _, err = jsplugin.NewRegistry().Register(target.Source, jsplugin.Options{Key: target.Key, Version: target.Version}); err != nil {
|
||||
common.ApiErrorMsg(c, err.Error())
|
||||
taskPluginCompileError(c, err)
|
||||
return
|
||||
}
|
||||
if err = model.ActivateTaskPlugin(target.Key, target.Version); err != nil {
|
||||
@@ -649,15 +659,18 @@ func GetTaskPluginOptions(c *gin.Context) {
|
||||
_, _, hasIcon = plugins.Icon(meta.Key)
|
||||
}
|
||||
options = append(options, gin.H{
|
||||
"key": meta.Key,
|
||||
"name": meta.Name,
|
||||
"icon": meta.Icon,
|
||||
"hasIcon": hasIcon,
|
||||
"baseUrl": meta.BaseURL,
|
||||
"sortPriority": meta.SortPriority,
|
||||
"website": meta.Website,
|
||||
"models": meta.Models,
|
||||
"usageSchema": meta.UsageSchema,
|
||||
"key": meta.Key,
|
||||
"name": meta.Name,
|
||||
"description": meta.Description,
|
||||
"icon": meta.Icon,
|
||||
"hasIcon": hasIcon,
|
||||
"baseUrl": meta.BaseURL,
|
||||
"sortPriority": meta.SortPriority,
|
||||
"website": meta.Website,
|
||||
"models": meta.Models,
|
||||
"channelTypes": meta.ChannelTypes,
|
||||
"usageSchema": meta.UsageSchema,
|
||||
"usageProfiles": meta.UsageProfiles,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/i18n"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/pkg/jsplugin"
|
||||
"github.com/QuantumNous/new-api/plugins"
|
||||
@@ -358,15 +359,18 @@ func TestMasterSwitchEmptiesOptionsAndKeepsList(t *testing.T) {
|
||||
assert.Equal(t, "kling", item.Meta.Key)
|
||||
}
|
||||
|
||||
func TestGetTaskPluginOptionsIncludesUsageSchemaIconAndBaseURL(t *testing.T) {
|
||||
func TestGetTaskPluginOptionsIncludesDescriptionUsageSchemaIconBaseURLAndChannelTypes(t *testing.T) {
|
||||
setupTaskPluginControllerTest(t)
|
||||
const key = "usage-options-probe"
|
||||
source := `
|
||||
export const meta = {
|
||||
apiVersion: 1, key: "usage-options-probe", name: "Usage Options", version: "1.0.0", author: {name: "Test"},
|
||||
description: {en: "Video generation via the vendor API", zh: "通过厂商接口生成视频"},
|
||||
icon: "text:UO", baseUrl: "http://localhost:9000/",
|
||||
channelTypes: [1990, 1991],
|
||||
models: ["usage-options-model"], fetchMode: "per_task",
|
||||
usageSchema: {seconds: {type: "number", unit: "second", description: "Generated media duration."}}
|
||||
usageSchema: {seconds: {type: "number", unit: "second", description: "Video generation unit price"}},
|
||||
usageProfiles: [{models:["usage-options-model"],schema:{image_count:{type:"number",unit:"count"}}}]
|
||||
};
|
||||
export function buildSubmitRequest() { return {}; }
|
||||
export function parseSubmitResponse() { return {}; }
|
||||
@@ -386,10 +390,13 @@ export function parseTaskResult() { return {}; }
|
||||
var response struct {
|
||||
Success bool `json:"success"`
|
||||
Data []struct {
|
||||
Key string `json:"key"`
|
||||
Icon string `json:"icon"`
|
||||
BaseURL string `json:"baseUrl"`
|
||||
UsageSchema map[string]jsplugin.UsageFieldSchema `json:"usageSchema"`
|
||||
Key string `json:"key"`
|
||||
Description jsplugin.LocalizedText `json:"description"`
|
||||
Icon string `json:"icon"`
|
||||
BaseURL string `json:"baseUrl"`
|
||||
ChannelTypes []int `json:"channelTypes"`
|
||||
UsageSchema map[string]jsplugin.UsageFieldSchema `json:"usageSchema"`
|
||||
UsageProfiles []jsplugin.UsageProfile `json:"usageProfiles"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
@@ -398,9 +405,14 @@ export function parseTaskResult() { return {}; }
|
||||
if option.Key != key {
|
||||
continue
|
||||
}
|
||||
assert.Equal(t, jsplugin.LocalizedText{"en": "Video generation via the vendor API", "zh": "通过厂商接口生成视频"}, option.Description)
|
||||
assert.Equal(t, "second", option.UsageSchema["seconds"].Unit)
|
||||
assert.Equal(t, "Generated media duration.", option.UsageSchema["seconds"].Description["en"])
|
||||
assert.Equal(t, "Video generation unit price", option.UsageSchema["seconds"].Description["en"])
|
||||
require.Len(t, option.UsageProfiles, 1)
|
||||
assert.Equal(t, []string{"usage-options-model"}, option.UsageProfiles[0].Models)
|
||||
assert.Equal(t, "count", option.UsageProfiles[0].Schema["image_count"].Unit)
|
||||
assert.Equal(t, "text:UO", option.Icon)
|
||||
assert.Equal(t, []int{1990, 1991}, option.ChannelTypes)
|
||||
assert.Equal(t, "http://localhost:9000", option.BaseURL, "the drawer prefills the normalized plugin default")
|
||||
return
|
||||
}
|
||||
@@ -888,6 +900,39 @@ func TestUploadTaskPluginPreflightConflict(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadTaskPluginLocalizesUnknownMetaField(t *testing.T) {
|
||||
require.NoError(t, i18n.Init())
|
||||
body, err := common.Marshal(map[string]any{
|
||||
"source": `export const meta = { futureField: true };`,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
for _, tc := range []struct {
|
||||
language string
|
||||
message string
|
||||
}{
|
||||
{"en", `Plugin metadata contains an unknown field "futureField". If this plugin was downloaded from the official marketplace, it may require a newer version of new-api. Try updating new-api and installing the plugin again.`},
|
||||
{"zh-CN", "插件元数据包含未知字段“futureField”。如果插件来自官方市场,可能需要更高版本的 new-api。请尝试更新 new-api 后重新安装插件。"},
|
||||
{"zh-TW", "外掛中繼資料包含未知欄位「futureField」。如果外掛來自官方市集,可能需要較新版本的 new-api。請嘗試更新 new-api 後重新安裝外掛。"},
|
||||
} {
|
||||
t.Run(tc.language, func(t *testing.T) {
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task", bytes.NewReader(body))
|
||||
c.Request.Header.Set("Content-Type", "application/json")
|
||||
c.Request.Header.Set("Accept-Language", tc.language)
|
||||
UploadTaskPlugin(c)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
var response struct {
|
||||
Success bool `json:"success"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
assert.False(t, response.Success)
|
||||
assert.Equal(t, tc.message, response.Message)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadTaskPluginRejectsMetaViolatingV1Schema(t *testing.T) {
|
||||
setupTaskPluginControllerTest(t)
|
||||
cases := []struct {
|
||||
|
||||
+152
-33
@@ -32,14 +32,34 @@ func (input *tokenAutoGroupsInput) UnmarshalJSON(data []byte) error {
|
||||
return common.Unmarshal(data, &input.Groups)
|
||||
}
|
||||
|
||||
type tokenGroupPoolInput struct {
|
||||
Set bool
|
||||
Entries []model.TokenGroupEntry
|
||||
}
|
||||
|
||||
func (input *tokenGroupPoolInput) UnmarshalJSON(data []byte) error {
|
||||
input.Set = true
|
||||
if strings.TrimSpace(string(data)) == "null" {
|
||||
input.Entries = nil
|
||||
return nil
|
||||
}
|
||||
return common.Unmarshal(data, &input.Entries)
|
||||
}
|
||||
|
||||
type tokenRequest struct {
|
||||
model.Token
|
||||
AutoGroups tokenAutoGroupsInput `json:"auto_groups"`
|
||||
AstraDetectionEnabled *bool `json:"astra_detection_enabled"`
|
||||
AstraAutoSwitch *bool `json:"astra_auto_switch"`
|
||||
RoutingPolicy *model.TokenRoutingPolicy `json:"routing_policy"`
|
||||
AutoGroups tokenAutoGroupsInput `json:"auto_groups"`
|
||||
GroupPool tokenGroupPoolInput `json:"group_pool"`
|
||||
}
|
||||
|
||||
type tokenResponse struct {
|
||||
*model.Token
|
||||
AutoGroups []string `json:"auto_groups"`
|
||||
RoutingPolicy *model.TokenRoutingPolicy `json:"routing_policy,omitempty"`
|
||||
AutoGroups []string `json:"auto_groups"`
|
||||
GroupPool []model.TokenGroupEntry `json:"group_pool,omitempty"`
|
||||
}
|
||||
|
||||
func maxTokenQuota() int {
|
||||
@@ -58,15 +78,20 @@ func buildMaskedTokenResponse(token *model.Token) *tokenResponse {
|
||||
}
|
||||
maskedToken := *token
|
||||
maskedToken.Key = token.GetMaskedKey()
|
||||
autoGroups, err := token.GetAutoGroups()
|
||||
groupPool, err := token.GetGroupPool()
|
||||
if err != nil {
|
||||
common.SysError(fmt.Sprintf("failed to parse auto groups for token %d: %v", token.Id, err))
|
||||
autoGroups = nil
|
||||
groupPool = nil
|
||||
}
|
||||
autoGroups := make([]string, 0, len(groupPool))
|
||||
for _, entry := range groupPool {
|
||||
autoGroups = append(autoGroups, entry.Group)
|
||||
}
|
||||
if len(autoGroups) == 0 {
|
||||
autoGroups = nil
|
||||
}
|
||||
return &tokenResponse{Token: &maskedToken, AutoGroups: autoGroups}
|
||||
policy, _ := token.GetRoutingPolicy()
|
||||
return &tokenResponse{Token: &maskedToken, AutoGroups: autoGroups, GroupPool: groupPool, RoutingPolicy: policy}
|
||||
}
|
||||
|
||||
func buildMaskedTokenResponses(tokens []*model.Token) []*tokenResponse {
|
||||
@@ -128,18 +153,56 @@ func setTokenAutoGroups(c *gin.Context, token *model.Token, groups []string) boo
|
||||
}
|
||||
|
||||
func setTokenBackupGroup(c *gin.Context, token *model.Token, groups []string) bool {
|
||||
if len(groups) > 1 {
|
||||
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsTooMany, map[string]any{"Max": 1})
|
||||
maxCount := setting.GetMaxTokenAutoGroups()
|
||||
if len(groups) > maxCount {
|
||||
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsTooMany, map[string]any{"Max": maxCount})
|
||||
return false
|
||||
}
|
||||
if len(groups) == 1 && groups[0] == token.Group {
|
||||
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsDuplicate, map[string]any{"Group": groups[0]})
|
||||
return false
|
||||
for _, group := range groups {
|
||||
if group == token.Group {
|
||||
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsDuplicate, map[string]any{"Group": group})
|
||||
return false
|
||||
}
|
||||
}
|
||||
if !setTokenAutoGroups(c, token, groups) {
|
||||
return false
|
||||
}
|
||||
token.CrossGroupRetry = len(groups) == 1
|
||||
token.CrossGroupRetry = len(groups) > 0
|
||||
return true
|
||||
}
|
||||
|
||||
func setTokenGroupPool(c *gin.Context, token *model.Token, entries []model.TokenGroupEntry) bool {
|
||||
maxCount := setting.GetMaxTokenAutoGroups()
|
||||
if len(entries) > maxCount {
|
||||
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsTooMany, map[string]any{"Max": maxCount})
|
||||
return false
|
||||
}
|
||||
userGroup, err := getTokenRequestUserGroup(c)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return false
|
||||
}
|
||||
seen := make(map[string]struct{}, len(entries))
|
||||
for _, entry := range entries {
|
||||
if entry.Group == token.Group {
|
||||
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsDuplicate, map[string]any{"Group": entry.Group})
|
||||
return false
|
||||
}
|
||||
if _, ok := seen[entry.Group]; ok {
|
||||
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsDuplicate, map[string]any{"Group": entry.Group})
|
||||
return false
|
||||
}
|
||||
seen[entry.Group] = struct{}{}
|
||||
if !service.IsUserSelectableGroup(userGroup, entry.Group) {
|
||||
common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsInvalid, map[string]any{"Group": entry.Group})
|
||||
return false
|
||||
}
|
||||
}
|
||||
if err := token.SetGroupPool(entries); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return false
|
||||
}
|
||||
token.CrossGroupRetry = len(entries) > 0
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -299,6 +362,16 @@ func AddToken(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
token := request.Token
|
||||
token.AstraDetectionEnabled = true
|
||||
if request.AstraDetectionEnabled != nil {
|
||||
token.AstraDetectionEnabled = *request.AstraDetectionEnabled
|
||||
}
|
||||
if request.AstraAutoSwitch != nil {
|
||||
token.AstraAutoSwitch = *request.AstraAutoSwitch
|
||||
}
|
||||
if !token.AstraDetectionEnabled {
|
||||
token.AstraAutoSwitch = false
|
||||
}
|
||||
if len(token.Name) > 50 {
|
||||
common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong)
|
||||
return
|
||||
@@ -331,15 +404,36 @@ func AddToken(c *gin.Context) {
|
||||
})
|
||||
return
|
||||
}
|
||||
if token.Group == "auto" {
|
||||
if request.RoutingPolicy != nil {
|
||||
if !setTokenRoutingPolicy(c, &token, request.RoutingPolicy) {
|
||||
return
|
||||
}
|
||||
} else if token.Group == "auto" {
|
||||
if !setTokenAutoGroups(c, &token, request.AutoGroups.Groups) {
|
||||
return
|
||||
}
|
||||
} else if request.GroupPool.Set {
|
||||
if !setTokenGroupPool(c, &token, request.GroupPool.Entries) {
|
||||
return
|
||||
}
|
||||
} else {
|
||||
if !setTokenBackupGroup(c, &token, request.AutoGroups.Groups) {
|
||||
return
|
||||
}
|
||||
}
|
||||
if token.RoutingPolicy == "" {
|
||||
policy, err := token.LegacyRoutingPolicy()
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
body, err := common.Marshal(policy)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
token.RoutingPolicy = string(body)
|
||||
}
|
||||
key, err := common.GenerateKey()
|
||||
if err != nil {
|
||||
common.ApiErrorI18n(c, i18n.MsgTokenGenerateFailed)
|
||||
@@ -347,20 +441,23 @@ func AddToken(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
cleanToken := model.Token{
|
||||
UserId: c.GetInt("id"),
|
||||
Name: token.Name,
|
||||
Key: key,
|
||||
CreatedTime: common.GetTimestamp(),
|
||||
AccessedTime: common.GetTimestamp(),
|
||||
ExpiredTime: token.ExpiredTime,
|
||||
RemainQuota: token.RemainQuota,
|
||||
UnlimitedQuota: token.UnlimitedQuota,
|
||||
ModelLimitsEnabled: token.ModelLimitsEnabled,
|
||||
ModelLimits: token.ModelLimits,
|
||||
AllowIps: token.AllowIps,
|
||||
Group: token.Group,
|
||||
CrossGroupRetry: token.CrossGroupRetry,
|
||||
AutoGroups: token.AutoGroups,
|
||||
UserId: c.GetInt("id"),
|
||||
Name: token.Name,
|
||||
Key: key,
|
||||
CreatedTime: common.GetTimestamp(),
|
||||
AccessedTime: common.GetTimestamp(),
|
||||
ExpiredTime: token.ExpiredTime,
|
||||
RemainQuota: token.RemainQuota,
|
||||
UnlimitedQuota: token.UnlimitedQuota,
|
||||
ModelLimitsEnabled: token.ModelLimitsEnabled,
|
||||
ModelLimits: token.ModelLimits,
|
||||
AllowIps: token.AllowIps,
|
||||
Group: token.Group,
|
||||
CrossGroupRetry: token.CrossGroupRetry,
|
||||
AstraAutoSwitch: token.AstraAutoSwitch,
|
||||
AstraDetectionEnabled: token.AstraDetectionEnabled,
|
||||
AutoGroups: token.AutoGroups,
|
||||
RoutingPolicy: token.RoutingPolicy,
|
||||
}
|
||||
err = cleanToken.Insert()
|
||||
if err != nil {
|
||||
@@ -454,15 +551,34 @@ func UpdateToken(c *gin.Context) {
|
||||
cleanToken.ModelLimitsEnabled = token.ModelLimitsEnabled
|
||||
cleanToken.ModelLimits = token.ModelLimits
|
||||
cleanToken.AllowIps = token.AllowIps
|
||||
cleanToken.Group = token.Group
|
||||
cleanToken.CrossGroupRetry = token.CrossGroupRetry
|
||||
if token.Group != "auto" {
|
||||
if !setTokenBackupGroup(c, cleanToken, request.AutoGroups.Groups) {
|
||||
if request.AstraDetectionEnabled != nil {
|
||||
cleanToken.AstraDetectionEnabled = *request.AstraDetectionEnabled
|
||||
}
|
||||
if request.AstraAutoSwitch != nil {
|
||||
cleanToken.AstraAutoSwitch = *request.AstraAutoSwitch
|
||||
}
|
||||
if !cleanToken.AstraDetectionEnabled {
|
||||
cleanToken.AstraAutoSwitch = false
|
||||
}
|
||||
if request.RoutingPolicy != nil {
|
||||
if !setTokenRoutingPolicy(c, cleanToken, request.RoutingPolicy) {
|
||||
return
|
||||
}
|
||||
} else if request.AutoGroups.Set {
|
||||
if !setTokenAutoGroups(c, cleanToken, request.AutoGroups.Groups) {
|
||||
return
|
||||
} else if cleanToken.RoutingPolicy == "" {
|
||||
cleanToken.Group = token.Group
|
||||
cleanToken.CrossGroupRetry = token.CrossGroupRetry
|
||||
if token.Group != "auto" {
|
||||
if request.GroupPool.Set {
|
||||
if !setTokenGroupPool(c, cleanToken, request.GroupPool.Entries) {
|
||||
return
|
||||
}
|
||||
} else if !setTokenBackupGroup(c, cleanToken, request.AutoGroups.Groups) {
|
||||
return
|
||||
}
|
||||
} else if request.AutoGroups.Set {
|
||||
if !setTokenAutoGroups(c, cleanToken, request.AutoGroups.Groups) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -490,7 +606,10 @@ func UpdateToken(c *gin.Context) {
|
||||
(previous.AllowIps != nil && cleanToken.AllowIps != nil && *previous.AllowIps != *cleanToken.AllowIps)},
|
||||
{"group", previous.Group != cleanToken.Group},
|
||||
{"cross_group_retry", previous.CrossGroupRetry != cleanToken.CrossGroupRetry},
|
||||
{"astra_auto_switch", previous.AstraAutoSwitch != cleanToken.AstraAutoSwitch},
|
||||
{"astra_detection_enabled", previous.AstraDetectionEnabled != cleanToken.AstraDetectionEnabled},
|
||||
{"auto_groups", previous.AutoGroups != cleanToken.AutoGroups},
|
||||
{"routing_policy", previous.RoutingPolicy != cleanToken.RoutingPolicy},
|
||||
} {
|
||||
if field.changed {
|
||||
changedFields = append(changedFields, field.name)
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
package controller
|
||||
|
||||
import (
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func setTokenRoutingPolicy(c *gin.Context, token *model.Token, policy *model.TokenRoutingPolicy) bool {
|
||||
body, err := common.Marshal(policy)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return false
|
||||
}
|
||||
if token.RoutingPolicy != "" && token.RoutingPolicy == string(body) {
|
||||
return true
|
||||
}
|
||||
userGroup, err := getTokenRequestUserGroup(c)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return false
|
||||
}
|
||||
if err := policy.Validate(setting.GetMaxTokenAutoGroups()+1, func(group string) bool { return service.IsUserSelectableGroup(userGroup, group) }); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return false
|
||||
}
|
||||
token.RoutingPolicy = string(body)
|
||||
token.CrossGroupRetry = true
|
||||
token.AutoGroups = ""
|
||||
if policy.Mode == "inherit" {
|
||||
token.Group = "auto"
|
||||
} else if policy.Mode == "user" {
|
||||
token.Group = ""
|
||||
} else {
|
||||
token.Group = policy.Resolve("", nil, func(string) bool { return true })[0].Group
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func PreviewTokenRouting(c *gin.Context) {
|
||||
var request struct {
|
||||
Policy model.TokenRoutingPolicy `json:"routing_policy"`
|
||||
Model string `json:"model"`
|
||||
}
|
||||
if err := c.ShouldBindJSON(&request); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
if err := request.Policy.Validate(setting.GetMaxTokenAutoGroups()+1, func(string) bool { return true }); err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
group, err := getTokenRequestUserGroup(c)
|
||||
if err != nil {
|
||||
common.ApiError(c, err)
|
||||
return
|
||||
}
|
||||
inherited := service.GetUserAutoGroup(group)
|
||||
if request.Policy.Mode == "user" {
|
||||
inherited = []string{group}
|
||||
}
|
||||
entries := request.Policy.Resolve(request.Model, inherited, func(candidate string) bool { return service.IsUserSelectableGroup(group, candidate) })
|
||||
common.ApiSuccess(c, entries)
|
||||
}
|
||||
@@ -118,6 +118,91 @@ func setupTokenControllerTestDB(t *testing.T) *gorm.DB {
|
||||
return db
|
||||
}
|
||||
|
||||
func TestTokenAstraMigration(t *testing.T) {
|
||||
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
|
||||
for _, upgrade := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("%s/upgrade=%t", dialect, upgrade), func(t *testing.T) {
|
||||
var db *gorm.DB
|
||||
if dialect == "sqlite" {
|
||||
db = openTokenControllerTestDB(t)
|
||||
} else {
|
||||
dsn := os.Getenv("TEST_" + strings.ToUpper(dialect) + "_DSN")
|
||||
if dsn == "" {
|
||||
t.Skip("external test database is not configured")
|
||||
}
|
||||
var managed *bool
|
||||
db, managed = openTokenControllerExternalDB(t, dialect, dsn)
|
||||
*managed = true
|
||||
if db.Migrator().HasTable("options") {
|
||||
t.Fatal("requires an empty test database")
|
||||
}
|
||||
}
|
||||
t.Cleanup(func() { require.NoError(t, db.Migrator().DropTable(&model.Option{})) })
|
||||
if upgrade {
|
||||
require.NoError(t, db.AutoMigrate(&legacyToken{}))
|
||||
require.NoError(t, db.Create(&legacyToken{UserId: 7, Key: "legacy-astra", Name: "preserved", RemainQuota: 123}).Error)
|
||||
}
|
||||
require.NoError(t, db.AutoMigrate(&model.Token{}, &model.Option{}))
|
||||
require.NoError(t, model.MigrateTokenAstraAutoSwitch())
|
||||
require.NoError(t, model.MigrateTokenAstraDetection())
|
||||
if upgrade {
|
||||
var legacy model.Token
|
||||
require.NoError(t, db.First(&legacy, 1).Error)
|
||||
require.True(t, legacy.AstraAutoSwitch)
|
||||
require.True(t, legacy.AstraDetectionEnabled)
|
||||
require.Equal(t, 123, legacy.RemainQuota)
|
||||
legacy.AstraAutoSwitch = false
|
||||
legacy.AstraDetectionEnabled = false
|
||||
require.NoError(t, legacy.Update())
|
||||
}
|
||||
fresh := model.Token{UserId: 7, Key: "new-astra"}
|
||||
require.NoError(t, db.Create(&fresh).Error)
|
||||
for range 2 {
|
||||
require.NoError(t, db.AutoMigrate(&model.Token{}, &model.Option{}))
|
||||
require.NoError(t, model.MigrateTokenAstraAutoSwitch())
|
||||
require.NoError(t, model.MigrateTokenAstraDetection())
|
||||
}
|
||||
var tokens []model.Token
|
||||
require.NoError(t, db.Find(&tokens).Error)
|
||||
for _, token := range tokens {
|
||||
require.False(t, token.AstraAutoSwitch, "restart must preserve opt-outs and new-token defaults")
|
||||
require.False(t, token.AstraDetectionEnabled)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenAstraSwitchAPIPreservesOmittedAndPersistsFalse(t *testing.T) {
|
||||
db := setupTokenControllerTestDB(t)
|
||||
token := seedToken(t, db, 7, "astra", "astra-api-key")
|
||||
token.AstraAutoSwitch = true
|
||||
token.AstraDetectionEnabled = true
|
||||
require.NoError(t, token.Update())
|
||||
for _, enabled := range []any{nil, false, true} {
|
||||
body := map[string]any{"id": token.Id, "name": "astra", "unlimited_quota": true, "group": "default", "expired_time": -1}
|
||||
if enabled != nil {
|
||||
body["astra_auto_switch"] = enabled
|
||||
}
|
||||
c, writer := newAuthenticatedContext(t, http.MethodPut, "/api/token/", body, 7)
|
||||
UpdateToken(c)
|
||||
require.True(t, decodeAPIResponse(t, writer).Success, writer.Body.String())
|
||||
require.NoError(t, db.First(token, token.Id).Error)
|
||||
want := enabled == nil || enabled == true
|
||||
require.Equal(t, want, token.AstraAutoSwitch)
|
||||
var response struct {
|
||||
AstraAutoSwitch bool `json:"astra_auto_switch"`
|
||||
}
|
||||
require.NoError(t, common.Unmarshal(decodeAPIResponse(t, writer).Data, &response))
|
||||
require.Equal(t, want, response.AstraAutoSwitch)
|
||||
}
|
||||
c, writer := newAuthenticatedContext(t, http.MethodPut, "/api/token/", map[string]any{"id": token.Id, "name": "foreign", "astra_auto_switch": false}, 8)
|
||||
UpdateToken(c)
|
||||
require.False(t, decodeAPIResponse(t, writer).Success)
|
||||
require.NoError(t, db.First(token, token.Id).Error)
|
||||
require.True(t, token.AstraAutoSwitch)
|
||||
}
|
||||
|
||||
func openTokenControllerExternalDB(t *testing.T, dialect string, dsn string) (*gorm.DB, *bool) {
|
||||
t.Helper()
|
||||
|
||||
@@ -892,3 +977,22 @@ func verifyAPITokenAudit(t *testing.T) {
|
||||
assert.EqualValues(t, 1, count)
|
||||
})
|
||||
}
|
||||
|
||||
func TestTokenDetectionOffDisablesSwitchAndSurvivesOmittedUpdate(t *testing.T) {
|
||||
db := setupTokenControllerTestDB(t)
|
||||
token := seedToken(t, db, 7, "detection", "detection-api-key")
|
||||
token.AstraDetectionEnabled, token.AstraAutoSwitch = true, true
|
||||
require.NoError(t, token.Update())
|
||||
for _, detection := range []any{false, nil, true} {
|
||||
body := map[string]any{"id": token.Id, "name": "detection", "unlimited_quota": true, "group": "default", "expired_time": -1}
|
||||
if detection != nil {
|
||||
body["astra_detection_enabled"] = detection
|
||||
}
|
||||
c, writer := newAuthenticatedContext(t, http.MethodPut, "/api/token/", body, 7)
|
||||
UpdateToken(c)
|
||||
require.True(t, decodeAPIResponse(t, writer).Success, writer.Body.String())
|
||||
require.NoError(t, db.First(token, token.Id).Error)
|
||||
require.Equal(t, detection == true, token.AstraDetectionEnabled)
|
||||
require.False(t, token.AstraAutoSwitch)
|
||||
}
|
||||
}
|
||||
+83
-146
@@ -2,9 +2,7 @@ package controller
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -15,29 +13,31 @@ import (
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"golang.org/x/sync/errgroup"
|
||||
)
|
||||
|
||||
const (
|
||||
requestTimeout = 30 * time.Second
|
||||
httpTimeout = 10 * time.Second
|
||||
uptimeKeySuffix = "_24"
|
||||
apiStatusPath = "/api/status-page/"
|
||||
apiHeartbeatPath = "/api/status-page/heartbeat/"
|
||||
)
|
||||
const requestTimeout = 30 * time.Second
|
||||
|
||||
type Monitor struct {
|
||||
Metrics *perfmetrics.GroupResult `json:"metrics,omitempty"`
|
||||
Name string `json:"name"`
|
||||
Uptime float64 `json:"uptime"`
|
||||
Status int `json:"status"`
|
||||
Group string `json:"group,omitempty"`
|
||||
Ping int `json:"ping,omitempty"`
|
||||
LastChecked string `json:"last_checked,omitempty"`
|
||||
History []MonitorSample `json:"history,omitempty"`
|
||||
CacheStats *model.GroupCacheStats `json:"cache_stats,omitempty"`
|
||||
ProbeStatus int `json:"probe_status"`
|
||||
ChannelDisabled bool `json:"channel_disabled"`
|
||||
ProbeHistory []service.ProbeHistoryPoint `json:"probe_history"`
|
||||
ID string `json:"id"`
|
||||
Models []string `json:"models"`
|
||||
Alias string `json:"alias,omitempty"`
|
||||
Hidden bool `json:"hidden,omitempty"`
|
||||
Pelican *service.PelicanSummary `json:"pelican,omitempty"`
|
||||
StatusSource string `json:"status_source"`
|
||||
StatusCheckedAt int64 `json:"status_checked_at,omitempty"`
|
||||
Fingerprint *service.AstraSamplingSummary `json:"fingerprint,omitempty"`
|
||||
Metrics *perfmetrics.GroupResult `json:"metrics,omitempty"`
|
||||
Name string `json:"name"`
|
||||
Uptime float64 `json:"uptime"`
|
||||
Status int `json:"status"`
|
||||
Group string `json:"group,omitempty"`
|
||||
Ping int `json:"ping,omitempty"`
|
||||
LastChecked string `json:"last_checked,omitempty"`
|
||||
History []MonitorSample `json:"history,omitempty"`
|
||||
CacheStats *model.GroupCacheStats `json:"cache_stats,omitempty"`
|
||||
ProbeStatus int `json:"probe_status"`
|
||||
ChannelDisabled bool `json:"channel_disabled"`
|
||||
}
|
||||
|
||||
type MonitorSample struct {
|
||||
@@ -51,113 +51,6 @@ type UptimeGroupResult struct {
|
||||
Monitors []Monitor `json:"monitors"`
|
||||
}
|
||||
|
||||
func getAndDecode(ctx context.Context, client *http.Client, url string, dest any) error {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return errors.New("non-200 status")
|
||||
}
|
||||
|
||||
return common.DecodeJson(resp.Body, dest)
|
||||
}
|
||||
|
||||
func fetchGroupData(ctx context.Context, client *http.Client, groupConfig map[string]any) UptimeGroupResult {
|
||||
url, _ := groupConfig["url"].(string)
|
||||
slug, _ := groupConfig["slug"].(string)
|
||||
categoryName, _ := groupConfig["categoryName"].(string)
|
||||
|
||||
result := UptimeGroupResult{
|
||||
CategoryName: categoryName,
|
||||
Monitors: []Monitor{},
|
||||
}
|
||||
|
||||
if url == "" || slug == "" {
|
||||
return result
|
||||
}
|
||||
|
||||
baseURL := strings.TrimSuffix(url, "/")
|
||||
|
||||
var statusData struct {
|
||||
PublicGroupList []struct {
|
||||
ID int `json:"id"`
|
||||
Name string `json:"name"`
|
||||
MonitorList []struct {
|
||||
ID int `json:"id"`
|
||||
Name string `json:"name"`
|
||||
} `json:"monitorList"`
|
||||
} `json:"publicGroupList"`
|
||||
}
|
||||
|
||||
var heartbeatData struct {
|
||||
HeartbeatList map[string][]struct {
|
||||
Status int `json:"status"`
|
||||
Ping int `json:"ping"`
|
||||
Time string `json:"time"`
|
||||
} `json:"heartbeatList"`
|
||||
UptimeList map[string]float64 `json:"uptimeList"`
|
||||
}
|
||||
|
||||
g, gCtx := errgroup.WithContext(ctx)
|
||||
g.Go(func() error {
|
||||
return getAndDecode(gCtx, client, baseURL+apiStatusPath+slug, &statusData)
|
||||
})
|
||||
g.Go(func() error {
|
||||
return getAndDecode(gCtx, client, baseURL+apiHeartbeatPath+slug, &heartbeatData)
|
||||
})
|
||||
|
||||
if g.Wait() != nil {
|
||||
return result
|
||||
}
|
||||
|
||||
for _, pg := range statusData.PublicGroupList {
|
||||
if len(pg.MonitorList) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
for _, m := range pg.MonitorList {
|
||||
monitor := Monitor{
|
||||
Name: m.Name,
|
||||
Group: pg.Name,
|
||||
Status: -1,
|
||||
}
|
||||
|
||||
monitorID := strconv.Itoa(m.ID)
|
||||
|
||||
if uptime, exists := heartbeatData.UptimeList[monitorID+uptimeKeySuffix]; exists {
|
||||
monitor.Uptime = uptime
|
||||
}
|
||||
|
||||
if heartbeats, exists := heartbeatData.HeartbeatList[monitorID]; exists && len(heartbeats) > 0 {
|
||||
latest := heartbeats[len(heartbeats)-1]
|
||||
monitor.Status = latest.Status
|
||||
monitor.Ping = latest.Ping
|
||||
monitor.LastChecked = latest.Time
|
||||
monitor.History = make([]MonitorSample, 0, len(heartbeats))
|
||||
for _, heartbeat := range heartbeats {
|
||||
monitor.History = append(monitor.History, MonitorSample{
|
||||
Status: heartbeat.Status,
|
||||
Ping: heartbeat.Ping,
|
||||
Time: heartbeat.Time,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
result.Monitors = append(result.Monitors, monitor)
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
var groupCacheSnapshot struct {
|
||||
sync.Mutex
|
||||
updated time.Time
|
||||
@@ -168,18 +61,35 @@ var groupCacheSnapshot struct {
|
||||
|
||||
const requestStatusFreshness = 2 * time.Hour
|
||||
|
||||
func currentGroupMonitors(ctx context.Context, results []UptimeGroupResult) ([]UptimeGroupResult, error) {
|
||||
func currentGroupMonitors(ctx context.Context, admin bool) ([]UptimeGroupResult, error) {
|
||||
channels, err := model.GetStatusChannels(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
settings, err := model.GetStatusPageGroups(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
models, err := model.GetStatusPageModels(ctx, channels)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
usable := service.GetUserUsableGroups("")
|
||||
now := time.Now()
|
||||
fingerprints, fingerprintErr := service.AstraSamplingSummaries(ctx, now, settings)
|
||||
if fingerprintErr != nil {
|
||||
return nil, fingerprintErr
|
||||
}
|
||||
detections, detectionErr := service.GroupDetectionSummaries(ctx, now, fingerprints, settings)
|
||||
if detectionErr != nil {
|
||||
return nil, detectionErr
|
||||
}
|
||||
active := make(map[string]bool)
|
||||
order := make([]string, 0)
|
||||
for _, channel := range channels {
|
||||
for _, group := range strings.Split(channel.Group, ",") {
|
||||
group = strings.TrimSpace(group)
|
||||
if _, visible := usable[group]; !visible {
|
||||
if _, visible := usable[group]; (!visible || settings[group].Hidden) && !admin {
|
||||
continue
|
||||
}
|
||||
if _, exists := active[group]; !exists {
|
||||
@@ -202,8 +112,8 @@ func currentGroupMonitors(ctx context.Context, results []UptimeGroupResult) ([]U
|
||||
}
|
||||
groupCacheSnapshot.stats = stats
|
||||
names := []string{}
|
||||
for _, item := range model.GetPricing() {
|
||||
names = append(names, item.ModelName)
|
||||
if err := model.DB.WithContext(ctx).Model(&model.Ability{}).Where("enabled = ?", true).Distinct().Pluck("model", &names).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
performance, metricsErr := perfmetrics.QueryGroups(names, 24)
|
||||
if metricsErr != nil {
|
||||
@@ -216,21 +126,26 @@ func currentGroupMonitors(ctx context.Context, results []UptimeGroupResult) ([]U
|
||||
groupCacheSnapshot.updated = now
|
||||
groupCacheSnapshot.channels = string(encoded)
|
||||
}
|
||||
byName := make(map[string]Monitor)
|
||||
for _, result := range results {
|
||||
for _, monitor := range result.Monitors {
|
||||
if previous, exists := byName[monitor.Name]; !exists || (monitor.Status == 1 && previous.Status != 1) {
|
||||
byName[monitor.Name] = monitor
|
||||
}
|
||||
}
|
||||
}
|
||||
successes := perfmetrics.LatestUserSuccess(ctx, order)
|
||||
monitors := make([]Monitor, 0, len(order))
|
||||
for _, group := range order {
|
||||
monitor, exists := byName[group]
|
||||
if !exists {
|
||||
monitor = Monitor{Name: group, Status: -1}
|
||||
monitor := Monitor{ID: model.DetectionSecretHash(group)[:16], Name: group, Status: -1, Models: models[group]}
|
||||
if monitor.Models == nil {
|
||||
monitor.Models = []string{}
|
||||
}
|
||||
if settings[group].Alias != "" {
|
||||
monitor.Name = settings[group].Alias
|
||||
}
|
||||
if admin {
|
||||
monitor.Alias, monitor.Hidden = settings[group].Alias, settings[group].Hidden
|
||||
}
|
||||
if admin {
|
||||
monitor.Group = group
|
||||
}
|
||||
monitor.Fingerprint = fingerprints[group]
|
||||
if monitor.Fingerprint == nil {
|
||||
monitor.Fingerprint = &service.AstraSamplingSummary{Enabled: service.AstraSamplingEnabled(), Available: fingerprintErr == nil, History: []service.AstraSamplingHistory{}}
|
||||
}
|
||||
monitor.Group = group
|
||||
monitor.ProbeStatus = monitor.Status
|
||||
monitor.ChannelDisabled = !active[group]
|
||||
if !active[group] {
|
||||
@@ -239,16 +154,30 @@ func currentGroupMonitors(ctx context.Context, results []UptimeGroupResult) ([]U
|
||||
stats := groupCacheSnapshot.stats[group]
|
||||
monitor.CacheStats = &stats
|
||||
applyRequestMetrics(&monitor, groupCacheSnapshot.metrics[group])
|
||||
if !admin && monitor.Metrics != nil {
|
||||
monitor.Metrics.Group = monitor.Name
|
||||
}
|
||||
monitor.Pelican = detections[group].Pelican
|
||||
monitor.StatusSource = "requests"
|
||||
if monitor.ChannelDisabled {
|
||||
monitor.StatusSource = "disabled"
|
||||
} else if monitor.Status < 0 {
|
||||
monitor.Status = detections[group].Status
|
||||
monitor.StatusSource = "detection"
|
||||
monitor.StatusCheckedAt = detections[group].At
|
||||
}
|
||||
monitor.ProbeHistory = detections[group].ProbeHistory
|
||||
applyUserSuccess(&monitor, successes[group], detections[group].At, now.Unix())
|
||||
monitors = append(monitors, monitor)
|
||||
}
|
||||
return []UptimeGroupResult{{CategoryName: "Groups", Monitors: monitors}}, nil
|
||||
}
|
||||
|
||||
func GetUptimeKumaStatus(c *gin.Context) {
|
||||
func GetGroupStatus(c *gin.Context) {
|
||||
ctx, cancel := context.WithTimeout(c.Request.Context(), requestTimeout)
|
||||
defer cancel()
|
||||
|
||||
results, err := currentGroupMonitors(ctx, nil)
|
||||
results, err := currentGroupMonitors(ctx, false)
|
||||
if err != nil {
|
||||
common.SysError("group status statistics: " + err.Error())
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"success": false, "message": "Status data is temporarily unavailable"})
|
||||
@@ -267,6 +196,14 @@ func requestStatus(rate float64) int {
|
||||
return 0
|
||||
}
|
||||
|
||||
func applyUserSuccess(monitor *Monitor, success, probe, now int64) {
|
||||
if !monitor.ChannelDisabled && success > 0 && success >= now-900 && success >= probe {
|
||||
monitor.Status = 1
|
||||
monitor.StatusSource = "requests"
|
||||
monitor.StatusCheckedAt = success
|
||||
}
|
||||
}
|
||||
|
||||
func applyRequestMetrics(monitor *Monitor, metrics perfmetrics.GroupResult) {
|
||||
monitor.Status = -1
|
||||
monitor.Uptime = 0
|
||||
|
||||
+191
-31
@@ -3,39 +3,199 @@ package controller
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/middleware"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
|
||||
"github.com/QuantumNous/new-api/service"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/glebarez/sqlite"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"gorm.io/driver/mysql"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/schema"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestFetchGroupDataUsesLatestHeartbeat(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/api/status-page/status":
|
||||
_, _ = fmt.Fprint(w, `{"publicGroupList":[{"id":1,"name":"hidden","monitorList":[{"id":7,"name":"GPT Mix"}]}]}`)
|
||||
case "/api/status-page/heartbeat/status":
|
||||
_, _ = fmt.Fprint(w, `{"heartbeatList":{"7":[{"status":0,"time":"2026-09-05 10:00:00","ping":0},{"status":1,"time":"2026-09-05 10:03:00","ping":321}]},"uptimeList":{"7_24":0.95}}`)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
t.Cleanup(server.Close)
|
||||
|
||||
result := fetchGroupData(context.Background(), server.Client(), map[string]interface{}{
|
||||
"url": server.URL,
|
||||
"slug": "status",
|
||||
"categoryName": "API",
|
||||
})
|
||||
|
||||
require.Len(t, result.Monitors, 1)
|
||||
monitor := result.Monitors[0]
|
||||
assert.Equal(t, "GPT Mix", monitor.Name)
|
||||
assert.Equal(t, 1, monitor.Status)
|
||||
assert.Equal(t, 321, monitor.Ping)
|
||||
assert.Equal(t, "2026-09-05 10:03:00", monitor.LastChecked)
|
||||
assert.Equal(t, 0.95, monitor.Uptime)
|
||||
assert.Len(t, monitor.History, 2)
|
||||
func TestStatusPageDatabaseMatrix(t *testing.T) {
|
||||
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
|
||||
t.Run(dialect, func(t *testing.T) {
|
||||
var driver gorm.Dialector
|
||||
switch dialect {
|
||||
case "sqlite":
|
||||
driver = sqlite.Open(":memory:")
|
||||
case "mysql":
|
||||
if os.Getenv("TEST_MYSQL_DSN") == "" {
|
||||
t.Skip("TEST_MYSQL_DSN unset")
|
||||
}
|
||||
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
|
||||
case "postgres":
|
||||
if os.Getenv("TEST_POSTGRES_DSN") == "" {
|
||||
t.Skip("TEST_POSTGRES_DSN unset")
|
||||
}
|
||||
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
|
||||
}
|
||||
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("status_%d_", time.Now().UnixNano())}})
|
||||
require.NoError(t, err)
|
||||
conn, err := db.DB()
|
||||
require.NoError(t, err)
|
||||
conn.SetMaxOpenConns(1)
|
||||
oldDB, oldLog := model.DB, model.LOG_DB
|
||||
oldType, oldLogType := common.MainDatabaseType(), common.LogDatabaseType()
|
||||
oldRedis, oldSecret, oldSecure := common.RedisEnabled, common.SessionSecret, common.SessionCookieSecure
|
||||
oldGroups := setting.UserUsableGroups2JSONString()
|
||||
model.DB, model.LOG_DB = db, db
|
||||
common.SetDatabaseTypes(common.DatabaseType(dialect), common.DatabaseType(dialect))
|
||||
common.RedisEnabled, common.SessionSecret, common.SessionCookieSecure = false, "status-test-only", true
|
||||
rows := []any{&model.Option{}, &model.Channel{}, &model.Ability{}, &model.SystemTask{}, &model.DetectionRun{}, &model.DetectionCall{}, &model.DetectionSchedule{}, &model.Log{}, &model.PerfMetric{}, &model.User{}, &model.UserSession{}}
|
||||
require.NoError(t, db.AutoMigrate(rows...))
|
||||
require.NoError(t, db.Table("audit_logs").AutoMigrate(&model.AuditLog{}))
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, db.Migrator().DropTable("audit_logs"))
|
||||
for _, row := range rows {
|
||||
require.NoError(t, db.Migrator().DropTable(row))
|
||||
}
|
||||
model.DB, model.LOG_DB = oldDB, oldLog
|
||||
common.SetDatabaseTypes(oldType, oldLogType)
|
||||
common.RedisEnabled, common.SessionSecret, common.SessionCookieSecure = oldRedis, oldSecret, oldSecure
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(oldGroups))
|
||||
groupCacheSnapshot.updated = time.Time{}
|
||||
require.NoError(t, conn.Close())
|
||||
})
|
||||
var version string
|
||||
q := "select version()"
|
||||
if dialect == "sqlite" {
|
||||
q = "select sqlite_version()"
|
||||
}
|
||||
require.NoError(t, db.Raw(q).Scan(&version).Error)
|
||||
t.Log("database version:", version)
|
||||
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"internal":"","shared":"","hidden":""}`))
|
||||
now := time.Now().Unix()
|
||||
channels := []model.Channel{{Id: 1, Name: "private-channel", Key: "private-key", Group: "internal,shared", Models: fingerprint.Model, Status: 1}, {Id: 2, Group: "hidden", Models: fingerprint.Model, Status: 1}, {Id: 3, Group: "internal", Models: "disabled-model", Status: 2}}
|
||||
require.NoError(t, db.Create(&channels).Error)
|
||||
abilities := []model.Ability{{Group: "internal", Model: fingerprint.Model, ChannelId: 1, Enabled: true}, {Group: "shared", Model: fingerprint.Model, ChannelId: 1, Enabled: true}, {Group: "hidden", Model: fingerprint.Model, ChannelId: 2, Enabled: true}, {Group: "internal", Model: "disabled-model", ChannelId: 3, Enabled: true}, {Group: "internal", Model: "disabled-ability", ChannelId: 1}}
|
||||
require.NoError(t, db.Create(&abilities).Error)
|
||||
state, err := common.Marshal(service.AstraSamplingState{Rounds: []service.AstraSamplingRound{{ID: "old-round", ChannelID: 1, StartedAt: now - 20, FinishedAt: now - 10, Result: fingerprint.Result{Status: "match", Closest: fingerprint.Model, ValidSamples: 3}}}})
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, db.Create(&model.SystemTask{TaskID: "old-task", Type: service.AstraSamplingTaskType, Status: model.SystemTaskStatusSucceeded, CreatedAt: now - 20, State: string(state)}).Error)
|
||||
runs := []model.DetectionRun{{ID: "old-work", RequestKey: "old-work", ChannelID: 1, GroupName: "internal", Kind: "pelican", Status: "succeeded", Quality: "normal", CreatedAt: now - 20, Settled: true, HTML: "<h1>art</h1>", Charged: 17}, {ID: "in-flight", RequestKey: "in-flight", ChannelID: 1, Kind: "pelican", Status: "running", CreatedAt: now - 10}, {ID: "hidden-art", RequestKey: "hidden-art", ChannelID: 2, Kind: "pelican", Status: "succeeded", CreatedAt: now - 20, HTML: "hidden"}}
|
||||
require.NoError(t, db.Create(&runs).Error)
|
||||
tokens := map[int]string{}
|
||||
for _, role := range []int{1, 10} {
|
||||
require.NoError(t, db.Create(&model.User{Id: role, Username: fmt.Sprintf("role-%d", role), Role: role, Status: 1, AuthVersion: 1, AffCode: fmt.Sprintf("role-%d", role)}).Error)
|
||||
sid := fmt.Sprintf("session-%d", role)
|
||||
require.NoError(t, db.Create(&model.UserSession{SID: sid, UserID: role, Version: 1, UserAuthVersion: 1, Status: model.UserSessionStatusActive, RefreshHash: sid, ExpiresAt: now + 3600}).Error)
|
||||
tokens[role], _, err = service.IssueAccessToken(service.AuthIdentity{UserID: role, SessionID: sid, UserAuthVersion: 1, SessionVersion: 1})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
engine := gin.New()
|
||||
engine.GET("/api/uptime/status", GetGroupStatus)
|
||||
engine.GET("/api/group-status", GetSiteGroupStatus)
|
||||
engine.PUT("/api/uptime/group", middleware.AdminAuth(), middleware.SessionCookieOriginGuard(), UpdateStatusPageGroup)
|
||||
engine.POST("/api/uptime/group/reset", middleware.AdminAuth(), middleware.SessionCookieOriginGuard(), UpdateStatusPageGroup)
|
||||
engine.GET("/api/detections/artwork/:id", GetDetectionArtwork)
|
||||
engine.GET("/api/group-status/artwork/:id", GetSiteDetectionArtwork)
|
||||
request := func(method, path, body, token, origin string) *httptest.ResponseRecorder {
|
||||
req := httptest.NewRequest(method, "https://status.test"+path, strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if token != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
}
|
||||
if origin != "" {
|
||||
req.Header.Set("Origin", origin)
|
||||
}
|
||||
out := httptest.NewRecorder()
|
||||
engine.ServeHTTP(out, req)
|
||||
return out
|
||||
}
|
||||
for _, token := range []string{"", tokens[1]} {
|
||||
assert.Contains(t, []int{401, 403}, request("POST", "/api/uptime/group/reset", `{"group":"internal"}`, token, "https://status.test").Code)
|
||||
}
|
||||
assert.Equal(t, 403, request("POST", "/api/uptime/group/reset", `{"group":"internal"}`, tokens[10], "https://evil.test").Code)
|
||||
assert.Equal(t, 400, request("PUT", "/api/uptime/group", `{"group":"internal"}`, tokens[10], "https://status.test").Code)
|
||||
assert.Equal(t, 404, request("POST", "/api/uptime/group/reset", `{"group":"missing"}`, tokens[10], "https://status.test").Code)
|
||||
for _, body := range []string{`{"group":"internal","alias":"#688","hidden":false}`, `{"group":"hidden","alias":"secret-alias","hidden":true}`} {
|
||||
res := request("PUT", "/api/uptime/group", body, tokens[10], "https://status.test")
|
||||
require.Equal(t, 200, res.Code, res.Body.String())
|
||||
}
|
||||
var auditCount int64
|
||||
require.NoError(t, db.Table("audit_logs").Count(&auditCount).Error)
|
||||
assert.Greater(t, auditCount, int64(0))
|
||||
groupCacheSnapshot.updated = time.Time{}
|
||||
res := request("GET", "/api/uptime/status", "", "", "")
|
||||
require.Equal(t, 200, res.Code, res.Body.String())
|
||||
for _, secret := range []string{"internal", "secret-alias", "private-channel", "private-key", "disabled-model", "disabled-ability"} {
|
||||
assert.NotContains(t, res.Body.String(), secret)
|
||||
}
|
||||
assert.Contains(t, res.Body.String(), "#688")
|
||||
fallback := model.DetectionRun{ID:"fallback-completed",RequestKey:"fallback-completed",Kind:"pelican",ChannelID:1,GroupName:"internal",Status:"running",Quality:"normal",HTML:"<html>pelican</html>",CreatedAt:now,Deadline:now+600}
|
||||
require.NoError(t,db.Create(&fallback).Error)
|
||||
require.NoError(t,db.Create(&model.DetectionCall{ID:fallback.ID,RunID:fallback.ID,SecretHash:"fallback-only-test",State:"finished"}).Error)
|
||||
processLocalPelican(context.Background(),fallback)
|
||||
require.NoError(t,db.First(&fallback,"id = ?",fallback.ID).Error)
|
||||
assert.Equal(t,"succeeded",fallback.Status)
|
||||
assert.Equal(t,"unknown",fallback.Quality)
|
||||
assert.Equal(t,"<html>pelican</html>",fallback.HTML)
|
||||
assert.True(t,fallback.Settled)
|
||||
require.NoError(t,db.Delete(&model.DetectionCall{},"id = ?",fallback.ID).Error)
|
||||
require.NoError(t,db.Delete(&model.DetectionRun{},"id = ?",fallback.ID).Error)
|
||||
visitorSite := request("GET", "/api/group-status", "", "", "")
|
||||
require.Equal(t, 200, visitorSite.Code)
|
||||
assert.Contains(t, visitorSite.Body.String(), `"group":"hidden"`)
|
||||
assert.Contains(t, visitorSite.Body.String(), `"name":"internal"`)
|
||||
assert.NotContains(t, visitorSite.Body.String(), "#688")
|
||||
site := request("GET", "/api/group-status", "", tokens[1], "")
|
||||
require.Equal(t, 200, site.Code)
|
||||
assert.Contains(t, site.Body.String(), `"group":"internal"`)
|
||||
assert.Contains(t, site.Body.String(), `"name":"internal"`)
|
||||
assert.NotContains(t, site.Body.String(), "#688")
|
||||
assert.Contains(t, site.Body.String(), `"group":"hidden"`)
|
||||
adminSite := request("GET", "/api/group-status", "", tokens[10], "")
|
||||
assert.Contains(t, adminSite.Body.String(), `"group":"hidden"`)
|
||||
assert.Equal(t, 404, request("GET", "/api/detections/artwork/hidden-art", "", "", "").Code)
|
||||
assert.Equal(t, 200, request("GET", "/api/group-status/artwork/hidden-art", "", "", "").Code)
|
||||
require.Equal(t, 200, request("POST", "/api/uptime/group/reset", `{"group":"internal"}`, tokens[10], "https://status.test").Code)
|
||||
require.NoError(t, db.Model(&model.DetectionRun{}).Where("id = ?", "in-flight").Updates(map[string]any{"settled": true, "status": "succeeded", "quality": "normal", "finished_at": now + 10}).Error)
|
||||
groups, err := currentGroupMonitors(context.Background(), true)
|
||||
require.NoError(t, err)
|
||||
byGroup := map[string]Monitor{}
|
||||
for _, m := range groups[0].Monitors {
|
||||
byGroup[m.Group] = m
|
||||
}
|
||||
assert.Empty(t, byGroup["internal"].Fingerprint.History)
|
||||
assert.Nil(t, byGroup["internal"].Fingerprint.Rate)
|
||||
assert.Empty(t, byGroup["internal"].Pelican.History)
|
||||
assert.Len(t, byGroup["shared"].Fingerprint.History, 1)
|
||||
assert.Len(t, byGroup["shared"].Pelican.History, 2)
|
||||
assert.True(t, byGroup["hidden"].Hidden)
|
||||
assert.Equal(t, []string{fingerprint.Model}, byGroup["internal"].Models)
|
||||
require.NoError(t, model.SaveStatusPageGroup(context.Background(), "internal", &model.StatusPageGroup{Alias: "#689", Hidden: true}))
|
||||
settings, err := model.GetStatusPageGroups(context.Background())
|
||||
require.NoError(t, err)
|
||||
assert.GreaterOrEqual(t, settings["internal"].ResetBefore, now)
|
||||
assert.Equal(t, "#689", settings["internal"].Alias)
|
||||
require.NoError(t, model.SaveStatusPageGroup(context.Background(), "shared", &model.StatusPageGroup{Hidden: true}))
|
||||
assert.Equal(t, 404, request("GET", "/api/detections/artwork/old-work", "", "", "").Code)
|
||||
var retained model.DetectionRun
|
||||
require.NoError(t, db.First(&retained, "id = ?", "old-work").Error)
|
||||
assert.Equal(t, 17, retained.Charged)
|
||||
assert.Equal(t, "<h1>art</h1>", retained.HTML)
|
||||
require.NoError(t, db.Create(&model.DetectionRun{ID: "new-work", RequestKey: "new-work", ChannelID: 1, Kind: "pelican", Status: "succeeded", Quality: "normal", CreatedAt: settings["internal"].ResetBefore + 1, Settled: true}).Error)
|
||||
groups, err = currentGroupMonitors(context.Background(), true)
|
||||
require.NoError(t, err)
|
||||
for _, m := range groups[0].Monitors {
|
||||
if m.Group == "internal" {
|
||||
require.Len(t, m.Pelican.History, 1)
|
||||
assert.Equal(t, "new-work", m.Pelican.History[0].ID)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
+11
-4
@@ -96,7 +96,7 @@ func Login(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
setupLogin(&user, c)
|
||||
setupLogin(&user, nil, c)
|
||||
}
|
||||
|
||||
// loginMethodFromContext 根据请求路径推导登录方式,用于登录审计日志。
|
||||
@@ -144,9 +144,11 @@ func recordLoginAudit(user *model.User, c *gin.Context) {
|
||||
}
|
||||
|
||||
// setupLogin evaluates the shared login policy after primary authentication.
|
||||
// Only a completed Passkey ceremony may go directly to session issuance.
|
||||
func setupLogin(user *model.User, c *gin.Context) {
|
||||
challenge, err := service.StartLoginVerification(user, loginMethodFromContext(c))
|
||||
// Only a completed Passkey ceremony may go directly to session issuance. A
|
||||
// pending legacy GitHub binding rewrite travels inside the challenge and is
|
||||
// written only when the verification completes.
|
||||
func setupLogin(user *model.User, migration *service.LegacyGitHubMigration, c *gin.Context) {
|
||||
challenge, err := service.StartLoginVerification(user, loginMethodFromContext(c), migration)
|
||||
if err != nil {
|
||||
writeSecurityOperationError(c, err)
|
||||
return
|
||||
@@ -1278,6 +1280,7 @@ type UpdateUserSettingRequest struct {
|
||||
UpstreamModelUpdateNotifyEnabled *bool `json:"upstream_model_update_notify_enabled,omitempty"`
|
||||
AcceptUnsetModelRatioModel bool `json:"accept_unset_model_ratio_model"`
|
||||
RecordIpLog bool `json:"record_ip_log"`
|
||||
AstraCostNoticeAck *bool `json:"astra_cost_notice_ack,omitempty"`
|
||||
}
|
||||
|
||||
func UpdateUserSetting(c *gin.Context) {
|
||||
@@ -1380,6 +1383,10 @@ func UpdateUserSetting(c *gin.Context) {
|
||||
UpstreamModelUpdateNotifyEnabled: upstreamModelUpdateNotifyEnabled,
|
||||
AcceptUnsetRatioModel: req.AcceptUnsetModelRatioModel,
|
||||
RecordIpLog: req.RecordIpLog,
|
||||
AstraCostNoticeAck: existingSettings.AstraCostNoticeAck,
|
||||
}
|
||||
if req.AstraCostNoticeAck != nil {
|
||||
settings.AstraCostNoticeAck = *req.AstraCostNoticeAck
|
||||
}
|
||||
|
||||
// 如果是webhook类型,添加webhook相关设置
|
||||
|
||||
@@ -121,7 +121,7 @@ func WeChatAuth(c *gin.Context) {
|
||||
})
|
||||
return
|
||||
}
|
||||
setupLogin(&user, c)
|
||||
setupLogin(&user, nil, c)
|
||||
}
|
||||
|
||||
type wechatBindRequest struct {
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
# Astra 返回模型校验
|
||||
|
||||
API 令牌的 `astra_auto_switch` 控制原始模型为 `gpt-6-astra` 的 HTTP `/v1/responses` 错配拦截和安全切换。独立字段 `astra_detection_enabled` 控制检测。关闭检测时不执行 Astra 模型比较、不拦截,并将自动切换保存为关闭;重新打开检测不会自动重新开启切换。检测开启且切换关闭时放行并记录 `mismatch_allowed`,开启时拦截并记录 `mismatch_rejected`。渠道 `responses_astra_model_guard` 字段仅保留存储兼容,界面标记停用,不再决定请求行为。
|
||||
|
||||
开启时,明确模型错配可以按现有路由策略安全切换,最多 3 次上游尝试。显式渠道固定、已交付内容、客户端取消、携带 previous_response_id 或 conversation 的请求不会透明切换。每次在 STREAMING_TIMEOUT 内、最多暂存 1 MiB,确认原始上游声明去除首尾空白后等于 `gpt-6-astra`,才允许提交响应头、ping 和内容。网络重试与 strict_session 策略保持现状。
|
||||
|
||||
接受后的流实时转发;后续错误模型事件不会交付,转为 `response.failed`,不再透明重试。每次有消耗的上游尝试独立结算,包括被拦截的结果;切换后的新尝试另行计费。沿用持久化尝试账本、usage 收集和结算,不修改价格或退款算法。
|
||||
|
||||
无法交付已验证结果时,客户端收到 HTTP 409 和 `response_model_rejected`;若已开始 SSE,则通过 `response.failed` 发送错误。提示包含请求及返回模型、每次消耗独立计费的说明,不包含渠道名称、地址或鉴权信息。拒收后仍收集 usage 以完成账本结算。
|
||||
|
||||
管理员日志详情的 `admin_info.request_details.model_guard.attempts` 保存尝试序号、渠道 ID、期望及返回模型、原因、耗时、是否交付以及已经取得的 usage。没有取得 usage 表示未知成本,不表示上游没有收费。普通用户日志不返回此字段。
|
||||
|
||||
请求所属用户可以查看尝试历史、实际费用、模型声明与校验结果;渠道 ID 仅管理员可见。尝试 API 增加 `charged`(已结算且费用大于零)及 `charge_explanation`。待结算尝试不会被误标为已扣费。
|
||||
|
||||
验证入口:`relay/response_model_test.go`;前端 `admin-request-details.test.tsx` 和 `responses-websocket-setting.test.tsx`。隔离运行验收还覆盖真实 HTTP 转发、429 混合重试、退款及单次扣费、超时关闭连接、ping 门控和客户端断线排空。
|
||||
|
||||
本功能仅验证上游响应的模型声明,不能识别伪报模型。首次启动新增令牌列并以迁移标记将已有令牌设为开启;新建令牌默认检测开启、自动切换关闭;已有令牌迁移后保留检测开启及原切换设置,重复启动保留用户选择。旧 Redis 哈希补齐新字段时不覆盖实时额度。启用并保存需要费用确认;令牌页的首次费用提示通过用户 setting 中的 `astra_cost_notice_ack` 记录已读。API 更新省略开关时保留原值,显式 false 正常持久化。
|
||||
|
||||
鉴权继续使用现有令牌所有权检查与认证缓存失效机制,新增审计字段不含凭据。安全参考:[OWASP Authentication Cheat Sheet](https://cheatsheetseries.owasp.org/cheatsheets/Authentication_Cheat_Sheet.html)、[OWASP Session Management Cheat Sheet](https://cheatsheetseries.owasp.org/cheatsheets/Session_Management_Cheat_Sheet.html)。
|
||||
@@ -0,0 +1,9 @@
|
||||
# custom.19: 错误日志实际分组
|
||||
|
||||
- 修复错误日志使用用户默认分组代替本次实际分组的问题。优先采用 RelayInfo 的实际计费分组,兼容 Auto/统一池及固定分组。
|
||||
- 未选择到渠道时,分组留空,不再伪装为 default。
|
||||
- 历史错误日志在 API 输出时,依据保存的 request_policy 中匹配渠道的分组纠正显示;在用户权限裁剪之前完成,仍不泄露私有路由信息。
|
||||
- 无证据的旧错误记录不猜测分组。消费日志、金额、token 用量和数据库中的历史记录不改写。
|
||||
- 历史记录的数据库分组筛选仍按原存储值工作;新的错误记录保存正确分组,筛选正常。
|
||||
|
||||
测试:实际错误分组、重试分组、无渠道、固定分组;历史读时投影及普通用户日志隐私;真实候选 HTTP 回归;旧版 Astra 守卫/逐次结算保持不变。
|
||||
@@ -54,11 +54,15 @@ query is redacted before request logging.
|
||||
Deployment boundaries and concurrency environment variables are documented in
|
||||
[v1.md](./v1.md#generic-task-management-api).
|
||||
|
||||
The host treats `protocols.openai_video.render` as a standard DTO, not an arbitrary
|
||||
JSON passthrough. Unknown top-level fields and legacy `task_id` are removed,
|
||||
`id` is forced to the public task id, and case-insensitive `url` entries are
|
||||
removed from metadata. Provider output URLs belong only behind artifact
|
||||
capabilities.
|
||||
The host requires `protocols.openai_video.render` to return a JSON object and
|
||||
preserves provider extensions, including output URLs and metadata. Plugins may
|
||||
return `task.data` to expose the latest persisted upstream snapshot. The host
|
||||
still overwrites `id`, `object`, `model`, `status`, `progress`, `created_at`, and
|
||||
`completed_at`, and removes legacy `task_id`. This is field-level projection,
|
||||
not byte-for-byte HTTP forwarding: known private task IDs are replaced before
|
||||
rendering, polling may redact inline video data, and retrieval reads the latest
|
||||
saved snapshot rather than issuing a new upstream request. Artifact content
|
||||
endpoints remain available for proxied downloads.
|
||||
|
||||
Provider-authenticated content URLs must use the channel base host or a
|
||||
plugin-declared `meta.allowedHosts` entry. A public dynamic CDN URL may instead
|
||||
|
||||
Vendored
+35
-5
@@ -1,4 +1,20 @@
|
||||
export type JSONValue = null | boolean | number | string | readonly JSONValue[] | {readonly [key: string]: JSONValue};
|
||||
export type HostCapability = "json-clone@1" | "submit-sse-delta@1";
|
||||
export type MutableJSON<T> = T extends readonly (infer Item)[] ? MutableJSON<Item>[] : T extends object ? {-readonly [Key in keyof T]: MutableJSON<T[Key]>} : T;
|
||||
export interface HostUtils {
|
||||
hasCapability(name: string): boolean;
|
||||
/** Independent native JS containers; at most 1 MiB encoded JSON, depth 32 and 32,768 nodes. */
|
||||
json: {clone<T extends JSONValue>(value: T): MutableJSON<T>};
|
||||
unixNow(): number;
|
||||
uuid(): string;
|
||||
hmacSHA256(message: string, secret: string): string;
|
||||
jwtSignHS256(claims: Record<string, JSONValue>, secret: string): string;
|
||||
base64(value: string): string;
|
||||
base64URL(value: string): string;
|
||||
base64URLDecode(value: string): string;
|
||||
volcSignV4(request: {method: string; url: string; headers?: Record<string, string>; body?: string; accessKey: string; secretKey: string; region?: string; service?: string; timestamp?: number}): Record<string, string>;
|
||||
}
|
||||
declare global {const utils: HostUtils;}
|
||||
export type FileReference = Readonly<{ref: string; field: string; filename: string; mimeType: string; size: number}>;
|
||||
export type FilePlaceholder = Readonly<{__fileRef: string; encoding: "base64" | "dataUrl"; mimeType?: string; maxBytes?: number}>;
|
||||
export type DecodedBody =
|
||||
@@ -20,15 +36,20 @@ export type ProtocolClaim =
|
||||
| {name: "openai_responses"; supports: readonly ResponsesMode[]; models?: readonly string[]}
|
||||
| {name: "openai_video"; models?: readonly string[]};
|
||||
export type LocalizedText = string | ({ en: string } & Record<string, string>);
|
||||
export type UsageFieldSchema = {type: "number"; unit: "second" | "count" | "token" | "credit"; description?: LocalizedText} | {type: "boolean"; description?: LocalizedText} | {enum: readonly string[]; description?: LocalizedText; enumLabels?: Readonly<Record<string, LocalizedText>>};
|
||||
export type UsageExample = {label: string; facts: Readonly<Record<string, string | number>>};
|
||||
export interface Meta {sortPriority?: number; website?: string; apiVersion: 1; key: string; name: string; icon?: string; description?: LocalizedText; version: string; author: {name: string; url?: string}; baseUrl?: string; channelTypes?: readonly number[]; models: readonly string[]; fetchMode: "per_task" | "batch"; allowedHosts?: readonly string[]; routes?: readonly NativeRoute[]; protocols?: readonly ProtocolClaim[]; usageSchema?: Readonly<Record<string, UsageFieldSchema>>; usageExamples?: readonly UsageExample[]; auth?: "none" | "api_key" | "vertex_oauth" | {type: "none" | "api_key" | "oauth2_jwt"}}
|
||||
export type UsageFieldSchema =
|
||||
| {type: "number"; unit: "count"; unitLabel?: LocalizedText; description?: LocalizedText}
|
||||
| {type: "number"; unit: "second" | "token" | "credit"; unitLabel?: never; description?: LocalizedText}
|
||||
| {type: "boolean"; unitLabel?: never; description?: LocalizedText}
|
||||
| {enum: readonly string[]; unitLabel?: never; description?: LocalizedText; enumLabels?: Readonly<Record<string, LocalizedText>>};
|
||||
export type UsageExample = {label: string; facts: Readonly<Record<string, string | number | boolean>>};
|
||||
export type UsageProfile = {models: readonly string[]; schema: Readonly<Record<string, UsageFieldSchema>>; examples?: readonly UsageExample[]};
|
||||
export interface Meta {requiredCapabilities?: readonly HostCapability[]; submitResponseTypes?: readonly ("json" | "sse")[]; sortPriority?: number; website?: string; apiVersion: 1; key: string; name: string; icon?: string; description?: LocalizedText; version: string; author: {name: string; url?: string}; baseUrl?: string; channelTypes?: readonly number[]; models: readonly string[]; fetchMode: "per_task" | "batch"; allowedHosts?: readonly string[]; routes?: readonly NativeRoute[]; protocols?: readonly ProtocolClaim[]; usageSchema?: Readonly<Record<string, UsageFieldSchema>>; usageExamples?: readonly UsageExample[]; usageProfiles?: readonly UsageProfile[]; auth?: "none" | "api_key" | "vertex_oauth" | {type: "none" | "api_key" | "oauth2_jwt"}}
|
||||
export interface TaskView {task_id: string; status: string; progress?: string; fail_reason?: string; created_at?: number; updated_at?: number; data?: unknown; properties?: Record<string, unknown>}
|
||||
export interface DriverContext {requestBody: unknown; requestHeaders: Readonly<Record<string, string>>; action: string; model: string; upstreamModel: string; baseUrl: string; apiKey?: string; authHeader: string; files: readonly FileReference[]; publicTaskId: string; originTasks?: readonly {taskId: string; upstreamTaskId: string; action: string; status: string; data: unknown}[]}
|
||||
export interface TaskQueryContext {taskId: string; publicTaskId: string; action: string; model: string; upstreamModel: string; baseUrl: string; apiKey?: string; authHeader: string; auth?: unknown; data: unknown; state: unknown}
|
||||
export interface BatchQueryContext {baseUrl: string; apiKey?: string; authHeader: string; auth?: unknown; tasks: readonly TaskQueryContext[]}
|
||||
export type HookHTTPResponse = {readonly status: number; readonly headers: Readonly<Record<string, string>>}
|
||||
export interface RequestDescriptor {url: string; method?: string; headers?: Record<string, string>; /** JSON body may contain FilePlaceholder objects at any depth; the host replaces each with a Base64 or data-URL string. */ body?: unknown; credentialless?: boolean; action?: string; model?: string; rewriteModel?: string; bodyType?: "json" | "multipart"; parts?: readonly {name: string; value?: unknown; fileRef?: string; filename?: string}[]}
|
||||
export interface RequestDescriptor {responseType?: "json" | "sse"; url: string; method?: string; headers?: Record<string, string>; /** JSON body may contain FilePlaceholder objects at any depth; the host replaces each with a Base64 or data-URL string. */ body?: unknown; credentialless?: boolean; action?: string; model?: string; rewriteModel?: string; bodyType?: "json" | "multipart"; parts?: readonly {name: string; value?: unknown; fileRef?: string; filename?: string}[]}
|
||||
export interface UpstreamResponse {statusCode: number; headers: Readonly<Record<string, readonly string[]>>; body: unknown}
|
||||
export interface NormalizedTaskResult {taskId?: string; status: "NOT_START" | "SUBMITTED" | "QUEUED" | "IN_PROGRESS" | "SUCCESS" | "FAILURE" | "UNKNOWN"; progress?: string; reason?: string; url?: string; remoteUrl?: string; completionTokens?: number; totalTokens?: number}
|
||||
export interface TaskArtifact {key: string; type: "video" | "audio" | "image" | "file"; mimeType?: string}
|
||||
@@ -39,6 +60,15 @@ export declare const protocols: {
|
||||
openai_video?: {decodeRequest(ctx: ProtocolDecodeContext): SubmitIntent; render(ctx: unknown, task: TaskView): unknown};
|
||||
};
|
||||
export declare function buildSubmitRequest(ctx: DriverContext): RequestDescriptor;
|
||||
export interface SubmitEvent {event: string; id: string; data: string}
|
||||
export type JSONPath = readonly (string | number)[];
|
||||
export type JSONChange =
|
||||
| {op: "set" | "append"; path: JSONPath; value: JSONValue}
|
||||
| {op: "appendText"; path: JSONPath; value: string};
|
||||
export interface SubmitEventDeltaResult {changes: readonly JSONChange[]; state: JSONValue; done: boolean}
|
||||
/** With submit-sse-delta@1, state is small control data; changes build the separate response body. */
|
||||
export declare function parseSubmitEventDelta(ctx: DriverContext, event: SubmitEvent, previousState: JSONValue | null): SubmitEventDeltaResult;
|
||||
export declare function parseSubmitEvent(ctx: DriverContext, event: SubmitEvent, previousState: JSONValue | null): {state: JSONValue; done: boolean};
|
||||
export declare function parseSubmitResponse(ctx: DriverContext, response: UpstreamResponse): {taskId: string; taskData?: unknown; immediate?: NormalizedTaskResult; state?: unknown};
|
||||
export declare function buildQueryRequest(ctx: TaskQueryContext): RequestDescriptor;
|
||||
export declare function buildBatchQueryRequest(ctx: BatchQueryContext, tasks: readonly TaskQueryContext[]): RequestDescriptor;
|
||||
@@ -46,6 +76,6 @@ export declare function parseTaskResult(ctx: TaskQueryContext, body: unknown, re
|
||||
export declare function parseBatchResult(ctx: BatchQueryContext, body: unknown, response: HookHTTPResponse): readonly (NormalizedTaskResult & {taskId: string; data?: unknown; state?: unknown})[];
|
||||
export declare function extractUsage(ctx: DriverContext & {usagePurpose?: "facts" | "billing_ratios"}): Readonly<Record<string, string | number | boolean>> | null;
|
||||
export declare function extractUsageOnSubmit(ctx: DriverContext, taskData: unknown): Readonly<Record<string, string | number | boolean>> | null;
|
||||
export declare function extractUsageOnComplete(task: TaskView, result: NormalizedTaskResult, data: unknown): Readonly<Record<string, string | number | boolean>> | null;
|
||||
export declare function extractUsageOnComplete(task: TaskQueryContext, result: NormalizedTaskResult, data: unknown): Readonly<Record<string, string | number | boolean>> | null;
|
||||
export declare function listArtifacts(task: {taskId: string; status: string; action: string; data: unknown; producerVersion: string}): readonly TaskArtifact[];
|
||||
export declare function buildContentRequest(ctx: DriverContext & {artifactKey: string; data: unknown; state?: unknown; upstreamTaskId: string; clientRequest: {method: "GET" | "HEAD"; headers: Readonly<Record<string, string>>}}): RequestDescriptor;
|
||||
+247
-4
@@ -28,9 +28,9 @@ export const meta = {
|
||||
|
||||
Each `protocols` entry claims a host protocol. A protocol that defines modes must be claimed in object form with an explicit `supports`; the bare-string form is valid only for protocols without modes (`openai_video`). `supports` declares which client request forms the plugin accepts on `openai_responses`: `"stream"` (`stream: true`), `"sync"` (neither flag; the call blocks until the terminal Response), `"background"` (`background: true`; the create returns a pending Response immediately). An unsupported request form is rejected with a 400 at channel selection, before any plugin hook runs and before billing. Retrieval (`GET /v1/responses/:response_id`) is not a mode: every created response is always retrievable.
|
||||
|
||||
`{name, models}` still narrows that protocol's endpoint bindings to a subset of `meta.models` and composes with `supports`. Unlisted models never enter the plugin on that protocol path — they fall through to the built-in Go relay. Cross-plugin endpoint conflicts are judged on the narrowed set, so two plugins may share one protocol path by claiming disjoint model sets.
|
||||
`{name, models}` still narrows that protocol's endpoint bindings to a subset of `meta.models` and composes with `supports`. Unlisted models never enter the plugin on that protocol path — they fall through to the built-in Go relay. Multiple plugins may claim the same model on the same host protocol path, including type-61 plugins and plugins with legacy `channelTypes`. Candidates are ordered by ascending plugin key. The selected channel determines the executing plugin: its channel type for legacy channels, or `setting.task_plugin_key` for type-61 channels.
|
||||
|
||||
Enabled uploads pre-flight the candidate against the live routing generation and reject the first channel-type, native-route, or protocol-model conflict (the error names the counterpart plugin). Set `force: true` or `enabled: false` to store the plugin anyway.
|
||||
Enabled uploads pre-flight the candidate against the live routing generation and reject channel-type, native-route, or incompatible protocol conflicts (the error names the counterpart plugin). Sharing a model on the same host protocol is allowed. Set `force: true` or `enabled: false` to store the plugin anyway.
|
||||
|
||||
`endpoints`, `routes[].renderer`, global `resolveRequest`, global `renderError`, and global `renderers` are rejected. `parseSubmitResponse` returns only `{taskId, taskData, immediate?, state?}`; `clientResponse` is rejected.
|
||||
|
||||
@@ -57,13 +57,97 @@ A plugin whose vendor has no LobeHub icon ships its logo as a sidecar file, `ico
|
||||
|
||||
`channelTypes` lists the legacy channel types this plugin's driver can drive (for example, sora declares `[55, 1]` because the same OpenAI-type base URL and bearer key serve both chat and video). Every entry equally participates in channel selection, historical `Task.Platform` matching, and the `byChannelType` routing index; the same type value may not appear on two plugins. Third-party plugins normally omit `channelTypes` and live on type-61 "Task Plugin" channels bound by `task_plugin_key`. The previous split identity/compatibility field names are rejected.
|
||||
|
||||
Before channel distribution, the host calls every mode-compatible candidate's
|
||||
`decodeRequest`. Only candidates returning a valid submit intent for the pinned
|
||||
model remain eligible. If all reject the request, the host returns the candidate
|
||||
failures, grouping identical messages with their plugin keys. A single candidate
|
||||
retains its original failure message. After channel selection, the executing plugin decodes again to build
|
||||
its own normalized request. Decoders must therefore tolerate repeated calls.
|
||||
Channel-type ownership remains unique; sharing models does not relax this rule.
|
||||
|
||||
Shared models may use different usage schemas and prices for each plugin.
|
||||
Administrators can save a provider expression in
|
||||
`billing_setting.plugin_billing_expr`, keyed by `<pluginKey>::<model>`.
|
||||
Without an override, the provider inherits the model-level expression. If the
|
||||
effective expression references a literal `u()` key absent from that provider's
|
||||
selected schema, the shared model is unconfigured for that provider; submission
|
||||
returns 400 `model_price_error` until compatible pricing is saved. Pricing pages
|
||||
show the provider variants and their configuration status. Model-level usage
|
||||
labels continue to come from the first plugin in ascending key order. Saving a
|
||||
model expression checks every candidate without its own override in the same
|
||||
draft. Provider overrides are not included in model-price synchronization.
|
||||
|
||||
Saved overrides remain visible and effective when only one provider is left.
|
||||
An override whose plugin is unavailable or no longer declares the model is
|
||||
marked `stale` in the administration snapshot. Its unchanged value can be kept
|
||||
while saving other prices, or explicitly removed; adding or changing an override
|
||||
still requires a current plugin/model binding. Public variants distinguish
|
||||
`ratio` fallback pricing from `tiered_expr`, so a provider using a configured
|
||||
per-call price is not marked unconfigured merely because it has no expression.
|
||||
|
||||
Numeric `usageSchema` fields declare a host-owned unit of `second`, `count`, `token`, or `credit`. Boolean fields declare `{type: "boolean"}`.
|
||||
|
||||
Number fields with `unit: "count"` may additionally declare `unitLabel`, an
|
||||
optional LocalizedText display unit. This distinguishes counted objects without
|
||||
changing the canonical unit, usage facts, expression keys, prices, or settlement:
|
||||
|
||||
```js
|
||||
image_count: {
|
||||
type: "number",
|
||||
unit: "count",
|
||||
unitLabel: {en: "image", zh: "张", "zh-TW": "張"},
|
||||
description: {en: "Image generation unit price", zh: "图片生成单价"},
|
||||
}
|
||||
```
|
||||
|
||||
Labels are short unit names or counting words, such as `image` / `张` or
|
||||
`song` / `首`, without prices, quantities, slashes, or trailing punctuation.
|
||||
They follow the same locale normalization, required English fallback, and
|
||||
256-rune-per-locale limit as field descriptions. A bare string is normalized to
|
||||
an English map. The UI uses the current language, its primary language, then
|
||||
English; omitted labels retain the host's existing count-unit display.
|
||||
Labels are plugin data, not frontend translation keys. They apply equally to
|
||||
`usageSchema` and `usageProfiles[].schema`, including pricing, calculators,
|
||||
plugin details, and usage-log pricing. Other units, booleans, and enums cannot
|
||||
declare `unitLabel`; token prices retain their per-million display scale.
|
||||
|
||||
This is an additive extension of plugin API version 1. Updated hosts accept
|
||||
existing plugins unchanged; older hosts reject the unknown `unitLabel` property.
|
||||
Upgrade the host before installing a plugin that declares it.
|
||||
|
||||
`meta.usageExamples` is an optional display-only list of pricing examples. Each `label` is a human-readable spec name. Each `facts` object must be a complete vector over `usageSchema`: every declared key present, no undeclared keys, and values that pass the same validation as usage facts. Plugins whose schema contains a `unit: "token"` number field must declare at least one example. Examples never affect billing. Labels are spec names such as `std · 10s`, not storefront prices.
|
||||
|
||||
### Usage profiles by model
|
||||
|
||||
`meta.usageProfiles` optionally contains `{models, schema, examples?}` entries.
|
||||
Each profile provides the complete schema and display examples for its models,
|
||||
replacing `usageSchema` and `usageExamples` without merging or inheriting them.
|
||||
Unmatched models use the plugin defaults. Models must use the declared spelling
|
||||
from `meta.models` and belong to at most one profile. Empty model lists and
|
||||
missing/null schemas are rejected; an empty schema object is valid. Token-unit
|
||||
profiles require their own examples. The same field, description, and example
|
||||
validation applies inside a profile as at the plugin level.
|
||||
|
||||
Define constants by use case (for example `IMAGE_USAGE_SCHEMA` and
|
||||
`VIDEO_USAGE_SCHEMA`) and reference them from the profiles. These are ordinary
|
||||
JavaScript constants: the host executes the module before reading `meta`.
|
||||
Static consumers can use the compiled display metadata in the marketplace index.
|
||||
|
||||
Model-level pricing and save validation select the declared target of an alias;
|
||||
ambiguous aliases retain the plugin defaults. Runtime validation selects the
|
||||
executing plugin's final upstream model, including each persisted task during
|
||||
batch polling. Existing saved expressions are not migrated. Single-plugin models retain their
|
||||
existing submission checks; shared models additionally enforce the per-provider
|
||||
field-reference compatibility rule above. Normal evaluation and quota safety
|
||||
checks continue to apply. Saving an expression uses its model's selected schema.
|
||||
|
||||
This remains plugin API version 1. New hosts accept older plugins unchanged;
|
||||
older hosts reject the new `usageProfiles` field. Upgrade the host to a build
|
||||
supporting profiles before installing a plugin that declares them.
|
||||
|
||||
## LocalizedText
|
||||
|
||||
`meta.description` and each `usageSchema` field `description` accept LocalizedText. A bare string is equivalent to `{en: <string>}`. A map must include a non-empty `en` value. The host normalizes both forms to a map; API responses always emit an object.
|
||||
`meta.description` and each `usageSchema` field `description` and `unitLabel` accept LocalizedText. A bare string is equivalent to `{en: <string>}`. A map must include a non-empty `en` value. The host normalizes both forms to a map; API responses always emit an object.
|
||||
|
||||
Enum fields may additionally declare `enumLabels`, a map from enum values to LocalizedText display names:
|
||||
|
||||
@@ -116,6 +200,37 @@ Descriptions are short, user-facing phrases without trailing punctuation in any
|
||||
- For example, replace `Requested video duration in seconds.` / `请求的视频时长,单位为秒。` with `Video generation unit price` / `视频生成单价`; replace `Whether audio is generated. Default true.` / `是否生成音频。默认为 true。` with `Whether audio is generated` / `是否生成音频`.
|
||||
- These are authoring conventions for plugin metadata, not additional runtime validation rules. Existing `LocalizedText` wire types, usage quantities, enum values and billing semantics are unchanged; a numeric field still contains usage, not a price.
|
||||
|
||||
## Host utilities and capabilities
|
||||
|
||||
`utils` exposes synchronous host tools. Existing utilities are `unixNow`, `uuid`,
|
||||
`hmacSHA256`, `jwtSignHS256`, `base64`, `base64URL`, `base64URLDecode`, and
|
||||
`volcSignV4`; their signatures are declared in `v1.d.ts`.
|
||||
|
||||
Plugins using new APIs declare exact versioned capabilities:
|
||||
|
||||
```js
|
||||
requiredCapabilities: ["json-clone@1", "submit-sse-delta@1"]
|
||||
```
|
||||
|
||||
The host rejects unknown, duplicate, or unsupported capability versions at load.
|
||||
Omitting the field preserves existing behavior. `utils.hasCapability(name)`
|
||||
supports optional feature detection on hosts that provide this utility; it
|
||||
does not grant a declared capability. These are additive API v1 extensions;
|
||||
older strict hosts reject `requiredCapabilities`, so upgrade the host first.
|
||||
|
||||
`json-clone@1` provides `utils.json.clone(value)`: an independent JSON snapshot
|
||||
with native mutable JavaScript objects and arrays. It accepts plain objects,
|
||||
arrays, null, booleans, finite numbers and strings. Undefined, functions,
|
||||
BigInt, dates, sparse arrays and cyclic structures are rejected. Object keys
|
||||
such as `__proto__` remain ordinary own data properties. The tool uses the
|
||||
host's JSON number and string normalization and preserves explicit zero and
|
||||
false values. Limits are 1 MiB of encoded JSON, 32 levels below the root and
|
||||
32,768 total nodes, including containers and scalar values. It performs no I/O.
|
||||
|
||||
Host functions are synchronous and must bound their own work: the JavaScript
|
||||
execution interrupt cannot preempt a running Go function. Request-scoped data
|
||||
must remain owned by the host request, not by pooled runtime globals.
|
||||
|
||||
## Request body
|
||||
|
||||
Every decoder receives one host-parsed body:
|
||||
@@ -180,7 +295,7 @@ The host registry owns these bindings:
|
||||
|
||||
`background: true` on create returns the pending Response immediately (host-synthesized, no plugin hook). Retrieval renders via `renderFinal` only at terminal status. Plugins declaring `"background"` need no create-time hook; the deliverable is rendered at retrieval via `renderFinal`. Plugins cannot observe the `background` field. Unlike the upstream OpenAI API, a plugin that does not declare `"stream"` rejects `stream: true` with a 400; the error names the supported forms.
|
||||
|
||||
Protocol media uses host-injected `ctx.artifacts[key].url`. Provider URLs from `task.data` are not protocol output. OpenAI Video projections are DTO-whitelisted and the host overwrites identity, lifecycle, timestamps, and removes URL-like metadata.
|
||||
Responses protocol media uses host-injected `ctx.artifacts[key].url`. OpenAI Video renderers may return provider fields and URLs from `task.data`; the host preserves those extensions and metadata, but overwrites `id`, `object`, `model`, `status`, `progress`, `created_at`, and `completed_at` and removes legacy `task_id`. Renderers must return a JSON object. Returning `task.data` exposes the latest persisted upstream snapshot, subject to the host's task-ID replacement and inline-video redaction; it does not perform a fresh upstream query or forward HTTP headers/status codes.
|
||||
|
||||
## Persisted data and driver hooks
|
||||
|
||||
@@ -224,3 +339,131 @@ The host classifies the HTTP status before trusting a non-terminal parse:
|
||||
A valid 2xx non-terminal parse resets `PollFailures` to 0. After `TASK_POLL_MAX_FAILURES` (default 20) consecutive failures the task becomes `FAILURE` and follows the existing refund chain. The 24h `TASK_TIMEOUT_MINUTES` sweep remains the outer deadline.
|
||||
|
||||
Marketplace index v1 plugin entries also accept optional `sortPriority` and `website` fields. Each source is sorted independently by descending priority and ascending key. Missing or invalid priorities fall back to `0`; invalid websites are hidden. Installed plugins and channel binding options use plugin metadata. Website links open only on user interaction; the host does not fetch them.
|
||||
|
||||
## Immediate completion and submission streams
|
||||
|
||||
A driver may return `immediate: {status: "SUCCESS" | "FAILURE", ...}` from
|
||||
`parseSubmitResponse`. For immediate success the host calls the existing
|
||||
`extractUsageOnComplete(ctx, result, taskData)` with a `TaskQueryContext`, using
|
||||
both model identities, the submitted task ID, parsed data and returned plugin
|
||||
state. The same model-specific usage validation as polling applies. Plugins
|
||||
return usage quantities, never host quota or a charge.
|
||||
|
||||
For task expressions the host overlays valid completion facts onto the frozen
|
||||
submission facts, evaluates the frozen expression and reserves any increase
|
||||
before persistence. The final quota, usage and tier are recorded before the
|
||||
single BillingSession settlement. Zero usage is retained. Invalid facts or an
|
||||
expression evaluation failure retain the reservation and produce a warning.
|
||||
Immediate failure settles to zero. Terminal tasks are not polled. Legacy
|
||||
per-call pricing keeps its existing meaning and submission-ratio hook.
|
||||
|
||||
### Optional SSE response reader
|
||||
|
||||
`meta.submitResponseTypes` defaults to `["json"]`. A plugin that needs SSE
|
||||
includes `"sse"` and exports `parseSubmitEvent`. `buildSubmitRequest` selects
|
||||
`responseType: "sse"` on its request descriptor (default `"json"`); the selected
|
||||
type must be declared. This optional v1 manifest field requires a supporting
|
||||
host: older strict hosts reject it instead of silently misreading the stream.
|
||||
|
||||
```ts
|
||||
parseSubmitEvent(ctx: DriverContext,
|
||||
event: {event: string; id: string; data: string},
|
||||
previousState: JSONValue | null
|
||||
): {state: JSONValue; done: boolean}
|
||||
```
|
||||
|
||||
The host accepts `text/event-stream`, joins multiple `data:` lines with newlines,
|
||||
handles CRLF, comment heartbeats, event names (default `message`) and event IDs.
|
||||
It does not parse provider JSON, interpret provider completion markers or sum
|
||||
usage. The plugin does those operations and returns explicit JSON state on every
|
||||
call; pooled runtimes do not guarantee module globals survive between events.
|
||||
Treat `previousState` as read-only JSON input. Construct owned arrays/objects
|
||||
when accumulating, rather than mutating host-backed arrays in place.
|
||||
|
||||
Idle runtimes survive garbage collection in a pool bounded by the plugin's
|
||||
configured concurrency. The host creates additional instances on demand and
|
||||
discards interrupted instances. Cancellation callbacks finish before an instance
|
||||
is returned to the pool. Hook bindings remain live; the host does not freeze
|
||||
exported functions or object members. No JavaScript execution slot is held while
|
||||
waiting for the next upstream event.
|
||||
|
||||
The host may copy plain JSON containers without re-encoding their string values;
|
||||
hook inputs and successive states remain separate snapshots. State size is still
|
||||
checked against its actual encoded JSON bytes after every event, including
|
||||
intermediate state that a later event might replace with a smaller value.
|
||||
|
||||
A complete, blank-line-delimited event whose hook returns `done: true` terminates
|
||||
the reader and closes the upstream response. The last state is passed as
|
||||
`response.body` to the existing `parseSubmitResponse`, retaining the upstream
|
||||
HTTP status and headers. The normal immediate-completion path then persists,
|
||||
settles and presents the result. SSE state is request-local; durable plugin
|
||||
state is still only the `state` returned by `parseSubmitResponse`.
|
||||
|
||||
EOF without `done`, invalid framing/state, read failures, hook failures and size
|
||||
limits fail submission. Events, accumulated JSON state and persisted task data
|
||||
are each limited to 1 MiB; data is never truncated. The existing streaming idle
|
||||
timeout applies (30 seconds if unset), as do the submission context deadline
|
||||
and per-hook execution timeout. Once a successful SSE HTTP response is accepted,
|
||||
failures are non-retryable, including 5xx-shaped local errors: starting the
|
||||
provider operation again could duplicate billable work.
|
||||
|
||||
No upstream events are forwarded to clients. Native presenters still return
|
||||
JSON, and Responses uses its existing presentation flow after the durable
|
||||
submission barrier. Native cancellation closes the upstream reader; Responses
|
||||
retains its existing independently bounded submission context. This interface
|
||||
does not add task creation before upstream completion or stream resumption.
|
||||
|
||||
### Optional host-managed SSE accumulation
|
||||
|
||||
Declaring `submit-sse-delta@1` requires `submitResponseTypes` to include `sse`
|
||||
and selects `parseSubmitEventDelta` instead of `parseSubmitEvent`. The descriptor
|
||||
still uses `responseType: "sse"`. Existing plugins without this capability keep
|
||||
the snapshot hook and its original contract.
|
||||
|
||||
```ts
|
||||
parseSubmitEventDelta(ctx: DriverContext,
|
||||
event: {event: string; id: string; data: string},
|
||||
previousState: JSONValue | null
|
||||
): {changes: JSONChange[]; state: JSONValue; done: boolean}
|
||||
```
|
||||
|
||||
`state` contains control data needed to interpret the next event, limited to
|
||||
64 KiB encoded JSON. Only this snapshot is passed back to JavaScript. The host
|
||||
maintains a separate result, initially null, and applies `changes` in order:
|
||||
|
||||
```js
|
||||
{op: "set", path: [], value: {text: "", images: []}}
|
||||
{op: "appendText", path: ["text"], value: "A new text fragment"}
|
||||
{op: "append", path: ["images"], value: {url: "https://cdn.example/image.png"}}
|
||||
```
|
||||
|
||||
- `set` replaces the root at an empty path, replaces an existing array element,
|
||||
or adds/replaces an object member. Parent containers must already exist.
|
||||
- `append` appends one JSON value to an existing array.
|
||||
- `appendText` appends a string to an existing string. The host retains text in
|
||||
Go and does not send accumulated text back to the plugin.
|
||||
- Paths are arrays of string object keys or non-negative integer array indices.
|
||||
Indices must address existing elements; sparse arrays and implicit parent
|
||||
creation are unsupported. Keys are data, never property-access expressions.
|
||||
|
||||
Each result contains exactly `changes`, `state`, and boolean `done`. Each change
|
||||
contains exactly `op`, `path`, and `value`. There are at most 256 changes per
|
||||
event, 32 path segments, 32 levels below the result root, 32,768 result nodes,
|
||||
and 1 MiB of encoded result JSON. Size accounting includes escaped keys and
|
||||
strings, punctuation and replacements, and is checked after every operation.
|
||||
An oversized intermediate operation fails even if a later operation would
|
||||
shrink the result. Invalid operations invalidate the submission; partial
|
||||
results are never published.
|
||||
|
||||
The host counts the encoded sizes of changed values rather than encoding the
|
||||
complete accumulated result on every event. At `done: true`, it materializes the
|
||||
ordinary JSON result and verifies its actual encoding against the recorded
|
||||
size before passing it as `response.body` to `parseSubmitResponse`. The result
|
||||
then follows the same persistence, usage validation and settlement path as the
|
||||
snapshot mode. State and builders belong to one request and are released on
|
||||
completion, cancellation or error; no JS slot is held during upstream I/O.
|
||||
|
||||
The plugin still interprets vendor payloads, decides which fields to update,
|
||||
distinguishes cumulative usage from increments, and decides completion. Alibaba
|
||||
uses this mode for interleaved image generation, keeping only content indices,
|
||||
text adjacency and finish reasons in its control state.
|
||||
Loaded 100 of 828 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user