diff --git a/.agents/github/ISSUE.md b/.agents/github/ISSUE.md new file mode 100644 index 000000000000..ec4e676a1d15 --- /dev/null +++ b/.agents/github/ISSUE.md @@ -0,0 +1,178 @@ + + +## Agent + +- Tool: +- Tool version: +- Model (full id): +- Host (CLI / IDE / GitHub coding agent / other): +- Date (UTC): + +## User request + +(verbatim or close paraphrase) + +## Out of scope — refuse + +If the request matches any item below, tell the user this repository does not +accept it, point them to the right place when there is one, and **do not file**. + +- Coding Plan +- Reverse-engineered channels +- 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) +- 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) + +- Matched: yes/no +- If yes, what was told to the user (stop here; do not file): + +## Kind + +- [ ] Bug +- [ ] Feature +- [ ] Investigation +- [ ] Other: + +## Usage / configuration / integration check + +Search these yourself before filing. Do not send the user to "read the docs first". If this is usage, configuration, or integration: answer the user and do not file. + +- https://docs.newapi.ai/ — what was searched, conclusion: +- https://deepwiki.com/QuantumNous/new-api — what was searched, conclusion: +- README / repo docs: +- Relevant code paths and conclusion: +- Can the current version already do this? (required for feature requests): +- Verdict: product bug or new feature / usage question (stop here): + +## Environment + +- new-api version / commit / image tag (not `latest` / `unknown`): +- Deploy source (repo release / official image / main source / other): +- Database (sqlite / mysql / postgres): + +## Problem facts + +Ask the user for every item and write it down: + +- Actual behavior: +- Impact: +- Frequency: +- Evidence that the problem is in new-api rather than the client or upstream: + +## Type-specific details + +Fill every applicable type. Write "not applicable" for the rest. Ask the user for missing items; do not invent them. + +### Relay / API + +- Request endpoint and method: +- Channel type: +- Model: +- Conversion format: +- Pass-through enabled: +- Evidence of upstream native support: +- Equivalent redacted request sent directly upstream: status, body, server logs: +- Same request through new-api: status, body, server logs: + +### Billing + +- Request endpoint and model: +- Response `usage`: +- Relevant ratio or pricing configuration: +- Consumption log: +- Expected charge and calculation basis: + +### Frontend + +- Page path: +- Browser and version: +- Active theme: +- Relevant browser Console / Network errors: + +### Deployment / upgrade + +- Deployment method: +- OS and architecture: +- Database type: +- Versions before and after the upgrade: +- Startup or migration logs: + +## Reproduction and expected result + +- Steps to reproduce: +- Expected result: +- Related screenshots (optional): + +## Feature (feature requests only) + +- Feature description: +- Use case: + +## Duplicate check + +- Search queries (issues, PRs, discussions): +- Closest existing threads: +- Why this is not a duplicate: + +## Research + +Open the docs and code. Do not write "already checked" without sources. + +### Docs + +- https://docs.newapi.ai/ : +- https://deepwiki.com/QuantumNous/new-api : +- README / other repo docs: +- Conclusions: + +### Code + +- Path — what it does, and how it relates: + +### Experiments + +- Command or redacted request: +- Direct upstream result: +- Result through new-api: +- Conclusion: + +## Working theory + +- What is broken or missing: +- Why: +- What would falsify this: + +## Scope + +- In scope for a later PR: +- Out of scope / not this repo: +- Large or directional feature? If yes, this issue is for maintainer alignment; do not open a PR yet. + +## Proposed direction + +(acceptance criteria, not an implementation dump) + +## Not verified + +(platforms, databases, providers, versions, paths not checked) + +## Related + +- Issues / PRs / upstream docs: diff --git a/.agents/github/PR.md b/.agents/github/PR.md new file mode 100644 index 000000000000..af60eb436a71 --- /dev/null +++ b/.agents/github/PR.md @@ -0,0 +1,136 @@ + + +## Agent + +- Tool: +- Tool version: +- Model (full id): +- Host (CLI / IDE / GitHub coding agent / other): +- Date (UTC): + +## Links + +- Closes # +- Related: + +## User request + +(verbatim or close paraphrase) + +## Out of scope — refuse + +If the change matches any item below, tell the user this repository does not +accept it and **do not open a PR**. + +- Coding Plan +- Reverse-engineered channels +- 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 +- 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): + +## Kind + +- [ ] Bug fix +- [ ] New feature +- [ ] Performance / refactor +- [ ] Docs +- [ ] Other: + +## Issue facts + +Take these from the linked issue. If a needed item is empty, ask the user that question. + +- Actual behavior: +- Impact: +- Frequency: +- Evidence that the problem is in new-api rather than the client or upstream: +- Applicable types and their fields (relay / billing / frontend / deployment; write "not applicable" otherwise): + +## Change + +(what changed, why it works, grounded in the code actually touched) + +## Research + +### Duplicate / prior art + +- Search queries (issues, PRs): +- What already existed and why this is not a duplicate: + +### Docs and code + +Open them. Do not write "already checked" without sources. + +- https://docs.newapi.ai/ : +- https://deepwiki.com/QuantumNous/new-api : +- README / repo docs: +- Code paths and what they imply for this change: + +### Alternatives considered + +- Option A: +- Option B: +- Why this approach: + +## Files + +| Path | Why | +| --- | --- | +| | | + +## Behavior + +- Before: +- After: +- Explicit non-goals / leftover work: + +## Verification + +Only what was actually run. + +- Commands and results: +- Manual steps and observed result: +- UI: screenshot or recording (or why none): +- Tests added or updated, or why none: +- Databases / providers / platforms exercised: +- Not verified: + +## Risks + +- Failure modes: +- Billing / quota / auth impact: +- Follow-ups: + +## Scope check + +- Single focused change: yes/no (if no, why): +- Secrets included: no +- Out of scope (Coding Plan / reverse-engineered channel / third-party wrapper / Codex): no diff --git a/.env.example b/.env.example index 3b8a2a5b9534..5ea44bc28958 100644 --- a/.env.example +++ b/.env.example @@ -35,6 +35,8 @@ # SQL_MAX_LIFETIME=60 # 慢查询日志阈值(毫秒),0 表示关闭慢查询日志,超出 0-3600000 范围回退默认值 200 # SQL_SLOW_THRESHOLD_MS=200 +# 跳过用户额度列 64 位 schema 检查(仅在已确认数据库列可容纳 64 位时启用) +# SKIP_64BIT_QUOTA_SCHEMA_CHECK=true # 缓存相关配置 @@ -54,12 +56,20 @@ # 任务和功能配置 # 更新任务启用 # UPDATE_TASK=true +# 异步任务硬超时(分钟),按提交时间计算,超时未完成的任务标记失败并退款;0 表示禁用 +# TASK_TIMEOUT_MINUTES=1440 +# 异步任务连续轮询失败阈值(上游 429/5xx/401/403、网络错误、无法识别的响应),达到后任务标记失败并退款;正常轮询成功一次即归零 +# TASK_POLL_MAX_FAILURES=20 # 对话超时设置 # 所有请求超时时间,单位秒,默认为0,表示不限制 # RELAY_TIMEOUT=0 # Relay HTTP 客户端空闲连接超时时间,单位秒,默认跟随 Go 标准库,设置为0表示不限制 # RELAY_IDLE_CONN_TIMEOUT=90 +# 等待上游返回响应头的超时时间,单位秒,默认 1800,设置为 0 表示不限制。 +# 仅约束「等待响应头」这一段;响应头返回之后的流式传输不受影响。 +# 注意:非流式请求通常要等上游生成完毕才会返回响应头,因此该值需留足余量。 +# RELAY_RESPONSE_HEADER_TIMEOUT=1800 # 流模式无响应超时时间,单位秒,如果出现空补全可以尝试改为更大值 # STREAMING_TIMEOUT=300 @@ -77,6 +87,8 @@ # 会话密钥 # SESSION_SECRET=random_string +# 登录密码请求体 RSA-OAEP 加密;默认关闭,且不能替代 HTTPS +# PASSWORD_LOGIN_ENCRYPTION_ENABLED=true # false/未配置:本地 HTTP 模式,关闭 refresh/logout OriginGuard,且不得设置 TRUSTED_URL;兼容本地开发代理。 # true:启用 Secure Refresh Cookie 和严格 OriginGuard,必须同时列出全部可信 HTTPS Origin。 # SESSION_COOKIE_TRUSTED_URL 多项用英文逗号分隔;不支持通配符、路径或域名后缀匹配。 diff --git a/.github/ISSUE_TEMPLATE/bug_report.yml b/.github/ISSUE_TEMPLATE/bug_report.yml index 8ee0cabeb18b..5f7bbdce18c2 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.yml +++ b/.github/ISSUE_TEMPLATE/bug_report.yml @@ -6,6 +6,9 @@ body: - type: markdown attributes: value: | + ## 提交前必读(请勿删除本节) - 文档:https://docs.newapi.ai/ diff --git a/.github/ISSUE_TEMPLATE/bug_report_en.yml b/.github/ISSUE_TEMPLATE/bug_report_en.yml index 67bf1404e0bb..a707f2857911 100644 --- a/.github/ISSUE_TEMPLATE/bug_report_en.yml +++ b/.github/ISSUE_TEMPLATE/bug_report_en.yml @@ -6,6 +6,9 @@ body: - type: markdown attributes: value: | + ## Read This First (Do Not Remove This Section) - Docs: https://docs.newapi.ai/ diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml index 2b1b118c231b..c5453ccd2892 100644 --- a/.github/ISSUE_TEMPLATE/feature_request.yml +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -6,6 +6,9 @@ body: - type: markdown attributes: value: | + ## 提交前必读(请勿删除本节) - 文档:https://docs.newapi.ai/ diff --git a/.github/ISSUE_TEMPLATE/feature_request_en.yml b/.github/ISSUE_TEMPLATE/feature_request_en.yml index a34070c6c201..c442682bee32 100644 --- a/.github/ISSUE_TEMPLATE/feature_request_en.yml +++ b/.github/ISSUE_TEMPLATE/feature_request_en.yml @@ -6,6 +6,9 @@ body: - type: markdown attributes: value: | + ## Read This First (Do Not Remove This Section) - Docs: https://docs.newapi.ai/ diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md index 7b1eb508ebef..237c3838256e 100644 --- a/.github/PULL_REQUEST_TEMPLATE.md +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -1,28 +1,40 @@ + # ⚠️ 提交说明 / PR Notice + +English template: `.github/PULL_REQUEST_TEMPLATE/en.md` + > [!IMPORTANT] > -> - 请提供**人工撰写**的简洁摘要,避免直接粘贴未经整理的 AI 输出。 +> - 描述可用 AI 辅助。提交前请审阅全文,并**声明对其负责**,避免未经核对的直接粘贴。 +> - 请按本模板填写后再提交。 -## 📝 变更描述 / Description -(简述:做了什么?为什么这样改能生效?请基于你对代码逻辑的理解来写,避免粘贴未经整理的内容) +## 🔗 关联任务 / Related Issue +- 新功能请填写下方 Issue 编号;若还没有对应 Issue,请先自行创建。功能讨论请放在 Issue 中进行。 +- 改动较大或方向性变更,请先在关联 Issue 中与维护者达成一致,再提交 PR。 +- Bug 修复请关联对应 Issue。设计取舍、理解偏差或预期不一致,更适合作为讨论或功能请求。 + +- Closes # ## 🚀 变更类型 / Type of change -- [ ] 🐛 Bug 修复 (Bug fix) - *请关联对应 Issue,避免将设计取舍、理解偏差或预期不一致直接归类为 bug* -- [ ] ✨ 新功能 (New feature) - *重大特性建议先通过 Issue 沟通* +- [ ] 🐛 Bug 修复 (Bug fix) +- [ ] ✨ 新功能 (New feature) - [ ] ⚡ 性能优化 / 重构 (Refactor) - [ ] 📝 文档更新 (Documentation) -## 🔗 关联任务 / Related Issue -- Closes # (如有) +## 📝 变更描述 / Description +(简述做了什么、为什么生效。如果难以简述,建议先拆分范围,或在 Issue 中与维护者对齐。) + +## 📸 运行证明 / Proof of Work +(请写明如何验证:实际步骤与观察结果。UI 变更请附截图或录屏;Bug 修复请说明复现过程与修复后结果。) ## ✅ 提交前检查项 / Checklist -- [ ] **人工确认:** 我已亲自整理并撰写此描述,没有直接粘贴未经处理的 AI 输出。 +- [ ] **人工确认:** 无论描述是否由 AI 生成,我已审阅全部内容,并声明对其准确性与完整性负责。 - [ ] **非重复提交:** 我已搜索现有的 [Issues](https://github.com/QuantumNous/new-api/issues) 与 [PRs](https://github.com/QuantumNous/new-api/pulls),确认不是重复提交。 -- [ ] **Bug fix 说明:** 若此 PR 标记为 `Bug fix`,我已提交或关联对应 Issue,且不会将设计取舍、预期不一致或理解偏差直接归类为 bug。 -- [ ] **变更理解:** 我已理解这些更改的工作原理及可能影响。 -- [ ] **范围聚焦:** 本 PR 未包含任何与当前任务无关的代码改动。 -- [ ] **本地验证:** 已在本地运行并通过测试或手动验证,维护者可以据此复核结果。 +- [ ] **新功能关联 Issue:** 若此 PR 标记为 `New feature`,我已关联对应 Issue;若尚无 Issue,我已先自行创建。 +- [ ] **事前沟通:** 若改动较大或涉及方向性变更,已在关联 Issue 中与维护者沟通并达成一致。 +- [ ] **功能范围:** 本 PR 不是 Coding Plan、逆向渠道、第三方封装接口,也不是对 Codex 渠道类型的改动。 +- [ ] **范围聚焦:** 本 PR 为一项聚焦改动,未包含无关代码。 +- [ ] **本地验证:** 已在本地运行并通过测试或手动验证,维护者可以据此复核。 - [ ] **安全合规:** 代码中无敏感凭据,且符合项目代码规范。 - -## 📸 运行证明 / Proof of Work -(请在此粘贴截图、关键日志或测试报告,以证明变更生效) diff --git a/.github/PULL_REQUEST_TEMPLATE/en.md b/.github/PULL_REQUEST_TEMPLATE/en.md new file mode 100644 index 000000000000..446e9cbd7962 --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE/en.md @@ -0,0 +1,40 @@ + +# PR Notice + +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. +> - Please complete this template before submitting. + +## Related Issue +- For new features, please fill in the Issue number below. If none exists yet, please create one first. Please discuss the feature in the Issue rather than using the PR in its place. +- For large or directional changes, please reach agreement with maintainers in the linked Issue before opening a PR. +- Bug fixes should link a corresponding Issue. Design trade-offs, misunderstandings, or mismatched expectations are a better fit for a discussion or feature request. + +- Closes # + +## Type of change +- [ ] Bug fix +- [ ] New feature +- [ ] Performance / Refactor +- [ ] 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.) + +## 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.) + +## 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. +- [ ] **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. +- [ ] **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. +- [ ] **Security:** This change does not include secrets and follows the project's coding guidelines. diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 98e2661d1925..b5578dfb1c3f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -76,7 +76,7 @@ jobs: - name: Set up Bun uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0 with: - bun-version: '1.3.14' + bun-version: '1.4.0' - name: Install dependencies run: bun install --frozen-lockfile @@ -85,4 +85,4 @@ jobs: run: bun run typecheck - name: Test - run: bun test + run: bun run test diff --git a/.github/workflows/electron-build.yml b/.github/workflows/electron-build.yml index ac4804a3b18d..08c57c9de0da 100644 --- a/.github/workflows/electron-build.yml +++ b/.github/workflows/electron-build.yml @@ -29,7 +29,7 @@ jobs: - name: Setup Bun uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0 with: - bun-version: latest + bun-version: '1.4.0' - name: Setup Node.js uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 diff --git a/.github/workflows/pr-check.yml b/.github/workflows/pr-check.yml deleted file mode 100644 index 67591702b4fd..000000000000 --- a/.github/workflows/pr-check.yml +++ /dev/null @@ -1,33 +0,0 @@ -name: PR Check - -permissions: - contents: read - issues: read - pull-requests: read - -on: - pull_request_target: - types: [opened, reopened] - -jobs: - pr-quality: - runs-on: ubuntu-latest - steps: - - uses: peakoss/anti-slop@85daca1880e9e1af197fc06ea03349daf08f4202 # v0.2.1 - with: - max-failures: 4 - require-description: true - - # require-linked-issue: false - blocked-terms: | - 🤖 Generated with Claude Code - - require-pr-template: true - strict-pr-template-sections: "✅ 提交前检查项 / Checklist" - - detect-spam-usernames: true - min-account-age: 30 - - failure-add-pr-labels: "pr-check-failed" - failure-pr-message: "感谢您的提交。由于该 PR 未遵循我们的贡献模板,且被识别为缺乏人工参与的纯 AI 生成内容 (AI Slop),我们将先予以关闭。我们更欢迎经过人工审核、验证并带有个人思考的贡献。如果您认为这其中存在误解,请回复告知。/ Thank you for your submission. This PR has been closed because it does not follow our contribution template and has been identified as purely AI-generated content (AI Slop) without meaningful human involvement. We prioritize contributions that are human-verified and reflect individual effort. If you believe this is a mistake, please let us know by replying to this comment." - close-pr: true diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index 547f5ed47c3c..b4bfed3c8bfa 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -24,11 +24,15 @@ jobs: fetch-depth: 0 - name: Determine Version run: | - VERSION=$(git describe --tags) - echo "VERSION=$VERSION" >> $GITHUB_ENV + if [[ "$GITHUB_REF" == refs/tags/* ]]; then + VERSION=${GITHUB_REF#refs/tags/} + else + VERSION=$(git describe --tags --match 'v[0-9]*') + fi + echo "VERSION=$VERSION" >> "$GITHUB_ENV" - uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0 with: - bun-version: latest + bun-version: '1.4.0' - name: Build Frontend env: CI: "" @@ -73,11 +77,15 @@ jobs: fetch-depth: 0 - name: Determine Version run: | - VERSION=$(git describe --tags) - echo "VERSION=$VERSION" >> $GITHUB_ENV + if [[ "$GITHUB_REF" == refs/tags/* ]]; then + VERSION=${GITHUB_REF#refs/tags/} + else + VERSION=$(git describe --tags --match 'v[0-9]*') + fi + echo "VERSION=$VERSION" >> "$GITHUB_ENV" - uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0 with: - bun-version: latest + bun-version: '1.4.0' - name: Build Frontend env: CI: "" @@ -121,11 +129,15 @@ jobs: fetch-depth: 0 - name: Determine Version run: | - VERSION=$(git describe --tags) - echo "VERSION=$VERSION" >> $GITHUB_ENV + if [[ "$GITHUB_REF" == refs/tags/* ]]; then + VERSION=${GITHUB_REF#refs/tags/} + else + VERSION=$(git describe --tags --match 'v[0-9]*') + fi + echo "VERSION=$VERSION" >> "$GITHUB_ENV" - uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0 with: - bun-version: latest + bun-version: '1.4.0' - name: Build Frontend env: CI: "" diff --git a/.github/workflows/sync-release-to-gitcode.yml b/.github/workflows/sync-release-to-gitcode.yml index 96aff722d538..d311f499eece 100644 --- a/.github/workflows/sync-release-to-gitcode.yml +++ b/.github/workflows/sync-release-to-gitcode.yml @@ -5,24 +5,25 @@ permissions: contents: read on: - push: - tags: - - '*' - - '!*-alpha*' workflow_dispatch: inputs: tag_name: description: GitHub release tag to sync required: true type: string + sync_files: + description: Sync GitHub release files to GitCode + required: false + default: false + type: boolean concurrency: - group: gitcode-release-${{ inputs.tag_name || github.ref_name }} + group: gitcode-release-${{ inputs.tag_name }} cancel-in-progress: false jobs: - prepare-release-assets: - name: Prepare GitHub release assets + prepare-release: + name: Create or update GitCode release if: ${{ vars.GITCODE_REPOSITORY != '' }} runs-on: ubuntu-latest timeout-minutes: 90 @@ -30,80 +31,125 @@ jobs: release_tag: ${{ steps.release.outputs.tag }} release_body: ${{ steps.release.outputs.body }} release_prerelease: ${{ steps.release.outputs.prerelease }} - bootstrap_asset: ${{ steps.release.outputs.bootstrap_asset }} - release_asset_matrix: ${{ steps.release.outputs.asset_matrix }} - has_bootstrap_asset: ${{ steps.release.outputs.has_bootstrap_asset }} - has_matrix_assets: ${{ steps.release.outputs.has_matrix_assets }} + release_asset_matrix: ${{ steps.assets.outputs.matrix }} + has_release_assets: ${{ steps.assets.outputs.has_assets }} steps: - - name: Wait for GitHub release workflows - if: ${{ github.event_name == 'push' }} + - name: Create or update GitCode release + id: release env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GITCODE_REPOSITORY: ${{ vars.GITCODE_REPOSITORY }} + GITCODE_TOKEN: ${{ secrets.GITCODE_TOKEN }} + RELEASE_TAG: ${{ inputs.tag_name }} run: | set -euo pipefail - wait_for_workflow() { - local workflow="$1" - local label="$2" - local last_state="" - - for attempt in $(seq 1 180); do - run_json="$( - gh run list \ - --repo "$GITHUB_REPOSITORY" \ - --workflow "$workflow" \ - --event push \ - --commit "$GITHUB_SHA" \ - --limit 1 \ - --json conclusion,status,url \ - --jq '.[0] // {}' - )" - status="$(jq -r '.status // empty' <<< "$run_json")" - conclusion="$(jq -r '.conclusion // empty' <<< "$run_json")" - state="${status:-not-found}/${conclusion:-pending}" - - if [[ "$state" != "$last_state" ]]; then - echo "$label: $state" - last_state="$state" - fi - - if [[ "$status" == "completed" ]]; then - if [[ "$conclusion" == "success" ]]; then - return 0 - fi - echo "::error::$label finished with conclusion: $conclusion" - return 1 - fi + release_json="$( + gh release view "$RELEASE_TAG" \ + --repo "$GITHUB_REPOSITORY" \ + --json body,isPrerelease + )" + release_body="$(jq -r '.body // ""' <<< "$release_json" | sed "s/'/’/g")" + release_prerelease="$(jq -r '.isPrerelease' <<< "$release_json")" + if [[ "$release_prerelease" == "true" ]]; then + release_status=pre + else + release_status=latest + fi - sleep 20 - done + gitcode_release_json="$(mktemp)" + gitcode_release_status="$( + curl -L -sS -o "$gitcode_release_json" -w "%{http_code}" \ + -H "PRIVATE-TOKEN: $GITCODE_TOKEN" \ + -H "Accept: application/json" \ + "https://api.gitcode.com/api/v5/repos/$GITCODE_REPOSITORY/releases/tags/$RELEASE_TAG" + )" - echo "::error::Timed out waiting for $label" - return 1 - } + request_json="$(mktemp)" + case "$gitcode_release_status" in + 200) + request_method=PATCH + request_url="https://api.gitcode.com/api/v5/repos/$GITCODE_REPOSITORY/releases/$RELEASE_TAG" + request_action=update + success_action=updated + jq -n \ + --arg tag_name "$RELEASE_TAG" \ + --arg name "$RELEASE_TAG" \ + --arg body "$release_body" \ + --arg release_status "$release_status" \ + '{ + tag_name: $tag_name, + name: $name, + body: $body, + release_status: $release_status + }' > "$request_json" + ;; + 404) + if ! git ls-remote --exit-code --tags \ + "https://gitcode.com/$GITCODE_REPOSITORY.git" \ + "refs/tags/$RELEASE_TAG" > /dev/null; then + echo "::error::Tag $RELEASE_TAG has not been mirrored to GitCode yet. Retry this workflow after the tag appears on GitCode." + exit 1 + fi - wait_for_workflow release.yml "Backend release" + request_method=POST + request_url="https://api.gitcode.com/api/v5/repos/$GITCODE_REPOSITORY/releases" + request_action=create + success_action=created + jq -n \ + --arg tag_name "$RELEASE_TAG" \ + --arg name "$RELEASE_TAG" \ + --arg body "$release_body" \ + --arg release_status "$release_status" \ + '{ + tag_name: $tag_name, + name: $name, + body: $body, + release_status: $release_status + }' > "$request_json" + ;; + *) + echo "::error::Failed to inspect GitCode release. Response code: $gitcode_release_status" + cat "$gitcode_release_json" + exit 1 + ;; + esac - if [[ "$GITHUB_REF_NAME" != *-* ]]; then - wait_for_workflow electron-build.yml "Electron release" + gitcode_response_json="$(mktemp)" + request_status="$( + curl -L -sS -o "$gitcode_response_json" -w "%{http_code}" \ + -X "$request_method" \ + -H "PRIVATE-TOKEN: $GITCODE_TOKEN" \ + -H "Content-Type: application/json" \ + -H "Accept: application/json" \ + --data-binary "@$request_json" \ + "$request_url" + )" + if [[ "$request_status" != "200" ]]; then + echo "::error::Failed to $request_action GitCode release. Response code: $request_status" + cat "$gitcode_response_json" + exit 1 fi + echo "GitCode release $success_action successfully" + + delimiter="release-body-$(openssl rand -hex 16)" + { + echo "tag=$RELEASE_TAG" + echo "body<<$delimiter" + echo "$release_body" + echo "$delimiter" + echo "prerelease=$release_prerelease" + } >> "$GITHUB_OUTPUT" - name: Download GitHub release assets - id: release + if: ${{ inputs.sync_files }} env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - GITCODE_REPOSITORY: ${{ vars.GITCODE_REPOSITORY }} - GITCODE_TOKEN: ${{ secrets.GITCODE_TOKEN }} - RELEASE_TAG: ${{ inputs.tag_name || github.ref_name }} + RELEASE_TAG: ${{ inputs.tag_name }} run: | set -euo pipefail mkdir release-assets - release_json="$( - gh release view "$RELEASE_TAG" \ - --repo "$GITHUB_REPOSITORY" \ - --json body,isPrerelease - )" gh release download "$RELEASE_TAG" \ --repo "$GITHUB_REPOSITORY" \ --dir release-assets @@ -128,6 +174,16 @@ jobs: find release-assets -maxdepth 1 -type f -print | sort + - name: Prepare GitCode release assets + id: assets + if: ${{ inputs.sync_files }} + env: + GITCODE_REPOSITORY: ${{ vars.GITCODE_REPOSITORY }} + GITCODE_TOKEN: ${{ secrets.GITCODE_TOKEN }} + RELEASE_TAG: ${{ inputs.tag_name }} + run: | + set -euo pipefail + local_assets="$( find release-assets -maxdepth 1 -type f -printf '%s\t%f\n' \ | sort -n -k1,1 \ @@ -142,23 +198,13 @@ jobs: -H "Accept: application/json" \ "https://api.gitcode.com/api/v5/repos/$GITCODE_REPOSITORY/releases/tags/$RELEASE_TAG" )" + if [[ "$gitcode_release_status" != "200" ]]; then + echo "::error::Failed to inspect GitCode release assets. Response code: $gitcode_release_status" + cat "$gitcode_release_json" + exit 1 + fi - case "$gitcode_release_status" in - 200) - release_exists=true - existing_assets="$(jq -c '[.assets[]?.name]' "$gitcode_release_json")" - ;; - 404) - release_exists=false - existing_assets='[]' - ;; - *) - echo "::error::Failed to inspect GitCode release. Response code: $gitcode_release_status" - cat "$gitcode_release_json" - exit 1 - ;; - esac - + existing_assets="$(jq -c '[.assets[]?.name]' "$gitcode_release_json")" missing_assets="$( jq -cn \ --argjson local_assets "$local_assets" \ @@ -170,37 +216,11 @@ jobs: --argjson missing_assets "$missing_assets" \ '$local_assets - $missing_assets | .[] | "Skipping existing GitCode asset: \(.)"' - if [[ "$release_exists" == "false" ]]; then - bootstrap_asset="$(jq -r '.[0] // ""' <<< "$missing_assets")" - asset_matrix="$(jq -c '.[1:]' <<< "$missing_assets")" - else - bootstrap_asset="" - asset_matrix="$missing_assets" - fi - - if [[ -n "$bootstrap_asset" ]]; then - has_bootstrap_asset=true - else - has_bootstrap_asset=false - fi - has_matrix_assets="$(jq -r 'length > 0' <<< "$asset_matrix")" - - delimiter="release-body-$(openssl rand -hex 16)" - { - echo "tag=$RELEASE_TAG" - echo "bootstrap_asset=$bootstrap_asset" - echo "asset_matrix=$asset_matrix" - echo "has_bootstrap_asset=$has_bootstrap_asset" - echo "has_matrix_assets=$has_matrix_assets" - echo "body<<$delimiter" - # sync_to_gitcode embeds this input in a single-quoted shell string. - # Replace ASCII apostrophes so release notes cannot break its script. - jq -r '.body // ""' <<< "$release_json" | sed "s/'/’/g" - echo "$delimiter" - echo "prerelease=$(jq -r '.isPrerelease' <<< "$release_json")" - } >> "$GITHUB_OUTPUT" + echo "matrix=$missing_assets" >> "$GITHUB_OUTPUT" + echo "has_assets=$(jq -r 'length > 0' <<< "$missing_assets")" >> "$GITHUB_OUTPUT" - name: Upload release assets for GitCode + if: ${{ inputs.sync_files && steps.assets.outputs.has_assets == 'true' }} uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: gitcode-release-assets @@ -208,41 +228,23 @@ jobs: if-no-files-found: error retention-days: 1 - gitcode-release-bootstrap: - name: Create GitCode release - needs: prepare-release-assets - if: ${{ vars.GITCODE_REPOSITORY != '' && needs.prepare-release-assets.outputs.has_bootstrap_asset == 'true' }} - uses: nvdacn/sync_to_gitcode/.github/workflows/CreateReleaseOnGitCode.yaml@18b70112d0e62260bc54085028e5510bbc6323b2 - with: - artifact_name: gitcode-release-assets - gitcode_repository: ${{ vars.GITCODE_REPOSITORY }} - default_branch: main - tag_name: ${{ needs.prepare-release-assets.outputs.release_tag }} - body: ${{ needs.prepare-release-assets.outputs.release_body }} - prerelease: ${{ needs.prepare-release-assets.outputs.release_prerelease == 'true' }} - file_name: ${{ needs.prepare-release-assets.outputs.bootstrap_asset }} - secrets: - GITCODE_TOKEN: ${{ secrets.GITCODE_TOKEN }} - gitcode-release-assets: name: Publish GitCode asset (${{ matrix.file_name }}) - needs: - - prepare-release-assets - - gitcode-release-bootstrap - if: ${{ always() && vars.GITCODE_REPOSITORY != '' && needs.prepare-release-assets.result == 'success' && (needs.gitcode-release-bootstrap.result == 'success' || needs.gitcode-release-bootstrap.result == 'skipped') && needs.prepare-release-assets.outputs.has_matrix_assets == 'true' }} + needs: prepare-release + if: ${{ inputs.sync_files && needs.prepare-release.result == 'success' && needs.prepare-release.outputs.has_release_assets == 'true' }} strategy: fail-fast: false max-parallel: 4 matrix: - file_name: ${{ fromJSON(needs.prepare-release-assets.outputs.release_asset_matrix) }} + file_name: ${{ fromJSON(needs.prepare-release.outputs.release_asset_matrix) }} uses: nvdacn/sync_to_gitcode/.github/workflows/CreateReleaseOnGitCode.yaml@18b70112d0e62260bc54085028e5510bbc6323b2 with: artifact_name: gitcode-release-assets gitcode_repository: ${{ vars.GITCODE_REPOSITORY }} default_branch: main - tag_name: ${{ needs.prepare-release-assets.outputs.release_tag }} - body: ${{ needs.prepare-release-assets.outputs.release_body }} - prerelease: ${{ needs.prepare-release-assets.outputs.release_prerelease == 'true' }} + tag_name: ${{ needs.prepare-release.outputs.release_tag }} + body: ${{ needs.prepare-release.outputs.release_body }} + prerelease: ${{ needs.prepare-release.outputs.release_prerelease == 'true' }} file_name: ${{ matrix.file_name }} secrets: GITCODE_TOKEN: ${{ secrets.GITCODE_TOKEN }} diff --git a/.gitignore b/.gitignore index dc328dd6c80c..042fef9a3cce 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ .idea +.review .vscode .zed .history @@ -7,6 +8,8 @@ upload *.db build *.db-journal +*.db-shm +*.db-wal logs web/dist web/node_modules @@ -20,7 +23,7 @@ tiktoken_cache .gocache .gomodcache/ .cache -plans +.plans .claude .cursor @@ -37,7 +40,7 @@ skills-lock.json # Local-only live probes and scratch test workspaces. .local-tests/ -service/relayconvert/chat_responses_live_local_test.go +relaykit/relayconvert/chat_responses_live_local_test.go service/openaicompat/chat_responses_live_local_test.go go.work go.work.sum diff --git a/AGENTS.md b/AGENTS.md index fa942c7159a6..2d572112fd17 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -81,6 +81,11 @@ Do NOT directly import or call `encoding/json` in business code. `json.RawMessag **Database compatibility:** All database code MUST work with SQLite, MySQL >= 5.7.8, and PostgreSQL >= 9.6 simultaneously. +- Any change that can affect database behavior MUST be verified before the work is considered complete. This includes ORM/database-driver dependency changes, connection/DSN/protocol or prepared-statement configuration, models and GORM tags, migrations and `AutoMigrate`, constraints and indexes, `Scanner`/`Valuer`/serializer behavior, raw SQL, transactions, and row locking. +- Required database verification MUST exercise real SQLite, MySQL, and PostgreSQL instances. Unit tests, mocks, a successful build, code inspection, or testing only one dialect are not substitutes. Use at least one supported version of each engine; changes that depend on version-specific behavior must also cover the minimum supported version. +- Treat GORM core and its database dialect/driver packages as a compatible version set. Any change to one of them requires checking upstream compatibility and running the complete three-database verification matrix; do not upgrade only the core package and infer that existing drivers remain compatible. +- Schema or migration changes MUST be tested both on a fresh database and by upgrading a representative database created by the latest released version. Run startup/migration at least twice to prove idempotency, and verify that existing data, indexes, constraints, and uniqueness guarantees are preserved. Cover the separately configured log database when the affected path is shared with or used by it. +- Record the exact database versions, commands, and results in the final handoff or pull request. If any required database verification cannot be run, report the blocker explicitly and do not claim the change is database-compatible or complete. - Prefer GORM methods (`Create`, `Find`, `Where`, `Updates`, etc.) over raw SQL. - Let GORM handle primary key generation; do not use `AUTO_INCREMENT` or `SERIAL` directly. - Standard `SELECT ... FOR UPDATE` row locks built with GORM query methods in `model/` MUST use `lockForUpdate(tx)`. Do not use the legacy GORM v1 pattern `tx.Set("gorm:query_option", "FOR UPDATE")`, because GORM v2 silently ignores it and no lock is acquired. Do not duplicate `clause.Locking{Strength: "UPDATE"}` at call sites; the shared helper emits `FOR UPDATE` for MySQL/PostgreSQL and skips it for SQLite, where the syntax is unsupported. Dialect-specific locking with different semantics (for example, a MySQL next-key/gap lock) may use raw SQL only behind explicit database-type branches with valid fallbacks for every supported database. @@ -102,12 +107,14 @@ Do NOT directly import or call `encoding/json` in business code. `json.RawMessag **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. Saturation bounds are int32 because quota columns (user/token/log) are 32-bit integers in the database, and every clamp/NaN fallback is logged via `common.SysError` since a single request should never approach those bounds. +- 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. @@ -116,6 +123,7 @@ Do NOT directly import or call `encoding/json` in business code. `json.RawMessag **Backend test quality:** Backend tests must protect real behavior, API contracts, billing/accounting invariants, data compatibility, or regression paths. +- **Do not scatter tests for a small change:** For a focused feature or fix, extend an existing suitable test file first. If a new test file is necessary, add at most one and consolidate the key regression cases there. MUST NOT create separate test files for the same small feature across `controller/`, `service/`, `setting/`, or other layers merely because its call chain crosses those layers. Do not repeat fixtures and assertions at each layer. Keep the cases compact and focused on observable behavior; the number of production files touched is not a reason to add more test files. - Do not add tests that only improve coverage numbers, prove that code happens to run, or lock in implementation details without a user-visible or cross-module contract. - Avoid fake fuzz/stress/smoke/performance tests built from random inputs, large loop counts, sleeps, timing comparisons, or log-only assertions. - Avoid duplicate tests that exercise the same branch with different names but no new invariant. @@ -149,8 +157,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. + **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. -- Always use the repository PR template at `.github/PULL_REQUEST_TEMPLATE.md` when drafting the PR title/body. Preserve the template structure and fill in the relevant sections instead of replacing it with an ad hoc format. +- 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. diff --git a/CLAUDE.md b/CLAUDE.md index ff3c01f0c766..97ba749380e3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -1,7 +1,15 @@ # CLAUDE.md — Project Conventions for new-api -@AGENTS.md +## MANDATORY: Read AGENTS.md with the Read tool -## Claude Code +Do not treat `@AGENTS.md` as loaded. Claude Code does not reliably inline that import. -- Follow the shared project instructions imported from `AGENTS.md`. \ No newline at end of file +Before any planning, coding, reviewing, or answering a project question, you MUST call the Read tool on the repo-root file `AGENTS.md` and wait for the full contents. This is the first action of every session and every new task. + +Rules: + +- Do not start from memory, summaries, or this file alone. +- Do not skip the Read because a previous turn mentioned AGENTS.md. +- 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. diff --git a/Dockerfile b/Dockerfile index 5be311b10e43..399a2ab85d94 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,4 @@ -FROM oven/bun:1@sha256:0733e50325078969732ebe3b15ce4c4be5082f18c4ac1a0f0ca4839c2e4e42a7 AS builder +FROM oven/bun:1.4.0@sha256:5ff609364c049b54eb0ff560ec96319729a972078ef2c755d758f0c6ef89c2d6 AS builder WORKDIR /build/web COPY web/package.json web/bun.lock ./ diff --git a/Dockerfile.dev b/Dockerfile.dev index bdc5be42da77..47475e34fc4f 100644 --- a/Dockerfile.dev +++ b/Dockerfile.dev @@ -12,6 +12,9 @@ ENV GOEXPERIMENT=greenteagc WORKDIR /build ADD go.mod go.sum ./ +# relaykit is a local submodule referenced via replace; its go.mod must be +# present for go mod download to resolve the main module graph. +ADD relaykit/go.mod ./relaykit/go.mod RUN go mod download COPY . . diff --git a/README.en.md b/README.en.md index 1204e1bc3dcc..cd0d551e1c01 100644 --- a/README.en.md +++ b/README.en.md @@ -281,7 +281,7 @@ docker run --name new-api -d --restart always \ - [Response Interface (Responses)](https://docs.newapi.pro/en/docs/api/ai-model/chat/openai/create-response) - [Image Interface (Image)](https://docs.newapi.pro/en/docs/api/ai-model/images/openai/v1-images-generations--post) - [Audio Interface (Audio)](https://docs.newapi.pro/en/docs/api/ai-model/audio/openai/create-transcription) -- [Video Interface (Video)](https://docs.newapi.pro/en/docs/api/ai-model/videos/create-video-generation) +- [Video Interface (Video)](https://docs.newapi.pro/en/docs/api/ai-model/videos/sora/createvideo) - [Embedding Interface (Embeddings)](https://docs.newapi.pro/en/docs/api/ai-model/embeddings/create-embedding) - [Rerank Interface (Rerank)](https://docs.newapi.pro/en/docs/api/ai-model/rerank/create-rerank) - [Realtime Conversation (Realtime)](https://docs.newapi.pro/en/docs/api/ai-model/realtime/create-realtime-session) diff --git a/README.fr.md b/README.fr.md index 937d306754e2..7870076eb6e1 100644 --- a/README.fr.md +++ b/README.fr.md @@ -282,7 +282,7 @@ docker run --name new-api -d --restart always \ - [Interface de réponse (Responses)](https://docs.newapi.pro/en/docs/api/ai-model/chat/openai/createresponse) - [Interface d'image (Image)](https://docs.newapi.pro/en/docs/api/ai-model/images/openai/post-v1-images-generations) - [Interface audio (Audio)](https://docs.newapi.pro/en/docs/api/ai-model/audio/openai/create-transcription) -- [Interface vidéo (Video)](https://docs.newapi.pro/en/docs/api/ai-model/audio/openai/createspeech) +- [Interface vidéo (Video)](https://docs.newapi.pro/en/docs/api/ai-model/videos/sora/createvideo) - [Interface d'incorporation (Embeddings)](https://docs.newapi.pro/en/docs/api/ai-model/embeddings/createembedding) - [Interface de rerank (Rerank)](https://docs.newapi.pro/en/docs/api/ai-model/rerank/creatererank) - [Conversation en temps réel (Realtime)](https://docs.newapi.pro/en/docs/api/ai-model/realtime/createrealtimesession) diff --git a/README.ja.md b/README.ja.md index 6c615349c165..24f88a333f17 100644 --- a/README.ja.md +++ b/README.ja.md @@ -284,7 +284,7 @@ docker run --name new-api -d --restart always \ - [レスポンスインターフェース (Responses)](https://docs.newapi.pro/ja/docs/api/ai-model/chat/openai/createresponse) - [イメージインターフェース (Image)](https://docs.newapi.pro/ja/docs/api/ai-model/images/openai/post-v1-images-generations) - [オーディオインターフェース (Audio)](https://docs.newapi.pro/ja/docs/api/ai-model/audio/openai/create-transcription) -- [ビデオインターフェース (Video)](https://docs.newapi.pro/ja/docs/api/ai-model/audio/openai/createspeech) +- [ビデオインターフェース (Video)](https://docs.newapi.pro/ja/docs/api/ai-model/videos/sora/createvideo) - [エンベッドインターフェース (Embeddings)](https://docs.newapi.pro/ja/docs/api/ai-model/embeddings/createembedding) - [再ランク付けインターフェース (Rerank)](https://docs.newapi.pro/ja/docs/api/ai-model/rerank/creatererank) - [リアルタイム対話インターフェース (Realtime)](https://docs.newapi.pro/ja/docs/api/ai-model/realtime/createrealtimesession) diff --git a/README.md b/README.md index 91778e659d3f..06924e547f45 100644 --- a/README.md +++ b/README.md @@ -282,7 +282,7 @@ docker run --name new-api -d --restart always \ - [Response Interface (Responses)](https://docs.newapi.pro/en/docs/api/ai-model/chat/openai/createresponse) - [Image Interface (Image)](https://docs.newapi.pro/en/docs/api/ai-model/images/openai/post-v1-images-generations) - [Audio Interface (Audio)](https://docs.newapi.pro/en/docs/api/ai-model/audio/openai/create-transcription) -- [Video Interface (Video)](https://docs.newapi.pro/en/docs/api/ai-model/audio/openai/createspeech) +- [Video Interface (Video)](https://docs.newapi.pro/en/docs/api/ai-model/videos/sora/createvideo) - [Embedding Interface (Embeddings)](https://docs.newapi.pro/en/docs/api/ai-model/embeddings/createembedding) - [Rerank Interface (Rerank)](https://docs.newapi.pro/en/docs/api/ai-model/rerank/creatererank) - [Realtime Conversation (Realtime)](https://docs.newapi.pro/en/docs/api/ai-model/realtime/createrealtimesession) @@ -327,6 +327,7 @@ docker run --name new-api -d --restart always \ | `SQL_DSN` | Database connection string | - | | `REDIS_CONN_STRING` | Redis connection string | - | | `RELAY_IDLE_CONN_TIMEOUT` | Idle keep-alive timeout for relay HTTP clients, seconds. Defaults to Go standard library behavior; set `0` to disable | `90` | +| `RELAY_RESPONSE_HEADER_TIMEOUT` | How long the relay waits for upstream **response headers**, seconds; set `0` to disable. Only bounds the header wait -- streaming after the headers arrive is unaffected. Note that non-streaming upstreams usually send headers only once generation finishes, so leave headroom | `1800` | | `STREAMING_TIMEOUT` | Streaming timeout (seconds) | `300` | | `STREAM_SCANNER_MAX_BUFFER_MB` | Max per-line buffer (MB) for the stream scanner; increase when upstream sends huge image/base64 payloads | `64` | | `MAX_REQUEST_BODY_MB` | Max request body size (MB, counted **after decompression**; prevents huge requests/zip bombs from exhausting memory). Exceeding it returns `413` | `32` | diff --git a/README.zh_CN.md b/README.zh_CN.md index b3bcd08dc96e..9ab237584ac4 100644 --- a/README.zh_CN.md +++ b/README.zh_CN.md @@ -282,7 +282,7 @@ docker run --name new-api -d --restart always \ - [响应接口 (Responses)](https://docs.newapi.pro/zh/docs/api/ai-model/chat/openai/createresponse) - [图像接口 (Image)](https://docs.newapi.pro/zh/docs/api/ai-model/images/openai/post-v1-images-generations) - [音频接口 (Audio)](https://docs.newapi.pro/zh/docs/api/ai-model/audio/openai/create-transcription) -- [视频接口 (Video)](https://docs.newapi.pro/zh/docs/api/ai-model/audio/openai/createspeech) +- [视频接口 (Video)](https://docs.newapi.pro/zh/docs/api/ai-model/videos/sora/createvideo) - [嵌入接口 (Embeddings)](https://docs.newapi.pro/zh/docs/api/ai-model/embeddings/createembedding) - [重排序接口 (Rerank)](https://docs.newapi.pro/zh/docs/api/ai-model/rerank/creatererank) - [实时对话 (Realtime)](https://docs.newapi.pro/zh/docs/api/ai-model/realtime/createrealtimesession) diff --git a/README.zh_TW.md b/README.zh_TW.md index ee9f952c3381..be7cf320a803 100644 --- a/README.zh_TW.md +++ b/README.zh_TW.md @@ -281,7 +281,7 @@ docker run --name new-api -d --restart always \ - [響應接口 (Responses)](https://docs.newapi.pro/zh/docs/api/ai-model/chat/openai/createresponse) - [圖像接口 (Image)](https://docs.newapi.pro/zh/docs/api/ai-model/images/openai/post-v1-images-generations) - [音訊接口 (Audio)](https://docs.newapi.pro/zh/docs/api/ai-model/audio/openai/create-transcription) -- [影片接口 (Video)](https://docs.newapi.pro/zh/docs/api/ai-model/audio/openai/createspeech) +- [影片接口 (Video)](https://docs.newapi.pro/zh/docs/api/ai-model/videos/sora/createvideo) - [嵌入接口 (Embeddings)](https://docs.newapi.pro/zh/docs/api/ai-model/embeddings/createembedding) - [重排序接口 (Rerank)](https://docs.newapi.pro/zh/docs/api/ai-model/rerank/creatererank) - [即時對話 (Realtime)](https://docs.newapi.pro/zh/docs/api/ai-model/realtime/createrealtimesession) diff --git a/THIRD-PARTY-LICENSES.md b/THIRD-PARTY-LICENSES.md index e04a8cd7028b..e67e61395f15 100644 --- a/THIRD-PARTY-LICENSES.md +++ b/THIRD-PARTY-LICENSES.md @@ -34,11 +34,13 @@ Transitive dependencies should be audited before a final external release. | backend | production | Go | `github.com/google/uuid` | `v1.6.0` | BSD-3-Clause | | backend | production | Go | `github.com/gorilla/websocket` | `v1.5.0` | BSD-2-Clause | | backend | production | Go | `github.com/grafana/pyroscope-go` | `v1.2.7` | Apache-2.0 | +| backend | production | Go | `github.com/grafana/sobek` | `v0.0.0-20260708062710-267a0e055bb4` | MIT | | backend | production | Go | `github.com/jfreymuth/oggvorbis` | `v1.0.5` | MIT | | backend | production | Go | `github.com/jinzhu/copier` | `v0.4.0` | MIT | | backend | production | Go | `github.com/joho/godotenv` | `v1.5.1` | MIT | | backend | production | Go | `github.com/mewkiz/flac` | `v1.0.13` | Unlicense | | backend | production | Go | `github.com/nicksnyder/go-i18n/v2` | `v2.6.1` | MIT | +| backend | test | Go | `github.com/openai/openai-go` | `v1.12.0` | Apache-2.0 | | backend | production | Go | `github.com/pkg/errors` | `v0.9.1` | BSD-2-Clause | | backend | production | Go | `github.com/pquerna/otp` | `v1.5.0` | Apache-2.0 | | backend | production | Go | `github.com/samber/hot` | `v0.11.0` | MIT | @@ -66,6 +68,7 @@ Transitive dependencies should be audited before a final external release. | backend | production | Go | `gorm.io/gorm` | `v1.25.2` | MIT | | backend | production | Go | `github.com/expr-lang/expr` | `v1.17.8` | MIT | | web | production | npm | `@base-ui/react` | `1.6.0` | MIT | +| web | production | npm | `@codemirror/lang-javascript` | `6.2.5` | MIT | | web | production | npm | `@codemirror/lang-markdown` | `6.5.1` | MIT | | web | production | npm | `@codemirror/language` | `6.12.4` | MIT | | web | production | npm | `@codemirror/state` | `6.7.1` | MIT | diff --git a/common/api_type.go b/common/api_type.go index 44841e11b8fc..b4ca7062b614 100644 --- a/common/api_type.go +++ b/common/api_type.go @@ -83,12 +83,17 @@ func ChannelType2APIType(channelType int) (int, bool) { apiType = constant.APITypeNewAPI } if apiType == -1 { + // Task plugin channels are served by the task relay and must never + // fall back to the OpenAI adaptor. + if channelType == constant.ChannelTypeTaskPlugin { + return -1, false + } return constant.APITypeOpenAI, false } return apiType, true } -func IsResponsesCompactAPIType(apiType int) bool { +func SupportsResponsesCompact(channelType, apiType int) bool { switch apiType { case constant.APITypeOpenAI, constant.APITypeCodex, diff --git a/common/api_type_task_plugin_test.go b/common/api_type_task_plugin_test.go new file mode 100644 index 000000000000..3a13e083bec3 --- /dev/null +++ b/common/api_type_task_plugin_test.go @@ -0,0 +1,14 @@ +package common + +import ( + "testing" + + "github.com/QuantumNous/new-api/constant" + "github.com/stretchr/testify/assert" +) + +func TestTaskPluginChannelHasNoOrdinaryAPIType(t *testing.T) { + apiType, ok := ChannelType2APIType(constant.ChannelTypeTaskPlugin) + assert.Equal(t, -1, apiType) + assert.False(t, ok) +} diff --git a/common/body_storage.go b/common/body_storage.go index 094dbda36d3f..515d749e10c9 100644 --- a/common/body_storage.go +++ b/common/body_storage.go @@ -20,6 +20,21 @@ type BodyStorage interface { Size() int64 // IsDisk 是否是磁盘存储 IsDisk() bool + // NewReader returns an independent reader positioned at the start of the + // stored payload. Each call returns a reader with its own cursor, so + // callers (e.g. http.Request.GetBody) can replay the body concurrently + // with, or after, other readers without sharing seek state. Closing the + // returned reader releases only that reader, never the storage itself; + // after the storage has been closed, NewReader returns ErrStorageClosed. + NewReader() (io.ReadCloser, error) +} + +// ReplayableBody is an outbound request body that can report its byte size and +// create independent readers for transport-level retries. +type ReplayableBody interface { + io.Reader + Size() int64 + NewReader() (io.ReadCloser, error) } // ErrStorageClosed 存储已关闭错误 @@ -80,6 +95,18 @@ func (m *memoryStorage) Bytes() ([]byte, error) { return m.data, nil } +func (m *memoryStorage) NewReader() (io.ReadCloser, error) { + m.mu.Lock() + defer m.mu.Unlock() + if atomic.LoadInt32(&m.closed) == 1 { + return nil, ErrStorageClosed + } + // A fresh bytes.Reader over the shared immutable backing array: an + // independent cursor at zero copy cost. NopCloser keeps Close a no-op, so + // the storage lifecycle stays owned by whoever holds the storage itself. + return io.NopCloser(bytes.NewReader(m.data)), nil +} + func (m *memoryStorage) Size() int64 { return m.size } @@ -229,6 +256,24 @@ func (d *diskStorage) Bytes() ([]byte, error) { return data, nil } +func (d *diskStorage) NewReader() (io.ReadCloser, error) { + d.mu.Lock() + defer d.mu.Unlock() + if atomic.LoadInt32(&d.closed) == 1 { + return nil, ErrStorageClosed + } + // A separate file descriptor over the same cache file: an independent + // cursor at zero copy cost. Closing the returned reader closes only that + // descriptor; the storage keeps owning the primary descriptor and the + // file's lifetime. Readers opened before Close stay usable even after the + // file is unlinked, as the descriptor keeps the inode alive. + file, err := os.Open(d.filePath) + if err != nil { + return nil, fmt.Errorf("failed to open body cache file for replay: %w", err) + } + return file, nil +} + func (d *diskStorage) Size() int64 { return d.size } @@ -302,10 +347,27 @@ func CreateBodyStorageFromReader(reader io.Reader, contentLength int64, maxBytes return storage, nil } -// ReaderOnly wraps an io.Reader to hide io.Closer, preventing http.NewRequest -// from type-asserting io.ReadCloser and closing the underlying BodyStorage. -func ReaderOnly(r io.Reader) io.Reader { - return struct{ io.Reader }{r} +type replayableBodyReader struct { + storage BodyStorage +} + +func (r replayableBodyReader) Read(p []byte) (int, error) { + return r.storage.Read(p) +} + +func (r replayableBodyReader) Size() int64 { + return r.storage.Size() +} + +func (r replayableBodyReader) NewReader() (io.ReadCloser, error) { + return r.storage.NewReader() +} + +// NewReplayableBodyReader exposes the replay capabilities of storage without +// exposing io.Closer. This keeps ownership of the storage lifecycle with the +// caller instead of allowing net/http to close it as the request body. +func NewReplayableBodyReader(storage BodyStorage) ReplayableBody { + return replayableBodyReader{storage: storage} } // CleanupOldCacheFiles 清理旧的缓存文件(用于启动时清理残留) diff --git a/common/body_storage_test.go b/common/body_storage_test.go new file mode 100644 index 000000000000..c3f877a60ff4 --- /dev/null +++ b/common/body_storage_test.go @@ -0,0 +1,37 @@ +package common + +import ( + "io" + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestNewReplayableBodyReaderKeepsStorageLifecycleWithCaller(t *testing.T) { + payload := []byte(`{"model":"test-model","input":"hello"}`) + storage, err := CreateBodyStorage(payload) + require.NoError(t, err) + defer storage.Close() + + body := NewReplayableBodyReader(storage) + assert.EqualValues(t, len(payload), body.Size()) + _, exposesCloser := any(body).(io.Closer) + assert.False(t, exposesCloser, "the request body must not expose the storage closer") + + req, err := http.NewRequest(http.MethodPost, "https://example.com", body) + require.NoError(t, err) + require.NoError(t, req.Body.Close()) + + replayBody, err := body.NewReader() + require.NoError(t, err, "closing the HTTP request body must not close the storage") + replay, err := io.ReadAll(replayBody) + require.NoError(t, err) + require.NoError(t, replayBody.Close()) + assert.Equal(t, payload, replay) + + require.NoError(t, storage.Close()) + _, err = body.NewReader() + require.ErrorIs(t, err, ErrStorageClosed) +} diff --git a/common/constants.go b/common/constants.go index d6b4fb52284c..2c7b90995094 100644 --- a/common/constants.go +++ b/common/constants.go @@ -60,6 +60,7 @@ var ItemsPerPage = 10 var MaxRecentItems = 1000 var PasswordLoginEnabled = true +var PasswordLoginEncryptionEnabled = false var PasswordRegisterEnabled = true var EmailVerificationEnabled = false var GitHubOAuthEnabled = false @@ -162,6 +163,16 @@ var BatchUpdateInterval int var RelayTimeout int // unit is second var RelayIdleConnTimeout int // unit is second + +// RelayResponseHeaderTimeout limits how long the relay transport waits for the +// upstream response headers after the request has been fully written. +// 0 disables it (previous behaviour: wait forever). +// +// Note this is NOT the same as RelayTimeout (http.Client.Timeout), which covers +// the whole response read and therefore breaks legitimate long streaming calls. +// ResponseHeaderTimeout only bounds the wait for the response headers; once the +// headers arrive, streaming is unaffected. +var RelayResponseHeaderTimeout int // unit is second var RelayMaxIdleConns int var RelayMaxIdleConnsPerHost int diff --git a/common/database.go b/common/database.go index 30d341f37268..44e8753139fa 100644 --- a/common/database.go +++ b/common/database.go @@ -41,4 +41,24 @@ func UsingLogDatabase(databaseType DatabaseType) bool { return logDatabaseType == databaseType } -var SQLitePath = "one-api.db?_busy_timeout=30000" +// SQLitePath is the DSN for the default SQLite database. It uses WAL journal +// mode so readers are never blocked by the single writer, plus a 30s busy +// timeout for writers to queue. +// +// Two details are non-obvious and both are required for concurrent correctness: +// +// 1. The busy timeout must be passed as a `_pragma=busy_timeout(30000)` DSN +// parameter. The pure-Go driver (modernc.org/sqlite, used through +// github.com/glebarez/sqlite) silently ignores the plain `_busy_timeout=` +// form, so without this the effective timeout stays at SQLite's 5s default +// and concurrent writes surface as "database is locked" (see #6805). +// +// 2. `_txlock=immediate` (BEGIN IMMEDIATE) must be enabled. Without it, a +// transaction that first SELECTs (establishing a read snapshot) and then +// writes can hit SQLITE_BUSY_SNAPSHOT when another connection commits in +// between; the busy handler does not cover that case, so the write fails +// instantly no matter the timeout. BEGIN IMMEDIATE takes the write lock up +// front, so writers serialize through the busy timeout instead of dying on +// a stale snapshot. Autocommit SELECTs stay concurrent because WAL keeps +// readers unlocked. +var SQLitePath = "one-api.db?_pragma=busy_timeout(30000)&_pragma=journal_mode(WAL)&_txlock=immediate" diff --git a/common/etag.go b/common/etag.go new file mode 100644 index 000000000000..cbae5a58654e --- /dev/null +++ b/common/etag.go @@ -0,0 +1,50 @@ +package common + +import ( + "crypto/sha256" + "encoding/hex" + "reflect" + "strings" +) + +type digestAnchor struct{} + +func modulePath() string { + return reflect.TypeOf(digestAnchor{}).PkgPath() +} + +var digestSeed = func() (s [sha256.Size]byte) { + return sha256.Sum256([]byte(modulePath())) +}() + +// ETagFor returns a weak ETag derived from the namespace and content. +func ETagFor(namespace, content string) string { + buf := make([]byte, 0, sha256.Size+1+len(namespace)+1+len(content)) + buf = append(buf, digestSeed[:]...) + buf = append(buf, 0) + buf = append(buf, namespace...) + buf = append(buf, 0) + buf = append(buf, content...) + digest := sha256.Sum256(buf) + return `W/"` + hex.EncodeToString(digest[:]) + `"` +} + +// ETagMatches reports whether an If-None-Match header matches etag under weak +// comparison (RFC 9110 §13.1.2): the W/ prefix is ignored on both sides, and +// "*" matches everything. +func ETagMatches(ifNoneMatch, etag string) bool { + ifNoneMatch = strings.TrimSpace(ifNoneMatch) + if ifNoneMatch == "" { + return false + } + if ifNoneMatch == "*" { + return true + } + etag = strings.TrimPrefix(etag, "W/") + for candidate := range strings.SplitSeq(ifNoneMatch, ",") { + if strings.TrimPrefix(strings.TrimSpace(candidate), "W/") == etag { + return true + } + } + return false +} diff --git a/common/init.go b/common/init.go index 4d4c62b27cac..d06accec7a49 100644 --- a/common/init.go +++ b/common/init.go @@ -87,6 +87,7 @@ func InitEnv() { DebugEnabled = os.Getenv("DEBUG") == "true" MemoryCacheEnabled = os.Getenv("MEMORY_CACHE_ENABLED") == "true" IsMasterNode = os.Getenv("NODE_TYPE") != "slave" + PasswordLoginEncryptionEnabled = GetEnvOrDefaultBool("PASSWORD_LOGIN_ENCRYPTION_ENABLED", false) initNodeNameIdentity() TLSInsecureSkipVerify = GetEnvOrDefaultBool("TLS_INSECURE_SKIP_VERIFY", false) if TLSInsecureSkipVerify { @@ -110,6 +111,7 @@ func InitEnv() { BatchUpdateInterval = GetEnvOrDefault("BATCH_UPDATE_INTERVAL", 5) RelayTimeout = GetEnvOrDefault("RELAY_TIMEOUT", 0) RelayIdleConnTimeout = GetEnvOrDefault("RELAY_IDLE_CONN_TIMEOUT", 90) + RelayResponseHeaderTimeout = GetEnvOrDefault("RELAY_RESPONSE_HEADER_TIMEOUT", 1800) RelayMaxIdleConns = GetEnvOrDefault("RELAY_MAX_IDLE_CONNS", 500) RelayMaxIdleConnsPerHost = GetEnvOrDefault("RELAY_MAX_IDLE_CONNS_PER_HOST", 100) @@ -187,6 +189,8 @@ func initConstantEnv() { constant.GetMediaToken = GetEnvOrDefaultBool("GET_MEDIA_TOKEN", true) constant.GetMediaTokenNotStream = GetEnvOrDefaultBool("GET_MEDIA_TOKEN_NOT_STREAM", false) constant.UpdateTask = GetEnvOrDefaultBool("UPDATE_TASK", true) + constant.TaskPluginEnabled = GetEnvOrDefaultBool("TASK_PLUGIN_ENABLED", true) + constant.TaskPluginOverrideEnabled = GetEnvOrDefaultBool("TASK_PLUGIN_OVERRIDE_ENABLED", true) constant.AzureDefaultAPIVersion = GetEnvOrDefaultString("AZURE_DEFAULT_API_VERSION", "2025-04-01-preview") constant.NotifyLimitCount = GetEnvOrDefault("NOTIFY_LIMIT_COUNT", 2) constant.NotificationLimitDurationMinute = GetEnvOrDefault("NOTIFICATION_LIMIT_DURATION_MINUTE", 10) @@ -198,6 +202,14 @@ func initConstantEnv() { constant.TaskQueryLimit = GetEnvOrDefault("TASK_QUERY_LIMIT", 1000) // 异步任务超时时间(分钟),超过此时间未完成的任务将被标记为失败并退款。0 表示禁用。 constant.TaskTimeoutMinutes = GetEnvOrDefault("TASK_TIMEOUT_MINUTES", 1440) + // Consecutive unrecognized/transient poll failures before the task is failed and refunded. + constant.TaskPollMaxFailures = GetEnvOrDefault("TASK_POLL_MAX_FAILURES", 20) + // 声明式任务协议桥只观察数据库;这些值控制一次客户端观察连接, + // 不改变后台轮询或结算生命周期。 + constant.TaskPluginProtocolTimeoutSeconds = GetEnvOrDefault("TASK_PLUGIN_PROTOCOL_TIMEOUT_SECONDS", 600) + constant.TaskPluginProtocolTickMilliseconds = GetEnvOrDefault("TASK_PLUGIN_PROTOCOL_TICK_MILLISECONDS", 2000) + constant.TaskPluginProtocolTickJitterMilliseconds = GetEnvOrDefault("TASK_PLUGIN_PROTOCOL_TICK_JITTER_MILLISECONDS", 500) + constant.TaskPluginProtocolHeartbeatSeconds = GetEnvOrDefault("TASK_PLUGIN_PROTOCOL_HEARTBEAT_SECONDS", 15) soraPatchStr := GetEnvOrDefaultString("TASK_PRICE_PATCH", "") if soraPatchStr != "" { diff --git a/common/json.go b/common/json.go index 1625be6d51f7..d7effa36ef32 100644 --- a/common/json.go +++ b/common/json.go @@ -22,6 +22,14 @@ func Marshal(v any) ([]byte, error) { return json.Marshal(v) } +func IndentJson(data []byte) ([]byte, error) { + var buffer bytes.Buffer + if err := json.Indent(&buffer, data, "", " "); err != nil { + return nil, err + } + return buffer.Bytes(), nil +} + func GetJsonType(data json.RawMessage) string { trimmed := bytes.TrimSpace(data) if len(trimmed) == 0 { diff --git a/common/password_crypto.go b/common/password_crypto.go new file mode 100644 index 000000000000..efbb97acbd34 --- /dev/null +++ b/common/password_crypto.go @@ -0,0 +1,115 @@ +package common + +import ( + "crypto/rand" + "crypto/rsa" + "crypto/sha256" + "crypto/x509" + "encoding/base64" + "encoding/hex" + "encoding/pem" + "errors" + "fmt" + "strings" + "sync" +) + +const passwordEncryptionKeyBits = 2048 + +var ErrPasswordEncryptionInvalid = errors.New("password encryption payload is invalid") + +var passwordEncryptionState struct { + sync.RWMutex + privateKey *rsa.PrivateKey + publicKey string + keyID string +} + +// GeneratePasswordEncryptionPrivateKey creates the server key used to decrypt +// browser login passwords. The caller is responsible for persisting the PEM. +func GeneratePasswordEncryptionPrivateKey() (string, error) { + privateKey, err := rsa.GenerateKey(rand.Reader, passwordEncryptionKeyBits) + if err != nil { + return "", fmt.Errorf("generate password encryption key: %w", err) + } + privateKeyDER, err := x509.MarshalPKCS8PrivateKey(privateKey) + if err != nil { + return "", fmt.Errorf("marshal password encryption key: %w", err) + } + return string(pem.EncodeToMemory(&pem.Block{ + Type: "PRIVATE KEY", + Bytes: privateKeyDER, + })), nil +} + +// LoadPasswordEncryptionPrivateKey validates a persisted key before replacing +// the active in-memory key used by request handlers. +func LoadPasswordEncryptionPrivateKey(privateKeyPEM string) error { + block, rest := pem.Decode([]byte(privateKeyPEM)) + if block == nil || block.Type != "PRIVATE KEY" || strings.TrimSpace(string(rest)) != "" { + return errors.New("password encryption key is not valid PKCS#8 PEM") + } + parsed, err := x509.ParsePKCS8PrivateKey(block.Bytes) + if err != nil { + return fmt.Errorf("parse password encryption key: %w", err) + } + privateKey, ok := parsed.(*rsa.PrivateKey) + if !ok { + return errors.New("password encryption key is not RSA") + } + if privateKey.N == nil || privateKey.N.BitLen() < passwordEncryptionKeyBits { + return fmt.Errorf("password encryption key must be at least %d bits", passwordEncryptionKeyBits) + } + if err := privateKey.Validate(); err != nil { + return fmt.Errorf("validate password encryption key: %w", err) + } + privateKey.Precompute() + + publicKeyDER, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey) + if err != nil { + return fmt.Errorf("marshal password encryption public key: %w", err) + } + publicKeyPEM := string(pem.EncodeToMemory(&pem.Block{ + Type: "PUBLIC KEY", + Bytes: publicKeyDER, + })) + keyDigest := sha256.Sum256(publicKeyDER) + keyID := hex.EncodeToString(keyDigest[:16]) + + passwordEncryptionState.Lock() + defer passwordEncryptionState.Unlock() + passwordEncryptionState.privateKey = privateKey + passwordEncryptionState.publicKey = publicKeyPEM + passwordEncryptionState.keyID = keyID + return nil +} + +// PasswordEncryptionPublicKey returns the active key identifier and SPKI PEM +// public key exposed to browser clients. +func PasswordEncryptionPublicKey() (keyID string, publicKeyPEM string) { + passwordEncryptionState.RLock() + defer passwordEncryptionState.RUnlock() + return passwordEncryptionState.keyID, passwordEncryptionState.publicKey +} + +// DecryptPassword decrypts a base64 RSA-OAEP/SHA-256 password submitted by a +// browser. All malformed inputs share one error so callers do not expose +// cryptographic details to unauthenticated clients. +func DecryptPassword(ciphertextBase64 string, keyID string) (string, error) { + passwordEncryptionState.RLock() + privateKey := passwordEncryptionState.privateKey + activeKeyID := passwordEncryptionState.keyID + passwordEncryptionState.RUnlock() + if privateKey == nil || keyID == "" || keyID != activeKeyID { + return "", ErrPasswordEncryptionInvalid + } + ciphertext, err := base64.StdEncoding.DecodeString(ciphertextBase64) + if err != nil || len(ciphertext) != privateKey.Size() { + return "", ErrPasswordEncryptionInvalid + } + plaintext, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, privateKey, ciphertext, nil) + if err != nil || len(plaintext) == 0 { + return "", ErrPasswordEncryptionInvalid + } + return string(plaintext), nil +} diff --git a/common/quota_math.go b/common/quota_math.go index 86c72492db70..fb0e71caa612 100644 --- a/common/quota_math.go +++ b/common/quota_math.go @@ -8,14 +8,24 @@ import ( ) // Quota conversions are centralized here so every billing path shares one -// saturation + logging policy. Quota columns (user/token/log) are 32-bit -// integers in the database, so an oversized product must clamp to the int32 -// range instead of wrapping around and turning a charge into a credit. +// saturation + logging policy. Single-request charges stay bounded to int32; +// top-ups and wallet-priced purchases use a JavaScript-safe 64-bit domain. const ( - MaxQuota = math.MaxInt32 - MinQuota = math.MinInt32 + MaxQuota = math.MaxInt32 + MinQuota = math.MinInt32 + MaxWalletQuota = 1<<53 - 1 ) +// ValidateWalletQuota enforces the upper bound shared by wallet mutations. +// Negative balances remain valid because billing can temporarily overdraw a +// wallet; callers that accept credits must apply their own positive check. +func ValidateWalletQuota(quota int) error { + if quota > MaxWalletQuota { + return fmt.Errorf("wallet quota exceeds %d", MaxWalletQuota) + } + return nil +} + // QuotaClampKind identifies why a quota conversion had to be saturated. type QuotaClampKind string @@ -27,11 +37,11 @@ const ( ) // QuotaClamp describes a single saturation event: a quota conversion whose -// input fell outside the representable int32 range (or was NaN) and was +// input fell outside its supported range (or was NaN) and was // therefore clamped. It is surfaced to billing callers so the event can be // recorded on the related consume/task log for admin auditing. type QuotaClamp struct { - Op string `json:"op"` // "QuotaFromFloat" | "QuotaRound" | "QuotaFromDecimal" + Op string `json:"op"` // "QuotaFromFloat" | "QuotaRound" | "QuotaFromDecimal" | "WalletQuotaFromDecimal" Kind QuotaClampKind `json:"kind"` // "overflow" | "underflow" | "nan" Original float64 `json:"original"` // best-effort pre-clamp value (decimal -> float64 approx) Clamped int `json:"clamped"` // the saturated result actually used @@ -61,23 +71,27 @@ func (c *QuotaClamp) AuditMap() map[string]interface{} { } } -// saturateQuota converts an already-rounded quota value to int, clamping to -// the int32 range. Whenever clamping (what would otherwise be an integer -// wraparound) or a NaN fallback is triggered it logs a warning, because in +// saturateQuota converts an already-rounded single-request quota to int. +// Whenever clamping (what would otherwise be an integer wraparound) or a NaN +// fallback is triggered it logs a warning, because in // normal operation a single request never approaches these bounds — hitting // them signals a bug or an abusive request. `op` names the caller. When a // clamp occurs it returns a non-nil *QuotaClamp so callers can additionally // record the event (e.g. on the consume log); the returned pointer is nil for // in-range values. func saturateQuota(value float64, op string) (int, *QuotaClamp) { + return saturateQuotaBounded(value, op, MaxQuota, MinQuota) +} + +func saturateQuotaBounded(value float64, op string, maxQuota int, minQuota int) (int, *QuotaClamp) { var clamp *QuotaClamp switch { case math.IsNaN(value): clamp = &QuotaClamp{Op: op, Kind: QuotaClampNaN, Original: value, Clamped: 0} - case value >= MaxQuota: - clamp = &QuotaClamp{Op: op, Kind: QuotaClampOverflow, Original: value, Clamped: MaxQuota} - case value <= MinQuota: - clamp = &QuotaClamp{Op: op, Kind: QuotaClampUnderflow, Original: value, Clamped: MinQuota} + case value > float64(maxQuota): + clamp = &QuotaClamp{Op: op, Kind: QuotaClampOverflow, Original: value, Clamped: maxQuota} + case value < float64(minQuota): + clamp = &QuotaClamp{Op: op, Kind: QuotaClampUnderflow, Original: value, Clamped: minQuota} default: return int(value), nil } @@ -146,3 +160,16 @@ func QuotaFromDecimalChecked(d decimal.Decimal) (int, *QuotaClamp) { f, _ := d.Round(0).Float64() return saturateQuota(f, "QuotaFromDecimal") } + +// QuotaFromDecimalStrict converts an in-range single-request quota and rejects +// a value that would otherwise be saturated at the int32 boundary. +func QuotaFromDecimalStrict(d decimal.Decimal) (int, error) { + return strictQuota(QuotaFromDecimalChecked(d)) +} + +// WalletQuotaFromDecimalStrict converts wallet and top-up values within the +// JavaScript-safe integer range, which is also exactly representable by float64. +func WalletQuotaFromDecimalStrict(d decimal.Decimal) (int, error) { + f, _ := d.Round(0).Float64() + return strictQuota(saturateQuotaBounded(f, "WalletQuotaFromDecimal", MaxWalletQuota, -MaxWalletQuota)) +} diff --git a/common/quota_math_test.go b/common/quota_math_test.go index 2d8742e6c176..efa8b9ebc2c2 100644 --- a/common/quota_math_test.go +++ b/common/quota_math_test.go @@ -1,6 +1,7 @@ package common import ( + "fmt" "math" "testing" @@ -21,6 +22,7 @@ func TestQuotaFromFloat(t *testing.T) { assert.Equal(t, 42, QuotaFromFloat(42.4)) assert.Equal(t, 42, QuotaFromFloat(42.9)) assert.Equal(t, -42, QuotaFromFloat(-42.9)) + assert.Equal(t, MaxQuota, QuotaFromFloat(float64(math.MaxInt32)+42)) assert.Equal(t, MaxQuota, QuotaFromFloat(overflowingProduct)) assert.Equal(t, MinQuota, QuotaFromFloat(-overflowingProduct)) assert.Equal(t, MaxQuota, QuotaFromFloat(math.Inf(1))) @@ -34,6 +36,7 @@ func TestQuotaRound(t *testing.T) { assert.Equal(t, 42, QuotaRound(41.5)) assert.Equal(t, 43, QuotaRound(42.5)) assert.Equal(t, -43, QuotaRound(-42.5)) + assert.Equal(t, MaxQuota, QuotaRound(float64(math.MaxInt32)+0.5)) assert.Equal(t, MaxQuota, QuotaRound(overflowingProduct)) assert.Equal(t, MinQuota, QuotaRound(-overflowingProduct)) assert.Equal(t, 0, QuotaRound(math.NaN())) @@ -93,7 +96,7 @@ func TestQuotaFromFloatStrictReturnsTypedClampError(t *testing.T) { assert.ErrorContains(t, err, "QuotaFromFloat") assert.ErrorContains(t, err, "overflow") assert.ErrorContains(t, err, "original=") - assert.ErrorContains(t, err, "clamped=2147483647") + assert.ErrorContains(t, err, fmt.Sprintf("clamped=%d", MaxQuota)) } // TestQuotaRoundChecked verifies the rounding entry point reports clamps the @@ -124,3 +127,20 @@ func TestQuotaFromDecimalChecked(t *testing.T) { assert.Equal(t, QuotaClampOverflow, clamp.Kind) } } + +func TestWalletQuotaFromDecimalStrict(t *testing.T) { + quota, err := WalletQuotaFromDecimalStrict(decimal.NewFromInt(4_294_500_000)) + require.NoError(t, err) + assert.Equal(t, 4_294_500_000, quota) + + quota, err = WalletQuotaFromDecimalStrict(decimal.NewFromInt(MaxWalletQuota)) + require.NoError(t, err) + assert.Equal(t, MaxWalletQuota, quota) + + quota, err = WalletQuotaFromDecimalStrict(decimal.NewFromInt(MaxWalletQuota + 1)) + assert.Zero(t, quota) + var clamp *QuotaClamp + require.ErrorAs(t, err, &clamp) + assert.Equal(t, "WalletQuotaFromDecimal", clamp.Op) + assert.Equal(t, QuotaClampOverflow, clamp.Kind) +} diff --git a/common/trusted_proxies.go b/common/trusted_proxies.go new file mode 100644 index 000000000000..363dc019fc28 --- /dev/null +++ b/common/trusted_proxies.go @@ -0,0 +1,54 @@ +package common + +import ( + "errors" + "fmt" + "strings" + + "github.com/gin-gonic/gin" +) + +var defaultTrustedProxyCIDRs = []string{ + "127.0.0.0/8", + "::1", + "10.0.0.0/8", + "172.16.0.0/12", + "192.168.0.0/16", + "fc00::/7", +} + +// ResolveTrustedProxies parses TRUSTED_PROXIES without applying it to an +// engine. The returned slice can be reused by the outer and plugin engines. +func ResolveTrustedProxies(raw string) (trustedProxies []string, usedDefaults bool, err error) { + raw = strings.TrimSpace(raw) + if raw == "" { + return append([]string(nil), defaultTrustedProxyCIDRs...), true, nil + } + if strings.EqualFold(raw, "none") { + return nil, false, nil + } + + parts := strings.Split(raw, ",") + trustedProxies = make([]string, 0, len(parts)) + for _, part := range parts { + trustedProxy := strings.TrimSpace(part) + if trustedProxy == "" { + continue + } + if strings.EqualFold(trustedProxy, "none") { + return nil, false, errors.New("TRUSTED_PROXIES=none must be used alone") + } + trustedProxies = append(trustedProxies, trustedProxy) + } + if len(trustedProxies) == 0 { + return nil, false, errors.New("TRUSTED_PROXIES does not contain an IP address or CIDR") + } + return trustedProxies, false, nil +} + +func ConfigureTrustedProxies(engine *gin.Engine, trustedProxies []string) error { + if err := engine.SetTrustedProxies(trustedProxies); err != nil { + return fmt.Errorf("invalid TRUSTED_PROXIES: %w", err) + } + return nil +} diff --git a/constant/channel.go b/constant/channel.go index 2a6c4a31c138..354dd0e14b2b 100644 --- a/constant/channel.go +++ b/constant/channel.go @@ -58,6 +58,7 @@ const ( ChannelTypeAdvancedCustom = 58 ChannelTypeSub2API = 59 ChannelTypeNewAPI = 60 + ChannelTypeTaskPlugin = 61 ChannelTypeDummy // this one is only for count, do not add any channel after this ) @@ -124,6 +125,14 @@ var ChannelBaseURLs = []string{ "", //58 "", //59 "", //60 + "", //61 +} + +func GetChannelBaseURL(channelType int) string { + if channelType < 0 || channelType >= len(ChannelBaseURLs) { + return "" + } + return ChannelBaseURLs[channelType] } var ChannelTypeNames = map[int]string{ @@ -184,6 +193,7 @@ var ChannelTypeNames = map[int]string{ ChannelTypeAdvancedCustom: "Advanced Custom", ChannelTypeSub2API: "Sub2API", ChannelTypeNewAPI: "New API", + ChannelTypeTaskPlugin: "Task Plugin", } func GetChannelTypeName(channelType int) string { diff --git a/constant/channel_test.go b/constant/channel_test.go new file mode 100644 index 000000000000..92faa26d075d --- /dev/null +++ b/constant/channel_test.go @@ -0,0 +1,12 @@ +package constant + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestGetChannelBaseURLIsBoundsSafe(t *testing.T) { + assert.Empty(t, GetChannelBaseURL(ChannelTypeTaskPlugin)) + assert.Empty(t, GetChannelBaseURL(9999)) +} diff --git a/constant/context_key.go b/constant/context_key.go index b856bc3dda14..93a18ba9af01 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -15,10 +15,12 @@ const ( ContextKeyTokenKey ContextKey = "token_key" ContextKeyTokenId ContextKey = "token_id" ContextKeyTokenGroup ContextKey = "token_group" - ContextKeyTokenSpecificChannelId ContextKey = "specific_channel_id" + 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" /* channel related keys */ ContextKeyChannelId ContextKey = "channel_id" diff --git a/constant/env.go b/constant/env.go index 512bfc31126b..4f8f6755b786 100644 --- a/constant/env.go +++ b/constant/env.go @@ -18,6 +18,11 @@ var GenerateDefaultToken bool var ErrorLogEnabled bool var TaskQueryLimit int var TaskTimeoutMinutes int +var TaskPollMaxFailures = 20 +var TaskPluginProtocolTimeoutSeconds int +var TaskPluginProtocolTickMilliseconds int +var TaskPluginProtocolTickJitterMilliseconds int +var TaskPluginProtocolHeartbeatSeconds int // temporary variable for sora patch, will be removed in future var TaskPricePatches []string diff --git a/constant/task.go b/constant/task.go index ecccf4dfe119..aee856831156 100644 --- a/constant/task.go +++ b/constant/task.go @@ -8,17 +8,35 @@ const ( ) const ( - SunoActionMusic = "MUSIC" - SunoActionLyrics = "LYRICS" - - TaskActionGenerate = "generate" - TaskActionTextGenerate = "textGenerate" - TaskActionFirstTailGenerate = "firstTailGenerate" - TaskActionReferenceGenerate = "referenceGenerate" - TaskActionRemix = "remixGenerate" + TaskActionImageToVideo = "image_to_video" + TaskActionTextToVideo = "text_to_video" + TaskActionFirstTailToVideo = "first_tail_to_video" + TaskActionReferenceToVideo = "reference_to_video" + TaskActionRemix = "remix" ) -var SunoModel2Action = map[string]string{ - "suno_music": SunoActionMusic, - "suno_lyrics": SunoActionLyrics, +var legacyTaskActionAliases = map[string]string{ + "generate": TaskActionImageToVideo, + "textGenerate": TaskActionTextToVideo, + "firstTailGenerate": TaskActionFirstTailToVideo, + "referenceGenerate": TaskActionReferenceToVideo, + "remixGenerate": TaskActionRemix, +} + +// TaskPluginEnabled is the master switch for the whole task-plugin system. +// When disabled, factory and override plugins both stop serving. +var TaskPluginEnabled = true + +// TaskPluginOverrideEnabled controls whether the database override layer is +// active. When disabled, uploaded plugins are ignored and factory plugins are +// used instead; the factory layer is unaffected. +var TaskPluginOverrideEnabled = true + +// NormalizeTaskAction maps persisted legacy action names to the canonical task +// action vocabulary. Unknown platform-specific actions pass through unchanged. +func NormalizeTaskAction(action string) string { + if canonical, ok := legacyTaskActionAliases[action]; ok { + return canonical + } + return action } diff --git a/constant/task_test.go b/constant/task_test.go new file mode 100644 index 000000000000..9b0065df69f1 --- /dev/null +++ b/constant/task_test.go @@ -0,0 +1,27 @@ +package constant + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestNormalizeTaskAction(t *testing.T) { + tests := map[string]string{ + "generate": TaskActionImageToVideo, + "textGenerate": TaskActionTextToVideo, + "firstTailGenerate": TaskActionFirstTailToVideo, + "referenceGenerate": TaskActionReferenceToVideo, + "remixGenerate": TaskActionRemix, + TaskActionTextToVideo: TaskActionTextToVideo, + "MUSIC": "MUSIC", + "custom_action": "custom_action", + "": "", + } + + for input, expected := range tests { + t.Run(input, func(t *testing.T) { + assert.Equal(t, expected, NormalizeTaskAction(input)) + }) + } +} diff --git a/controller/auth_flow_test.go b/controller/auth_flow_test.go index 3fe08b427830..5917d05374dc 100644 --- a/controller/auth_flow_test.go +++ b/controller/auth_flow_test.go @@ -46,6 +46,7 @@ func (*authFlowTestOAuthProvider) IsUserIDTaken(string) bool func (*authFlowTestOAuthProvider) FillUserByProviderID(*model.User, string) error { return nil } func (*authFlowTestOAuthProvider) SetProviderUserID(*model.User, string) {} func (*authFlowTestOAuthProvider) GetProviderPrefix() string { return "flow_" } +func (*authFlowTestOAuthProvider) ProviderUserIDColumn() string { return "" } func setupAuthFlowControllerTest(t *testing.T) *authFlowTestOAuthProvider { t.Helper() diff --git a/controller/billing_option_test.go b/controller/billing_option_test.go new file mode 100644 index 000000000000..cf6b37ea53f1 --- /dev/null +++ b/controller/billing_option_test.go @@ -0,0 +1,203 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "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) { + const pluginKey = "billing-save-probe" + const modelName = "billing-save-model" + source := ` +export const meta = { + apiVersion: 1, key: "billing-save-probe", name: "Billing Save Probe", version: "1.0.0", author: {name: "Test"}, + models: ["billing-save-model"], fetchMode: "per_task", + usageSchema: {seconds: {type: "number", unit: "second"}} +}; +export function buildSubmitRequest() { return {}; } +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +` + _, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{}) + require.NoError(t, err) + t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(pluginKey) }) + + tests := []struct { + name string + expression string + errorText string + }{ + { + name: "invalid syntax", + expression: `tier("base",`, + errorText: "expr compile error", + }, + { + name: "undeclared usage key", + expression: `tier("base", u("clips") * 0.1)`, + errorText: `usage key \"clips\" is not declared`, + }, + } + + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + expressions, marshalErr := common.Marshal(map[string]string{modelName: testCase.expression}) + require.NoError(t, marshalErr) + body, marshalErr := common.Marshal(OptionUpdateRequest{ + Key: "billing_setting.billing_expr", + Value: string(expressions), + }) + require.NoError(t, marshalErr) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodPut, "/api/option/", strings.NewReader(string(body))) + + UpdateOption(context) + + assert.Equal(t, http.StatusOK, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), modelName) + assert.Contains(t, recorder.Body.String(), testCase.errorText) + }) + } +} + +func TestUpdateOptionRejectsUsageExpressionWithoutTaskPlugin(t *testing.T) { + const modelName = "billing-save-model-without-plugin" + expressions, err := common.Marshal(map[string]string{ + modelName: `u("mode") == "std" ? 1 : 2`, + }) + require.NoError(t, err) + body, err := common.Marshal(OptionUpdateRequest{ + Key: "billing_setting.billing_expr", + Value: string(expressions), + }) + require.NoError(t, err) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest( + http.MethodPut, + "/api/option/", + strings.NewReader(string(body)), + ) + + UpdateOption(context) + + assert.Equal(t, http.StatusOK, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), modelName) + assert.Contains(t, recorder.Body.String(), "mode") + 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.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) + 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"}} +}; +export function buildSubmitRequest() { return {}; } +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +` + _, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{}) + require.NoError(t, err) + t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(pluginKey) }) + + mapping := `{"alias-model":"declared-model"}` + require.NoError(t, model.DB.Create(&model.Channel{ + Id: 1, + Type: 54, + Key: "key-1", + Status: common.ChannelStatusEnabled, + Name: "ch-1", + Group: "default", + Models: "alias-model,declared-model", + ModelMapping: &mapping, + }).Error) + model.InitChannelCache() + + saved := map[string]string{} + require.NoError(t, config.GlobalConfig.SaveToDB(func(key, value string) error { + saved[key] = value + return nil + })) + t.Cleanup(func() { + require.NoError(t, config.GlobalConfig.LoadFromDB(saved)) + }) + + putExpr := func(modelName, expression string) *httptest.ResponseRecorder { + t.Helper() + expressions, marshalErr := common.Marshal(map[string]string{modelName: expression}) + require.NoError(t, marshalErr) + body, marshalErr := common.Marshal(OptionUpdateRequest{ + Key: "billing_setting.billing_expr", + Value: string(expressions), + }) + require.NoError(t, marshalErr) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodPut, "/api/option/", strings.NewReader(string(body))) + UpdateOption(context) + return recorder + } + + accepted := putExpr("alias-model", `u("seconds")`) + assert.Equal(t, http.StatusOK, accepted.Code) + assert.Contains(t, accepted.Body.String(), `"success":true`) + + rejectedKey := putExpr("alias-model", `u("clips")`) + 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`) + + unresolvable := putExpr("unknown-alias-model", `u("seconds")`) + assert.Equal(t, http.StatusOK, unresolvable.Code) + assert.Contains(t, unresolvable.Body.String(), `"success":false`) + assert.Contains(t, unresolvable.Body.String(), "no task plugin usage schema") +} diff --git a/controller/channel-billing.go b/controller/channel-billing.go index 62982d2f5ceb..36cd887b805f 100644 --- a/controller/channel-billing.go +++ b/controller/channel-billing.go @@ -5,13 +5,19 @@ import ( "errors" "fmt" "io" + "math" "net/http" "strconv" + "strings" "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/relay/channel/advancedcustom" + relaycommon "github.com/QuantumNous/new-api/relay/common" + relayconstant "github.com/QuantumNous/new-api/relay/constant" + "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/operation_setting" @@ -47,6 +53,13 @@ type OpenAICreditGrants struct { TotalAvailable float64 `json:"total_available"` } +const maxAdvancedCustomBalanceResponseBytes = 256 << 10 + +type channelBalanceResult struct { + Balance float64 + RawResponse string +} + type OpenAIUsageResponse struct { Object string `json:"object"` //DailyCosts []OpenAIUsageDailyCost `json:"daily_costs"` @@ -174,7 +187,7 @@ func updateChannelCloseAIBalance(channel *model.Channel) (float64, error) { return 0, err } response := OpenAICreditGrants{} - err = json.Unmarshal(body, &response) + err = common.Unmarshal(body, &response) if err != nil { return 0, err } @@ -189,7 +202,7 @@ func updateChannelOpenAISBBalance(channel *model.Channel) (float64, error) { return 0, err } response := OpenAISBUsageResponse{} - err = json.Unmarshal(body, &response) + err = common.Unmarshal(body, &response) if err != nil { return 0, err } @@ -213,7 +226,7 @@ func updateChannelAIProxyBalance(channel *model.Channel) (float64, error) { return 0, err } response := AIProxyUserOverviewResponse{} - err = json.Unmarshal(body, &response) + err = common.Unmarshal(body, &response) if err != nil { return 0, err } @@ -232,7 +245,7 @@ func updateChannelAPI2GPTBalance(channel *model.Channel) (float64, error) { return 0, err } response := API2GPTUsageResponse{} - err = json.Unmarshal(body, &response) + err = common.Unmarshal(body, &response) if err != nil { return 0, err } @@ -247,7 +260,7 @@ func updateChannelSiliconFlowBalance(channel *model.Channel) (float64, error) { return 0, err } response := SiliconFlowUsageResponse{} - err = json.Unmarshal(body, &response) + err = common.Unmarshal(body, &response) if err != nil { return 0, err } @@ -269,7 +282,7 @@ func updateChannelDeepSeekBalance(channel *model.Channel) (float64, error) { return 0, err } response := DeepSeekUsageResponse{} - err = json.Unmarshal(body, &response) + err = common.Unmarshal(body, &response) if err != nil { return 0, err } @@ -298,7 +311,7 @@ func updateChannelAIGC2DBalance(channel *model.Channel) (float64, error) { return 0, err } response := APGC2DGPTUsageResponse{} - err = json.Unmarshal(body, &response) + err = common.Unmarshal(body, &response) if err != nil { return 0, err } @@ -313,7 +326,7 @@ func updateChannelOpenRouterBalance(channel *model.Channel) (float64, error) { return 0, err } response := OpenRouterCreditResponse{} - err = json.Unmarshal(body, &response) + err = common.Unmarshal(body, &response) if err != nil { return 0, err } @@ -343,7 +356,7 @@ func updateChannelMoonshotBalance(channel *model.Channel) (float64, error) { } response := MoonshotBalanceResponse{} - err = json.Unmarshal(body, &response) + err = common.Unmarshal(body, &response) if err != nil { return 0, err } @@ -356,8 +369,101 @@ func updateChannelMoonshotBalance(channel *model.Channel) (float64, error) { return availableBalanceUsd, nil } -func updateChannelBalance(channel *model.Channel) (float64, error) { - baseURL := constant.ChannelBaseURLs[channel.Type] +func fetchAdvancedCustomBalance(channel *model.Channel) (channelBalanceResult, error) { + key := strings.TrimSpace(channel.Key) + info := &relaycommon.RelayInfo{ + RelayFormat: types.RelayFormatOpenAI, + RelayMode: relayconstant.RelayModeUnknown, + RequestURLPath: dto.AdvancedCustomBalancePath, + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelType: constant.ChannelTypeAdvancedCustom, + ChannelBaseUrl: channel.GetBaseURL(), + ApiKey: key, + ChannelOtherSettings: channel.GetOtherSettings(), + }, + } + requestURL, headers, err := (&advancedcustom.Adaptor{}).BuildBalanceRequest(info) + if err != nil { + return channelBalanceResult{}, sanitizeFetchModelsError(err, key) + } + if err := applyFetchModelsHeaderOverrides(channel, key, headers); err != nil { + return channelBalanceResult{}, sanitizeFetchModelsError(err, key) + } + + request, err := http.NewRequest(http.MethodGet, requestURL, nil) + if err != nil { + return channelBalanceResult{}, sanitizeFetchModelsError(err, key) + } + for name, values := range headers { + for _, value := range values { + request.Header.Add(name, value) + } + if strings.EqualFold(name, "Host") { + request.Host = headers.Get(name) + } + } + client, err := service.GetHttpClientWithProxy(channel.GetSetting().Proxy) + if err != nil { + return channelBalanceResult{}, sanitizeFetchModelsError(err, key) + } + response, err := client.Do(request) + if err != nil { + return channelBalanceResult{}, sanitizeAdvancedCustomRequestError(err, key, requestURL) + } + defer response.Body.Close() + if response.StatusCode != http.StatusOK { + return channelBalanceResult{}, fmt.Errorf("status code: %d", response.StatusCode) + } + body, err := io.ReadAll(io.LimitReader(response.Body, maxAdvancedCustomBalanceResponseBytes+1)) + if err != nil { + return channelBalanceResult{}, sanitizeAdvancedCustomRequestError(err, key, requestURL) + } + if len(body) > maxAdvancedCustomBalanceResponseBytes { + return channelBalanceResult{}, fmt.Errorf("balance response exceeds %d bytes", maxAdvancedCustomBalanceResponseBytes) + } + + var validated json.RawMessage + if err := common.Unmarshal(body, &validated); err != nil { + return channelBalanceResult{}, fmt.Errorf("invalid balance JSON response: %w", err) + } + if common.GetJsonType(validated) == "object" { + var creditSummary struct { + Object string `json:"object"` + TotalAvailable json.RawMessage `json:"total_available"` + } + if err := common.Unmarshal(body, &creditSummary); err != nil { + return channelBalanceResult{}, fmt.Errorf("invalid balance JSON response: %w", err) + } + if creditSummary.Object == "credit_summary" && + common.GetJsonType(creditSummary.TotalAvailable) == "number" { + var balance float64 + if err := common.Unmarshal(creditSummary.TotalAvailable, &balance); err == nil && + balance >= 0 && + !math.IsNaN(balance) && + !math.IsInf(balance, 0) { + channel.UpdateBalance(balance) + return channelBalanceResult{Balance: balance}, nil + } + } + } + + formatted, err := common.IndentJson(body) + if err != nil { + return channelBalanceResult{}, fmt.Errorf("invalid balance JSON response: %w", err) + } + return channelBalanceResult{RawResponse: string(formatted)}, nil +} + +func updateChannelBalance(channel *model.Channel) (channelBalanceResult, error) { + if channel.Type == constant.ChannelTypeAdvancedCustom { + return fetchAdvancedCustomBalance(channel) + } + balance, err := updateStandardChannelBalance(channel) + return channelBalanceResult{Balance: balance}, err +} + +func updateStandardChannelBalance(channel *model.Channel) (float64, error) { + baseURL := constant.GetChannelBaseURL(channel.Type) if channel.GetBaseURL() == "" { channel.BaseURL = &baseURL } @@ -396,7 +502,7 @@ func updateChannelBalance(channel *model.Channel) (float64, error) { return 0, err } subscription := OpenAISubscriptionResponse{} - err = json.Unmarshal(body, &subscription) + err = common.Unmarshal(body, &subscription) if err != nil { return 0, err } @@ -412,7 +518,7 @@ func updateChannelBalance(channel *model.Channel) (float64, error) { return 0, err } usage := OpenAIUsageResponse{} - err = json.Unmarshal(body, &usage) + err = common.Unmarshal(body, &usage) if err != nil { return 0, err } @@ -432,6 +538,10 @@ func UpdateChannelBalance(c *gin.Context) { common.ApiError(c, err) return } + if channel.Type == constant.ChannelTypeTaskPlugin { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "Task Plugin channels do not support balance queries"}) + return + } if channel.ChannelInfo.IsMultiKey { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -439,16 +549,21 @@ func UpdateChannelBalance(c *gin.Context) { }) return } - balance, err := updateChannelBalance(channel) + result, err := updateChannelBalance(channel) if err != nil { common.ApiError(c, err) return } - c.JSON(http.StatusOK, gin.H{ + response := gin.H{ "success": true, "message": "", - "balance": balance, - }) + } + if result.RawResponse == "" { + response["balance"] = result.Balance + } else { + response["raw_response"] = result.RawResponse + } + c.JSON(http.StatusOK, response) } func updateAllChannelsBalance() error { @@ -467,12 +582,12 @@ func updateAllChannelsBalance() error { //if channel.Type != common.ChannelTypeOpenAI && channel.Type != common.ChannelTypeCustom { // continue //} - balance, err := updateChannelBalance(channel) + result, err := updateChannelBalance(channel) if err != nil { continue - } else { + } else if result.RawResponse == "" { // err is nil & balance <= 0 means quota is used up - if balance <= 0 { + if result.Balance <= 0 { service.DisableChannel(*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, "", channel.GetAutoBan()), "余额不足") } } diff --git a/controller/channel-test.go b/controller/channel-test.go index f6e6bd7f1163..f8fad94cd2de 100644 --- a/controller/channel-test.go +++ b/controller/channel-test.go @@ -7,11 +7,11 @@ import ( "errors" "fmt" "io" - "math" "net/http" "net/http/httptest" "strconv" "strings" + "sync" "time" "github.com/QuantumNous/new-api/common" @@ -27,7 +27,6 @@ import ( "github.com/QuantumNous/new-api/relaykit/types" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting/operation_setting" - "github.com/QuantumNous/new-api/setting/ratio_setting" hosttypes "github.com/QuantumNous/new-api/types" "github.com/samber/lo" @@ -42,14 +41,11 @@ type testResult struct { newAPIError *types.NewAPIError } -func normalizeChannelTestEndpoint(channel *model.Channel, modelName, endpointType string) string { +func normalizeChannelTestEndpoint(channel *model.Channel, endpointType string) string { normalized := strings.TrimSpace(endpointType) if normalized != "" { return normalized } - if strings.HasSuffix(modelName, ratio_setting.CompactModelSuffix) { - return string(constant.EndpointTypeOpenAIResponseCompact) - } if channel != nil && channel.Type == constant.ChannelTypeCodex { return string(constant.EndpointTypeOpenAIResponse) } @@ -86,6 +82,7 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te constant.ChannelTypeJimeng, constant.ChannelTypeDoubaoVideo, constant.ChannelTypeVidu, + constant.ChannelTypeTaskPlugin, } if lo.Contains(unsupportedTestChannelTypes, channel.Type) { channelTypeName := constant.GetChannelTypeName(channel.Type) @@ -111,7 +108,7 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te } } - endpointType = normalizeChannelTestEndpoint(channel, testModel, endpointType) + endpointType = normalizeChannelTestEndpoint(channel, endpointType) requestPath := "/v1/chat/completions" @@ -146,15 +143,12 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te requestPath = "/v1/responses" } - // responses compaction models (must use /v1/responses/compact) - if strings.HasSuffix(testModel, ratio_setting.CompactModelSuffix) { - requestPath = "/v1/responses/compact" - } } - if strings.HasPrefix(requestPath, "/v1/responses/compact") { - testModel = ratio_setting.WithCompactModelSuffix(testModel) + // Gemini 原生流式通过 URL action(:streamGenerateContent)表达而非请求体字段, + // GeminiChatRequest.IsStream 依据请求 URL 判定,合成请求路径需与生产入口保持一致 + if isStream && constant.EndpointType(endpointType) == constant.EndpointTypeGemini { + requestPath = strings.Replace(requestPath, ":generateContent", ":streamGenerateContent", 1) } - c.Request = httptest.NewRequestWithContext(ctx, http.MethodPost, requestPath, nil) cache, err := model.GetUserCache(testUserID) @@ -265,6 +259,13 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te newAPIError: types.NewError(err, types.ErrorCodeChannelModelMappedError), } } + if err := helper.ApplyReasoningModelSuffix(c, info, request); err != nil { + return testResult{ + context: c, + localErr: err, + newAPIError: types.NewErrorWithStatusCode(err, types.ErrorCodeConvertRequestFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()), + } + } testModel = info.UpstreamModelName // 更新请求中的模型名称 @@ -272,7 +273,7 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te apiType, _ := common.ChannelType2APIType(channel.Type) if info.RelayMode == relayconstant.RelayModeResponsesCompact && - !common.IsResponsesCompactAPIType(apiType) { + !common.SupportsResponsesCompact(channel.Type, apiType) { return testResult{ context: c, localErr: fmt.Errorf("responses compaction test is not supported for api type %d", apiType), @@ -371,14 +372,18 @@ func testChannel(ctx context.Context, channel *model.Channel, testUserID int, te } } default: - // Chat/Completion 等其他请求类型 - if generalReq, ok := request.(*dto.GeneralOpenAIRequest); ok { - convertedRequest, err = adaptor.ConvertOpenAIRequest(c, info, generalReq) - } else { + switch req := request.(type) { + case *dto.GeneralOpenAIRequest: + convertedRequest, err = adaptor.ConvertOpenAIRequest(c, info, req) + case *dto.ClaudeRequest: + convertedRequest, err = adaptor.ConvertClaudeRequest(c, info, req) + case *dto.GeminiChatRequest: + convertedRequest, err = adaptor.ConvertGeminiRequest(c, info, req) + default: return testResult{ context: c, - localErr: errors.New("invalid general request type"), - newAPIError: types.NewError(errors.New("invalid general request type"), types.ErrorCodeConvertRequestFailed), + localErr: errors.New("invalid chat request type"), + newAPIError: types.NewError(errors.New("invalid chat request type"), types.ErrorCodeConvertRequestFailed), } } } @@ -542,18 +547,19 @@ func settleTestQuota(info *relaycommon.RelayInfo, priceData hosttypes.PriceData, quota := 0 if !priceData.UsePrice { - quota = usage.PromptTokens + int(math.Round(float64(usage.CompletionTokens)*priceData.CompletionRatio)) - quota = int(math.Round(float64(quota) * priceData.ModelRatio)) + completionQuota := common.QuotaRound(float64(usage.CompletionTokens) * priceData.CompletionRatio) + quota = common.QuotaRound(float64(usage.PromptTokens) + float64(completionQuota)) + quota = common.QuotaRound(float64(quota) * priceData.ModelRatio) if priceData.ModelRatio != 0 && quota <= 0 { quota = 1 } return quota, nil } - return int(priceData.ModelPrice * common.QuotaPerUnit), nil + return common.QuotaFromFloat(priceData.ModelPrice * common.QuotaPerUnit), nil } -func buildTestLogOther(c *gin.Context, info *relaycommon.RelayInfo, priceData hosttypes.PriceData, usage *dto.Usage, tieredResult *billingexpr.TieredResult) map[string]interface{} { +func buildTestLogOther(c *gin.Context, info *relaycommon.RelayInfo, priceData hosttypes.PriceData, usage *dto.Usage, tieredResult *billingexpr.TieredResult) *model.LogOther { other := service.GenerateTextOtherInfo(c, info, priceData.ModelRatio, priceData.GroupRatioInfo.GroupRatio, priceData.CompletionRatio, usage.PromptTokensDetails.CachedTokens, priceData.CacheRatio, priceData.ModelPrice, priceData.GroupRatioInfo.GroupSpecialRatio) if tieredResult != nil { @@ -733,12 +739,31 @@ func buildTestRequest(model string, endpointType string, channel *model.Channel, Model: model, Input: testResponsesInput, } - case constant.EndpointTypeAnthropic, constant.EndpointTypeGemini, constant.EndpointTypeOpenAI: - // 返回 GeneralOpenAIRequest - maxTokens := uint(16) - if constant.EndpointType(endpointType) == constant.EndpointTypeGemini { - maxTokens = 3000 + case constant.EndpointTypeAnthropic: + return &dto.ClaudeRequest{ + Model: model, + Stream: lo.ToPtr(isStream), + MaxTokens: lo.ToPtr(uint(16)), + Messages: []dto.ClaudeMessage{ + { + Role: "user", + Content: "hi", + }, + }, } + case constant.EndpointTypeGemini: + return &dto.GeminiChatRequest{ + Contents: []dto.GeminiChatContent{ + { + Role: "user", + Parts: []dto.GeminiPart{{Text: "hi"}}, + }, + }, + GenerationConfig: dto.GeminiChatGenerationConfig{ + MaxOutputTokens: lo.ToPtr(uint(3000)), + }, + } + case constant.EndpointTypeOpenAI: req := &dto.GeneralOpenAIRequest{ Model: model, Stream: lo.ToPtr(isStream), @@ -748,7 +773,7 @@ func buildTestRequest(model string, endpointType string, channel *model.Channel, Content: "hi", }, }, - MaxTokens: lo.ToPtr(maxTokens), + MaxTokens: lo.ToPtr(uint(16)), } if isStream { req.StreamOptions = &dto.StreamOptions{IncludeUsage: true} @@ -778,14 +803,6 @@ func buildTestRequest(model string, endpointType string, channel *model.Channel, } } - // Responses compaction models (must use /v1/responses/compact) - if strings.HasSuffix(model, ratio_setting.CompactModelSuffix) { - return &dto.OpenAIResponsesCompactionRequest{ - Model: model, - Input: testResponsesInput, - } - } - // Responses-only models (e.g. codex series) if strings.Contains(strings.ToLower(model), "codex") { return &dto.OpenAIResponsesRequest{ @@ -900,92 +917,167 @@ type channelTestSummary struct { Enabled int `json:"enabled"` } -// performChannelTests runs the channel test loop synchronously, honoring ctx -// cancellation so a system-task runner that loses its lease stops promptly. When -// report is non-nil it is called after each channel with (processed, total) so -// the system task can surface progress. -func performChannelTests(ctx context.Context, channels []*model.Channel, testUserID int, allowDisable bool, report func(processed, total int)) channelTestSummary { +func testChannelForHealthCheck(ctx context.Context, channel *model.Channel, testUserID int, allowDisable bool, disableThreshold int64) channelTestSummary { summary := channelTestSummary{} - var disableThreshold = int64(common.ChannelDisableThreshold * 1000) - if disableThreshold == 0 { - disableThreshold = 10000000 // a impossible value + isChannelEnabled := channel.Status == common.ChannelStatusEnabled + tik := time.Now() + result := testChannel(ctx, channel, testUserID, "", "", shouldUseStreamForAutomaticChannelTest(channel)) + milliseconds := time.Since(tik).Milliseconds() + if ctx.Err() != nil { + return summary } - total := len(channels) - for index, channel := range channels { - if ctx != nil && ctx.Err() != nil { - break - } - if report != nil { - report(index, total) // channels completed before this one - } - if channel.Status == common.ChannelStatusManuallyDisabled { - continue - } - isChannelEnabled := channel.Status == common.ChannelStatusEnabled - tik := time.Now() - result := testChannel(ctx, channel, testUserID, "", "", shouldUseStreamForAutomaticChannelTest(channel)) - tok := time.Now() - milliseconds := tok.Sub(tik).Milliseconds() - if ctx != nil && ctx.Err() != nil { - break - } + summary.Tested++ - summary.Tested++ + shouldBanChannel := false + newAPIError := result.newAPIError + if newAPIError != nil { + shouldBanChannel = service.ShouldDisableChannel(result.newAPIError) + } - shouldBanChannel := false - newAPIError := result.newAPIError - // request error disables the channel - if newAPIError != nil { - shouldBanChannel = service.ShouldDisableChannel(result.newAPIError) + if common.AutomaticDisableChannelEnabled && !shouldBanChannel { + if milliseconds > disableThreshold { + err := fmt.Errorf("响应时间 %.2fs 超过阈值 %.2fs", float64(milliseconds)/1000.0, float64(disableThreshold)/1000.0) + newAPIError = types.NewOpenAIError(err, types.ErrorCodeChannelResponseTimeExceeded, http.StatusRequestTimeout) + shouldBanChannel = true } + } - // 当错误检查通过,才检查响应时间 - if common.AutomaticDisableChannelEnabled && !shouldBanChannel { - if milliseconds > disableThreshold { - err := fmt.Errorf("响应时间 %.2fs 超过阈值 %.2fs", float64(milliseconds)/1000.0, float64(disableThreshold)/1000.0) - newAPIError = types.NewOpenAIError(err, types.ErrorCodeChannelResponseTimeExceeded, http.StatusRequestTimeout) - shouldBanChannel = true - } - } + if newAPIError == nil { + summary.Succeeded++ + } else { + summary.Failed++ + } - if newAPIError == nil { - summary.Succeeded++ - } else { - summary.Failed++ - } + if allowDisable && isChannelEnabled && shouldBanChannel && channel.GetAutoBan() { + processChannelError(result.context, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(result.context, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError, nil) + summary.Disabled++ + } - // disable channel - if allowDisable && isChannelEnabled && shouldBanChannel && channel.GetAutoBan() { - processChannelError(result.context, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(result.context, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError) - summary.Disabled++ - } + if result.localErr == nil && !isChannelEnabled && service.ShouldEnableChannel(newAPIError, channel.Status) { + service.EnableChannel(channel.Id, common.GetContextKeyString(result.context, constant.ContextKeyChannelKey), channel.Name) + summary.Enabled++ + } - // enable channel - if result.localErr == nil && !isChannelEnabled && service.ShouldEnableChannel(newAPIError, channel.Status) { - service.EnableChannel(channel.Id, common.GetContextKeyString(result.context, constant.ContextKeyChannelKey), channel.Name) - summary.Enabled++ - } + channel.UpdateResponseTime(milliseconds) + return summary +} - channel.UpdateResponseTime(milliseconds) - if common.RequestInterval > 0 { - if ctx == nil { - time.Sleep(common.RequestInterval) - } else { +// runChannelTestWorkers executes independent channel tests with bounded +// concurrency. Results and progress are reduced by the caller goroutine, so +// summary counts and the progress reporter remain serialized. +func runChannelTestWorkers( + ctx context.Context, + channels []*model.Channel, + concurrency int, + run func(context.Context, *model.Channel) channelTestSummary, + report func(processed, total int), +) channelTestSummary { + if ctx == nil { + ctx = context.Background() + } + total := len(channels) + if report != nil { + report(0, total) + } + if total == 0 { + return channelTestSummary{} + } + + workerCount := min(operation_setting.NormalizeChannelTestConcurrency(concurrency), total) + jobs := make(chan *model.Channel) + results := make(chan channelTestSummary) + + var workers sync.WaitGroup + workers.Add(workerCount) + for range workerCount { + go func() { + defer workers.Done() + for { select { case <-ctx.Done(): - return summary - case <-time.After(common.RequestInterval): + return + case channel, ok := <-jobs: + if !ok { + return + } + if ctx.Err() != nil { + return + } + + result := channelTestSummary{} + if channel != nil && channel.Status != common.ChannelStatusManuallyDisabled { + result = run(ctx, channel) + } + + results <- result + + if common.RequestInterval > 0 { + select { + case <-ctx.Done(): + return + case <-time.After(common.RequestInterval): + } + } } } - } + }() } - if report != nil && (ctx == nil || ctx.Err() == nil) { - report(total, total) // mark complete only when the full set was tested + + go func() { + defer close(jobs) + for _, channel := range channels { + select { + case <-ctx.Done(): + return + case jobs <- channel: + } + } + }() + + go func() { + workers.Wait() + close(results) + }() + + summary := channelTestSummary{} + processed := 0 + for result := range results { + summary.Tested += result.Tested + summary.Succeeded += result.Succeeded + summary.Failed += result.Failed + summary.Disabled += result.Disabled + summary.Enabled += result.Enabled + processed++ + if report != nil && ctx.Err() == nil { + report(processed, total) + } } return summary } +// performChannelTests runs channel health checks with the configured bounded +// concurrency and honors cancellation when a system-task runner loses its +// lease. +func performChannelTests(ctx context.Context, channels []*model.Channel, testUserID int, allowDisable bool, concurrency int, report func(processed, total int)) channelTestSummary { + if ctx == nil { + ctx = context.Background() + } + disableThreshold := int64(common.ChannelDisableThreshold * 1000) + if disableThreshold == 0 { + disableThreshold = 10000000 // an impossible value + } + return runChannelTestWorkers( + ctx, + channels, + concurrency, + func(ctx context.Context, channel *model.Channel) channelTestSummary { + return testChannelForHealthCheck(ctx, channel, testUserID, allowDisable, disableThreshold) + }, + report, + ) +} + // runChannelTestTask runs one synchronous channel test cycle for the system task // runner (both the scheduled job and the manual "test all channels" trigger go // through here). It honors ctx cancellation so a runner that loses its lease @@ -1008,7 +1100,8 @@ func runChannelTestTask(ctx context.Context, mode string, notify bool, report fu } selected := selectChannelsForAutomaticTest(channels, mode) allowDisable := mode != operation_setting.ChannelTestModePassiveRecovery - summary := performChannelTests(ctx, selected, testUserID, allowDisable, report) + concurrency := operation_setting.GetMonitorSetting().ChannelTestConcurrency + summary := performChannelTests(ctx, selected, testUserID, allowDisable, concurrency, report) if notify && (ctx == nil || ctx.Err() == nil) { service.NotifyRootUser(dto.NotifyTypeChannelTest, "通道测试完成", "所有通道测试已完成") } @@ -1021,6 +1114,9 @@ func selectChannelsForAutomaticTest(channels []*model.Channel, mode string) []*m if channel.Status == common.ChannelStatusManuallyDisabled { continue } + if mode == operation_setting.ChannelTestModeAutoBanOnly && !channel.GetAutoBan() { + continue + } if mode == operation_setting.ChannelTestModePassiveRecovery && channel.Status != common.ChannelStatusAutoDisabled { continue } diff --git a/controller/channel.go b/controller/channel.go index 3a1e58328923..19ddca8e6a07 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -13,6 +13,7 @@ import ( "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" relaychannel "github.com/QuantumNous/new-api/relay/channel" "github.com/QuantumNous/new-api/relay/channel/ollama" relaycommon "github.com/QuantumNous/new-api/relay/common" @@ -480,6 +481,21 @@ func validateChannel(channel *model.Channel, isAdd bool) error { if err := channel.ValidateSettings(); err != nil { return fmt.Errorf("渠道额外设置[channel setting] 格式错误:%s", err.Error()) } + if channel.Type == constant.ChannelTypeTaskPlugin { + pluginKey := strings.TrimSpace(channel.GetSetting().TaskPluginKey) + if pluginKey == "" { + return fmt.Errorf("task plugin key is required") + } + if len(pluginKey) > 30 { + return fmt.Errorf("task plugin key must not exceed 30 characters") + } + if _, ok := jsplugin.DefaultRegistry.Get(pluginKey); !ok { + return fmt.Errorf("task plugin %q is not registered", pluginKey) + } + if channel.BaseURL == nil || strings.TrimSpace(*channel.BaseURL) == "" { + return fmt.Errorf("base URL is required for task plugin channels") + } + } if channel.Type == constant.ChannelTypeNewAPI && strings.TrimSpace(channel.GetBaseURL()) == "" { return fmt.Errorf("New API channel base URL cannot be empty") @@ -617,6 +633,15 @@ func AddChannel(c *gin.Context) { return } + if addChannelRequest.Channel != nil && addChannelRequest.Channel.Type == constant.ChannelTypeTaskPlugin && + !authz.Can(c.GetInt("id"), c.GetInt("role"), authz.TaskPluginBind) { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "task plugin channels require the task_plugin.bind permission", + }) + return + } + // 使用统一的校验函数 if err := validateChannel(addChannelRequest.Channel, true); err != nil { c.JSON(http.StatusOK, gin.H{ @@ -964,6 +989,15 @@ func UpdateChannel(c *gin.Context) { } clearChannelReadOnlyFields(&channel, requestData) + if channel.Type == constant.ChannelTypeTaskPlugin && + !authz.Can(c.GetInt("id"), c.GetInt("role"), authz.TaskPluginBind) { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "task plugin channels require the task_plugin.bind permission", + }) + return + } + // 使用统一的校验函数 if err := validateChannel(&channel.Channel, false); err != nil { c.JSON(http.StatusOK, gin.H{ @@ -1299,7 +1333,7 @@ func FetchModels(c *gin.Context) { baseURL = strings.TrimSpace(*req.BaseURL) } if baseURL == "" { - baseURL = constant.ChannelBaseURLs[req.Type] + baseURL = constant.GetChannelBaseURL(req.Type) } key := strings.TrimSpace(req.Key) @@ -1424,6 +1458,11 @@ func CopyChannel(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"success": false, "message": "获取渠道信息失败,请稍后重试"}) return } + if origin.Type == constant.ChannelTypeTaskPlugin && + !authz.Can(c.GetInt("id"), c.GetInt("role"), authz.TaskPluginBind) { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "task plugin channels require the task_plugin.bind permission"}) + return + } // clone channel clone := *origin // shallow copy is sufficient as we will overwrite primitives @@ -2010,7 +2049,7 @@ func OllamaPullModel(c *gin.Context) { return } - baseURL := constant.ChannelBaseURLs[channel.Type] + baseURL := constant.GetChannelBaseURL(channel.Type) if channel.GetBaseURL() != "" { baseURL = channel.GetBaseURL() } @@ -2073,7 +2112,7 @@ func OllamaPullModelStream(c *gin.Context) { return } - baseURL := constant.ChannelBaseURLs[channel.Type] + baseURL := constant.GetChannelBaseURL(channel.Type) if channel.GetBaseURL() != "" { baseURL = channel.GetBaseURL() } @@ -2155,7 +2194,7 @@ func OllamaDeleteModel(c *gin.Context) { return } - baseURL := constant.ChannelBaseURLs[channel.Type] + baseURL := constant.GetChannelBaseURL(channel.Type) if channel.GetBaseURL() != "" { baseURL = channel.GetBaseURL() } @@ -2204,7 +2243,7 @@ func OllamaVersion(c *gin.Context) { return } - baseURL := constant.ChannelBaseURLs[channel.Type] + baseURL := constant.GetChannelBaseURL(channel.Type) if channel.GetBaseURL() != "" { baseURL = channel.GetBaseURL() } diff --git a/controller/channel_pin_retry_test.go b/controller/channel_pin_retry_test.go new file mode 100644 index 000000000000..f79b7b8b1123 --- /dev/null +++ b/controller/channel_pin_retry_test.go @@ -0,0 +1,95 @@ +package controller + +import ( + "errors" + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/relaykit/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestShouldRetryHonorsPinRetryMode(t *testing.T) { + openaiErr := types.NewOpenAIError(errors.New("upstream"), types.ErrorCodeBadResponseStatusCode, http.StatusInternalServerError) + + c := newPinRetryContext() + assert.True(t, shouldRetry(c, openaiErr, 1)) + + origin := newPinRetryContext() + service.GetChannelConstraints(origin).AddPin(dto.ChannelPin{ + ChannelId: 2, + Source: dto.PinSourceOriginTask, + Rank: dto.PinRankOriginTask, + RetryMode: dto.PinRetrySameChannel, + }) + assert.True(t, shouldRetry(origin, openaiErr, 1), "origin pin retries on the same channel") + + token := newPinRetryContext() + service.GetChannelConstraints(token).AddPin(dto.ChannelPin{ + ChannelId: 1, + Source: dto.PinSourceToken, + Rank: dto.PinRankToken, + RetryMode: dto.PinRetrySingleAttempt, + }) + assert.False(t, shouldRetry(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)) + + origin := newPinRetryContext() + service.GetChannelConstraints(origin).AddPin(dto.ChannelPin{ + ChannelId: 2, + Source: dto.PinSourceOriginTask, + Rank: dto.PinRankOriginTask, + RetryMode: dto.PinRetrySameChannel, + }) + assert.True(t, shouldRetryTaskRelay(origin, 2, taskErr, 1)) + + token := newPinRetryContext() + service.GetChannelConstraints(token).AddPin(dto.ChannelPin{ + ChannelId: 1, + Source: dto.PinSourceToken, + Rank: dto.PinRankToken, + RetryMode: dto.PinRetrySingleAttempt, + }) + assert.False(t, shouldRetryTaskRelay(token, 1, taskErr, 1)) +} + +func TestSameChannelPinsMergeToStricterRetryMode(t *testing.T) { + c := newPinRetryContext() + constraints := service.GetChannelConstraints(c) + constraints.AddPin(dto.ChannelPin{ + ChannelId: 7, + Source: dto.PinSourceOriginTask, + Rank: dto.PinRankOriginTask, + RetryMode: dto.PinRetrySameChannel, + }) + constraints.AddPin(dto.ChannelPin{ + ChannelId: 7, + Source: dto.PinSourceToken, + Rank: dto.PinRankToken, + RetryMode: dto.PinRetrySingleAttempt, + }) + pin, found, overridden := constraints.ResolvedPin() + require.True(t, found) + 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)) +} + +func newPinRetryContext() *gin.Context { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + return c +} diff --git a/controller/channel_task_plugin_bind_test.go b/controller/channel_task_plugin_bind_test.go new file mode 100644 index 000000000000..b8f326661594 --- /dev/null +++ b/controller/channel_task_plugin_bind_test.go @@ -0,0 +1,131 @@ +package controller + +import ( + "fmt" + "net/http" + "net/http/httptest" + "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/pkg/jsplugin" + "github.com/QuantumNous/new-api/service/authz" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupTaskPluginBindChannelTest(t *testing.T) { + t.Helper() + wasMaster := common.IsMasterNode + common.IsMasterNode = true + previousRedisEnabled := common.RedisEnabled + common.RedisEnabled = false + originalDB, originalLogDB := model.DB, model.LOG_DB + database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := database.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.CasbinRule{}, &model.AuthzRole{}, &model.Log{}, &model.User{})) + model.DB = database + model.LOG_DB = database + require.NoError(t, authz.Init(database)) + t.Cleanup(func() { + common.IsMasterNode = wasMaster + common.RedisEnabled = previousRedisEnabled + model.DB = originalDB + model.LOG_DB = originalLogDB + }) +} + +func postAddChannel(t *testing.T, userID, role int, body string) *httptest.ResponseRecorder { + t.Helper() + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Set("id", userID) + context.Set("role", role) + context.Request = httptest.NewRequest(http.MethodPost, "/api/channel", strings.NewReader(body)) + context.Request.Header.Set("Content-Type", "application/json") + AddChannel(context) + return recorder +} + +func TestAddChannelTaskPluginRequiresBindPermission(t *testing.T) { + setupTaskPluginBindChannelTest(t) + const key = "channel-bind" + source := ` +export const meta = {apiVersion: 1, key: "channel-bind", name: "Bind", version: "1.0.0", author: {name: "Test"}, models: ["doc"], fetchMode: "per_task"}; +export function buildSubmitRequest() { return {}; } +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +` + _, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{}) + require.NoError(t, err) + t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(key) }) + + taskPluginBody := `{"mode":"single","channel":{"type":61,"name":"plugin-channel","key":"sk","models":"doc","group":"default","base_url":"https://example.com","setting":"{\"task_plugin_key\":\"channel-bind\"}"}}` + openaiBody := `{"mode":"single","channel":{"type":1,"name":"openai-channel","key":"sk","models":"gpt","group":"default"}}` + + adminDenied := postAddChannel(t, 2, common.RoleAdminUser, taskPluginBody) + assert.Contains(t, adminDenied.Body.String(), "task plugin channels require the task_plugin.bind permission") + assert.Contains(t, adminDenied.Body.String(), `"success":false`) + + rootAllowed := postAddChannel(t, 1, common.RoleRootUser, taskPluginBody) + assert.Contains(t, rootAllowed.Body.String(), `"success":true`) + assert.NotContains(t, rootAllowed.Body.String(), "task_plugin.bind") + + adminOtherType := postAddChannel(t, 2, common.RoleAdminUser, openaiBody) + assert.Contains(t, adminOtherType.Body.String(), `"success":true`) + assert.NotContains(t, adminOtherType.Body.String(), "task_plugin.bind") +} + +func TestUpdateChannelTaskPluginRequiresBindPermission(t *testing.T) { + setupTaskPluginBindChannelTest(t) + const key = "channel-bind-update" + source := ` +export const meta = {apiVersion: 1, key: "channel-bind-update", name: "Bind", version: "1.0.0", author: {name: "Test"}, models: ["doc"], fetchMode: "per_task"}; +export function buildSubmitRequest() { return {}; } +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +` + _, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{}) + require.NoError(t, err) + t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(key) }) + + baseURL := "https://example.com" + setting := `{"task_plugin_key":"channel-bind-update"}` + channel := model.Channel{ + Type: constant.ChannelTypeTaskPlugin, + Status: common.ChannelStatusEnabled, + Name: "existing-plugin", + Models: "doc", + Group: "default", + Key: "sk", + BaseURL: &baseURL, + Setting: &setting, + } + require.NoError(t, channel.Insert()) + + payload := fmt.Sprintf( + `{"id":%d,"type":61,"name":"existing-plugin","key":"sk","models":"doc","group":"default","base_url":"https://example.com","setting":"{\"task_plugin_key\":\"channel-bind-update\"}"}`, + channel.Id, + ) + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Set("id", 2) + context.Set("role", common.RoleAdminUser) + context.Request = httptest.NewRequest(http.MethodPut, "/api/channel", strings.NewReader(payload)) + context.Request.Header.Set("Content-Type", "application/json") + UpdateChannel(context) + assert.Contains(t, recorder.Body.String(), "task plugin channels require the task_plugin.bind permission") + assert.Contains(t, recorder.Body.String(), `"success":false`) +} diff --git a/controller/channel_task_plugin_validation_test.go b/controller/channel_task_plugin_validation_test.go new file mode 100644 index 000000000000..6dfe065040f5 --- /dev/null +++ b/controller/channel_task_plugin_validation_test.go @@ -0,0 +1,41 @@ +package controller + +import ( + "strings" + "testing" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/pkg/jsplugin" + "github.com/stretchr/testify/require" +) + +func TestValidateTaskPluginChannel(t *testing.T) { + source := ` +export const meta = {apiVersion: 1, key: "channel-validation", name: "Validation", version: "1.0.0", author: {name: "Test"}, models: ["doc"], fetchMode: "per_task"}; +export function buildSubmitRequest() { return {}; } +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +` + _, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{}) + require.NoError(t, err) + t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("channel-validation") }) + baseURL := "https://example.com" + + channel := &model.Channel{Type: constant.ChannelTypeTaskPlugin, BaseURL: &baseURL} + require.ErrorContains(t, validateChannel(channel, false), "task plugin key is required") + + missing := `{"task_plugin_key":"missing"}` + channel.Setting = &missing + require.ErrorContains(t, validateChannel(channel, false), "is not registered") + + longKey := `{"task_plugin_key":"` + strings.Repeat("x", 31) + `"}` + channel.Setting = &longKey + require.ErrorContains(t, validateChannel(channel, false), "must not exceed 30") + + valid := `{"task_plugin_key":"channel-validation"}` + channel.Setting = &valid + channel.BaseURL = nil + require.ErrorContains(t, validateChannel(channel, false), "base URL is required") +} diff --git a/controller/channel_test_internal_test.go b/controller/channel_test_internal_test.go index abbd2d237070..5f19aeb86610 100644 --- a/controller/channel_test_internal_test.go +++ b/controller/channel_test_internal_test.go @@ -2,9 +2,11 @@ package controller import ( "bytes" + "context" "fmt" "net/http" "net/http/httptest" + "sync/atomic" "testing" "github.com/QuantumNous/new-api/common" @@ -95,23 +97,25 @@ func TestNewAPIChannelRegistration(t *testing.T) { assert.Empty(t, constant.ChannelBaseURLs[constant.ChannelTypeNewAPI]) } -func TestResponsesCompactAPITypeSupport(t *testing.T) { +func TestResponsesCompactChannelSupport(t *testing.T) { tests := []struct { - name string - apiType int - want bool + name string + channelType int + apiType int + want bool }{ - {name: "OpenAI", apiType: constant.APITypeOpenAI, want: true}, - {name: "Codex", apiType: constant.APITypeCodex, want: true}, - {name: "Advanced Custom", apiType: constant.APITypeAdvancedCustom, want: true}, - {name: "Sub2API", apiType: constant.APITypeSub2API, want: true}, - {name: "New API", apiType: constant.APITypeNewAPI, want: true}, - {name: "Anthropic", apiType: constant.APITypeAnthropic, want: false}, + {name: "OpenAI", channelType: constant.ChannelTypeOpenAI, apiType: constant.APITypeOpenAI, want: true}, + {name: "Azure", channelType: constant.ChannelTypeAzure, apiType: constant.APITypeOpenAI, want: true}, + {name: "Codex", channelType: constant.ChannelTypeCodex, apiType: constant.APITypeCodex, want: true}, + {name: "Advanced Custom", channelType: constant.ChannelTypeAdvancedCustom, apiType: constant.APITypeAdvancedCustom, want: true}, + {name: "Sub2API", channelType: constant.ChannelTypeSub2API, apiType: constant.APITypeSub2API, want: true}, + {name: "New API", channelType: constant.ChannelTypeNewAPI, apiType: constant.APITypeNewAPI, want: true}, + {name: "Anthropic", channelType: constant.ChannelTypeAnthropic, apiType: constant.APITypeAnthropic, want: false}, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - assert.Equal(t, test.want, common.IsResponsesCompactAPIType(test.apiType)) + assert.Equal(t, test.want, common.SupportsResponsesCompact(test.channelType, test.apiType)) }) } } @@ -265,13 +269,21 @@ func TestBuildTestLogOtherInjectsTieredInfo(t *testing.T) { }, } + requestRules := []billingexpr.RequestRuleTrace{{ + Cond: `param("service_tier") == "fast"`, + Multiplier: 2, + Matched: true, + }} other := buildTestLogOther(ctx, info, priceData, usage, &billingexpr.TieredResult{ - MatchedTier: "base", + MatchedTier: "base", + RequestRules: requestRules, }) - require.Equal(t, "tiered_expr", other["billing_mode"]) - require.Equal(t, "base", other["matched_tier"]) - require.NotEmpty(t, other["expr_b64"]) + fields := other.Snapshot() + require.Equal(t, "tiered_expr", fields["billing_mode"]) + require.Equal(t, "base", fields["matched_tier"]) + require.Equal(t, requestRules, fields["request_rules"]) + require.NotEmpty(t, fields["expr_b64"]) } func TestResolveChannelTestUserIDUsesRequestUser(t *testing.T) { @@ -312,6 +324,129 @@ func TestSelectChannelsForAutomaticTestScheduledSkipsManualDisabled(t *testing.T require.Equal(t, 2, selected[1].Id) } +func TestSelectChannelsForAutomaticTestAutoBanOnlyUsesEligibleChannels(t *testing.T) { + autoBanEnabled := 1 + autoBanDisabled := 0 + channels := []*model.Channel{ + {Id: 1, Status: common.ChannelStatusEnabled, AutoBan: &autoBanEnabled}, + {Id: 2, Status: common.ChannelStatusEnabled, AutoBan: &autoBanDisabled}, + {Id: 3, Status: common.ChannelStatusAutoDisabled, AutoBan: &autoBanEnabled}, + {Id: 4, Status: common.ChannelStatusManuallyDisabled, AutoBan: &autoBanEnabled}, + {Id: 5, Status: common.ChannelStatusEnabled}, + } + + selected := selectChannelsForAutomaticTest(channels, operation_setting.ChannelTestModeAutoBanOnly) + + require.Len(t, selected, 2) + require.Equal(t, 1, selected[0].Id) + require.Equal(t, 3, selected[1].Id) +} + +func TestRunChannelTestWorkersHonorsConfiguredConcurrency(t *testing.T) { + originalInterval := common.RequestInterval + common.RequestInterval = 0 + t.Cleanup(func() { common.RequestInterval = originalInterval }) + + channels := []*model.Channel{ + {Id: 1, Status: common.ChannelStatusEnabled}, + {Id: 2, Status: common.ChannelStatusEnabled}, + {Id: 3, Status: common.ChannelStatusEnabled}, + {Id: 4, Status: common.ChannelStatusEnabled}, + } + started := make(chan struct{}, len(channels)) + release := make(chan struct{}) + var active atomic.Int32 + var maxActive atomic.Int32 + progress := make([]int, 0, len(channels)+1) + summaryResult := make(chan channelTestSummary, 1) + + go func() { + summaryResult <- runChannelTestWorkers( + context.Background(), + channels, + 2, + func(_ context.Context, _ *model.Channel) channelTestSummary { + current := active.Add(1) + defer active.Add(-1) + for { + observed := maxActive.Load() + if current <= observed || maxActive.CompareAndSwap(observed, current) { + break + } + } + started <- struct{}{} + <-release + return channelTestSummary{Tested: 1, Succeeded: 1} + }, + func(processed, _ int) { + progress = append(progress, processed) + }, + ) + }() + + <-started + <-started + select { + case <-started: + t.Fatal("started more channel tests than the configured concurrency") + default: + } + close(release) + + summary := <-summaryResult + + assert.Equal(t, int32(2), maxActive.Load()) + assert.Equal(t, channelTestSummary{Tested: 4, Succeeded: 4}, summary) + assert.Equal(t, []int{0, 1, 2, 3, 4}, progress) +} + +func TestRunChannelTestWorkersStopsAfterCancellation(t *testing.T) { + originalInterval := common.RequestInterval + common.RequestInterval = 0 + t.Cleanup(func() { common.RequestInterval = originalInterval }) + + ctx, cancel := context.WithCancel(context.Background()) + channels := []*model.Channel{ + {Id: 1, Status: common.ChannelStatusEnabled}, + {Id: 2, Status: common.ChannelStatusEnabled}, + {Id: 3, Status: common.ChannelStatusEnabled}, + {Id: 4, Status: common.ChannelStatusEnabled}, + } + started := make(chan struct{}, len(channels)) + progress := make([]int, 0, 1) + summaryResult := make(chan channelTestSummary, 1) + + go func() { + summaryResult <- runChannelTestWorkers( + ctx, + channels, + 2, + func(ctx context.Context, _ *model.Channel) channelTestSummary { + started <- struct{}{} + <-ctx.Done() + return channelTestSummary{Tested: 1, Succeeded: 1} + }, + func(processed, _ int) { + progress = append(progress, processed) + }, + ) + }() + + <-started + <-started + cancel() + + summary := <-summaryResult + + select { + case <-started: + t.Fatal("started another channel test after cancellation") + default: + } + assert.Equal(t, channelTestSummary{Tested: 2, Succeeded: 2}, summary) + assert.Equal(t, []int{0}, progress) +} + func TestTestAllChannelsRejectsExistingActiveTask(t *testing.T) { db := setupModelListControllerTestDB(t) require.NoError(t, db.AutoMigrate(&model.SystemTask{}, &model.SystemTaskLock{})) diff --git a/controller/channel_upstream_update.go b/controller/channel_upstream_update.go index 71ab0e53fafe..6817a0086ac6 100644 --- a/controller/channel_upstream_update.go +++ b/controller/channel_upstream_update.go @@ -16,6 +16,7 @@ import ( "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/pkg/jsplugin" "github.com/QuantumNous/new-api/relay/channel/advancedcustom" "github.com/QuantumNous/new-api/relay/channel/gemini" "github.com/QuantumNous/new-api/relay/channel/ollama" @@ -304,6 +305,34 @@ func sanitizeFetchModelsError(err error, key string) error { return errors.New(message) } +func sanitizeAdvancedCustomRequestError(err error, key string, requestURL string) error { + err = sanitizeFetchModelsError(err, key) + if err == nil { + return nil + } + parsedURL, parseErr := url.Parse(requestURL) + if parseErr != nil { + return err + } + message := err.Error() + for _, value := range parsedURL.Query() { + for _, secret := range value { + if secret == "" { + continue + } + message = strings.ReplaceAll(message, secret, "[REDACTED]") + message = strings.ReplaceAll(message, url.QueryEscape(secret), "[REDACTED]") + message = strings.ReplaceAll(message, url.PathEscape(secret), "[REDACTED]") + } + } + if key != "" { + message = strings.ReplaceAll(message, key, "[REDACTED]") + message = strings.ReplaceAll(message, url.QueryEscape(key), "[REDACTED]") + message = strings.ReplaceAll(message, url.PathEscape(key), "[REDACTED]") + } + return errors.New(message) +} + func getFetchModelsResponseBody(method string, requestURL string, channel *model.Channel, headers http.Header) ([]byte, error) { request, err := http.NewRequest(method, requestURL, nil) if err != nil { @@ -333,7 +362,14 @@ func getFetchModelsResponseBody(method string, requestURL string, channel *model } func fetchChannelUpstreamModelIDs(channel *model.Channel) ([]string, error) { - baseURL := constant.ChannelBaseURLs[channel.Type] + if channel.Type == constant.ChannelTypeTaskPlugin { + plugin, ok := jsplugin.DefaultRegistry.Get(channel.GetSetting().TaskPluginKey) + if !ok { + return nil, fmt.Errorf("task plugin %q is not registered", channel.GetSetting().TaskPluginKey) + } + return normalizeModelNames(plugin.Meta.Models), nil + } + baseURL := constant.GetChannelBaseURL(channel.Type) if channel.GetBaseURL() != "" { baseURL = channel.GetBaseURL() } @@ -409,7 +445,7 @@ func fetchChannelUpstreamModelIDs(channel *model.Channel) ([]string, error) { body, err := getFetchModelsResponseBody(http.MethodGet, url, channel, headers) if err != nil { - return nil, sanitizeFetchModelsError(err, key) + return nil, sanitizeAdvancedCustomRequestError(err, key, url) } var result OpenAIModelsResponse diff --git a/controller/channel_upstream_update_test.go b/controller/channel_upstream_update_test.go index 0a524d704bdf..5cb4fac4b0bf 100644 --- a/controller/channel_upstream_update_test.go +++ b/controller/channel_upstream_update_test.go @@ -168,6 +168,15 @@ func TestFetchAdvancedCustomModelsRedactsQueryKeyFromTransportErrors(t *testing. Err: errors.New("connection refused"), }, secret) require.EqualError(t, direct, "connection refused") + + queryValue := "prefix-" + secret + queryError := sanitizeAdvancedCustomRequestError( + errors.New("dial "+queryValue+": connection refused"), + queryValue, + baseURL+"/v1/models?custom-token="+url.QueryEscape(queryValue), + ) + require.NotContains(t, queryError.Error(), queryValue) + require.EqualError(t, queryError, "dial [REDACTED]: connection refused") } func TestFetchOrdinaryOpenAIModelsKeepsExistingEmptyDataBehavior(t *testing.T) { diff --git a/controller/log.go b/controller/log.go index 470c759fc1a1..b5484992c00e 100644 --- a/controller/log.go +++ b/controller/log.go @@ -27,6 +27,11 @@ func GetAllLogs(c *gin.Context) { common.ApiError(c, err) return } + if c.GetInt("role") < common.RoleRootUser { + model.FormatAdminLogs(logs) + } else { + model.FormatRootLogs(logs) + } pageInfo.SetTotal(int(total)) pageInfo.SetItems(logs) common.ApiSuccess(c, pageInfo) diff --git a/controller/midjourney.go b/controller/midjourney.go index bf52314a7581..0494386c84b5 100644 --- a/controller/midjourney.go +++ b/controller/midjourney.go @@ -213,22 +213,7 @@ func runMidjourneyTaskUpdateOnce(ctx context.Context, report func(processed, tot if err != nil { logger.LogError(ctx, "UpdateMidjourneyTask task error: "+err.Error()) } else if won && shouldReturnQuota { - err = model.IncreaseUserQuota(task.UserId, task.Quota, false) - if err != nil { - logger.LogError(ctx, "fail to increase user quota: "+err.Error()) - } - model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{ - UserId: task.UserId, - LogType: model.LogTypeRefund, - Content: "", - ChannelId: task.ChannelId, - ModelName: service.CovertMjpActionToModelName(task.Action), - Quota: task.Quota, - Other: map[string]interface{}{ - "task_id": task.MjId, - "reason": "构图失败", - }, - }) + service.RefundMidjourneyQuota(ctx, task, "构图失败") } } } diff --git a/controller/misc.go b/controller/misc.go index 7343b12f10a3..9f7480ef958e 100644 --- a/controller/misc.go +++ b/controller/misc.go @@ -94,6 +94,8 @@ func GetStatus(c *gin.Context) { "password_register_enabled": common.PasswordRegisterEnabled, "default_use_auto_group": setting.DefaultUseAutoGroup, + "password_login_encryption_enabled": common.PasswordLoginEncryptionEnabled, + "usd_exchange_rate": operation_setting.USDExchangeRate, "price": operation_setting.Price, "stripe_unit_price": setting.StripeUnitPrice, @@ -174,42 +176,24 @@ func GetStatus(c *gin.Context) { func GetNotice(c *gin.Context) { common.OptionMapRWMutex.RLock() - defer common.OptionMapRWMutex.RUnlock() - c.JSON(http.StatusOK, gin.H{ - "success": true, - "message": "", - "data": common.OptionMap["Notice"], - }) - return + notice := common.OptionMap["Notice"] + common.OptionMapRWMutex.RUnlock() + serveRevalidatedJSON(c, notice) } func GetAbout(c *gin.Context) { common.OptionMapRWMutex.RLock() - defer common.OptionMapRWMutex.RUnlock() - c.JSON(http.StatusOK, gin.H{ - "success": true, - "message": "", - "data": common.OptionMap["About"], - }) - return + about := common.OptionMap["About"] + common.OptionMapRWMutex.RUnlock() + serveRevalidatedJSON(c, about) } func GetUserAgreement(c *gin.Context) { - c.JSON(http.StatusOK, gin.H{ - "success": true, - "message": "", - "data": system_setting.GetLegalSettings().UserAgreement, - }) - return + serveRevalidatedJSON(c, system_setting.GetLegalSettings().UserAgreement) } func GetPrivacyPolicy(c *gin.Context) { - c.JSON(http.StatusOK, gin.H{ - "success": true, - "message": "", - "data": system_setting.GetLegalSettings().PrivacyPolicy, - }) - return + serveRevalidatedJSON(c, system_setting.GetLegalSettings().PrivacyPolicy) } func GetMidjourney(c *gin.Context) { @@ -225,13 +209,9 @@ func GetMidjourney(c *gin.Context) { func GetHomePageContent(c *gin.Context) { common.OptionMapRWMutex.RLock() - defer common.OptionMapRWMutex.RUnlock() - c.JSON(http.StatusOK, gin.H{ - "success": true, - "message": "", - "data": common.OptionMap["HomePageContent"], - }) - return + homePageContent := common.OptionMap["HomePageContent"] + common.OptionMapRWMutex.RUnlock() + serveRevalidatedJSON(c, homePageContent) } func SendEmailVerification(c *gin.Context) { diff --git a/controller/model.go b/controller/model.go index b32eebd7daac..54d71f256885 100644 --- a/controller/model.go +++ b/controller/model.go @@ -9,6 +9,7 @@ import ( "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/pkg/jsplugin" "github.com/QuantumNous/new-api/relay" "github.com/QuantumNous/new-api/relay/channel/ai360" "github.com/QuantumNous/new-api/relay/channel/lingyiwanwu" @@ -20,6 +21,7 @@ import ( "github.com/QuantumNous/new-api/relaykit/types" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/gin-gonic/gin" "github.com/samber/lo" ) @@ -96,6 +98,9 @@ func init() { for i := 1; i <= constant.ChannelTypeDummy; i++ { apiType, success := common.ChannelType2APIType(i) if !success || apiType == constant.APITypeAIProxyLibrary { + if plugin, ok := jsplugin.DefaultRegistry.GetByChannelType(i); ok { + channelId2Models[i] = append([]string(nil), plugin.Meta.Models...) + } continue } meta := &relaycommon.RelayInfo{ChannelMeta: &relaycommon.ChannelMeta{ @@ -104,6 +109,11 @@ func init() { adaptor := relay.GetAdaptor(apiType) adaptor.Init(meta) channelId2Models[i] = adaptor.GetModelList() + if len(channelId2Models[i]) == 0 { + if plugin, ok := jsplugin.DefaultRegistry.GetByChannelType(i); ok { + channelId2Models[i] = append([]string(nil), plugin.Meta.Models...) + } + } } openAIModels = lo.UniqBy(openAIModels, func(m dto.OpenAIModels) string { return m.Id @@ -190,7 +200,7 @@ func getModelListGroups(c *gin.Context) (modelListGroups, error) { return modelListGroups{ userGroup: userGroup, tokenGroup: tokenGroup, - ownerGroups: service.GetUserAutoGroup(userGroup), + ownerGroups: service.GetRequestAutoGroups(c, userGroup), }, nil } @@ -228,32 +238,28 @@ func ListModels(c *gin.Context, modelType int) { } ownerGroups := groups.ownerGroups modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled) + var tokenModelLimit map[string]bool if modelLimitEnable { s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit) - var tokenModelLimit map[string]bool if ok { - tokenModelLimit = s.(map[string]bool) - } else { + tokenModelLimit, _ = s.(map[string]bool) + } + if tokenModelLimit == nil { tokenModelLimit = map[string]bool{} } - for allowModel, _ := range tokenModelLimit { - if !acceptUnsetRatioModel { - if !helper.HasModelBillingConfig(allowModel) { - continue - } + } + models := service.GetGroupsEnabledModels(ownerGroups) + for _, modelName := range models { + if modelLimitEnable { + matchingName := ratio_setting.RoutingMatchModelName(modelName) + if !tokenModelLimit[modelName] && !tokenModelLimit[matchingName] { + continue } - userModelNames = append(userModelNames, allowModel) } - } else { - models := service.GetGroupsEnabledModels(ownerGroups) - for _, modelName := range models { - if !acceptUnsetRatioModel { - if !helper.HasModelBillingConfig(modelName) { - continue - } - } - userModelNames = append(userModelNames, modelName) + if !acceptUnsetRatioModel && !helper.HasModelBillingConfig(modelName) { + continue } + userModelNames = append(userModelNames, modelName) } ownerByModel := map[string]string{} @@ -276,11 +282,17 @@ func ListModels(c *gin.Context, modelType int) { Type: "model", } } + firstID := "" + lastID := "" + if len(useranthropicModels) > 0 { + firstID = useranthropicModels[0].ID + lastID = useranthropicModels[len(useranthropicModels)-1].ID + } c.JSON(200, gin.H{ "data": useranthropicModels, - "first_id": useranthropicModels[0].ID, + "first_id": firstID, "has_more": false, - "last_id": useranthropicModels[len(useranthropicModels)-1].ID, + "last_id": lastID, }) case constant.ChannelTypeGemini: userGeminiModels := make([]dto.GeminiModel, len(userOpenAiModels)) @@ -311,9 +323,18 @@ func ChannelListModels(c *gin.Context) { } func DashboardListModels(c *gin.Context) { + modelsByChannel := make(map[int][]string, len(channelId2Models)) + for channelType, models := range channelId2Models { + modelsByChannel[channelType] = append([]string(nil), models...) + } + for channelType := 1; channelType <= constant.ChannelTypeDummy; channelType++ { + if plugin, ok := jsplugin.DefaultRegistry.GetByChannelType(channelType); ok { + modelsByChannel[channelType] = append([]string(nil), plugin.Meta.Models...) + } + } c.JSON(200, gin.H{ "success": true, - "data": channelId2Models, + "data": modelsByChannel, }) } diff --git a/controller/model_list_test.go b/controller/model_list_test.go index b1fa9b956ce3..812207b8fd44 100644 --- a/controller/model_list_test.go +++ b/controller/model_list_test.go @@ -402,7 +402,13 @@ func TestListModelsTokenLimitIncludesTieredBillingModel(t *testing.T) { "zz-token-tiered-visible-model": `tier("base", p * 1 + c * 2)`, "zz-token-tiered-empty-expr-model": "", }) - setupModelListControllerTestDB(t) + db := setupModelListControllerTestDB(t) + require.NoError(t, db.Create(&[]model.Ability{ + {Group: "default", Model: "zz-token-tiered-visible-model", ChannelId: 1, Enabled: true}, + {Group: "default", Model: "zz-token-tiered-empty-expr-model", ChannelId: 1, Enabled: true}, + {Group: "default", Model: "zz-token-tiered-missing-expr-model", ChannelId: 1, Enabled: true}, + {Group: "default", Model: "zz-token-unpriced-model", ChannelId: 1, Enabled: true}, + }).Error) recorder := httptest.NewRecorder() ctx, _ := gin.CreateTestContext(recorder) @@ -425,6 +431,68 @@ func TestListModelsTokenLimitIncludesTieredBillingModel(t *testing.T) { require.NotContains(t, ids, "zz-token-unpriced-model") } +func TestListModelsTokenLimitUsesResolvedCustomAutoGroups(t *testing.T) { + withSelfUseModeEnabled(t) + originalMax := setting.GetMaxTokenAutoGroups() + originalUsableGroups := setting.UserUsableGroups2JSONString() + originalRatios := ratio_setting.GroupRatio2JSONString() + require.NoError(t, setting.UpdateMaxTokenAutoGroups("5")) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`)) + t.Cleanup(func() { + require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMax))) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios)) + }) + + db := setupModelListControllerTestDB(t) + require.NoError(t, db.Create(&[]model.Ability{ + {Group: "vip", Model: "zz-vip-allowed", ChannelId: 1, Enabled: true}, + {Group: "vip", Model: "zz-vip-denied", ChannelId: 1, Enabled: true}, + {Group: "default", Model: "zz-default-outside-snapshot", ChannelId: 1, Enabled: true}, + }).Error) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil) + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") + common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "auto") + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"}) + common.SetContextKey(ctx, constant.ContextKeyTokenModelLimitEnabled, true) + common.SetContextKey(ctx, constant.ContextKeyTokenModelLimit, map[string]bool{ + "zz-vip-allowed": true, + "zz-default-outside-snapshot": true, + "zz-not-enabled": true, + }) + + ListModels(ctx, constant.ChannelTypeOpenAI) + ids := decodeListModelsResponse(t, recorder) + require.Equal(t, map[string]struct{}{"zz-vip-allowed": {}}, ids) + + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`)) + emptyRecorder := httptest.NewRecorder() + emptyCtx, _ := gin.CreateTestContext(emptyRecorder) + emptyCtx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil) + common.SetContextKey(emptyCtx, constant.ContextKeyUserGroup, "default") + common.SetContextKey(emptyCtx, constant.ContextKeyTokenGroup, "auto") + common.SetContextKey(emptyCtx, constant.ContextKeyTokenAutoGroups, []string{"vip"}) + common.SetContextKey(emptyCtx, constant.ContextKeyTokenModelLimitEnabled, true) + common.SetContextKey(emptyCtx, constant.ContextKeyTokenModelLimit, map[string]bool{"zz-vip-allowed": true}) + + require.NotPanics(t, func() { + ListModels(emptyCtx, constant.ChannelTypeAnthropic) + }) + var anthropicResponse struct { + Data []dto.AnthropicModel `json:"data"` + FirstID string `json:"first_id"` + LastID string `json:"last_id"` + } + require.NoError(t, common.Unmarshal(emptyRecorder.Body.Bytes(), &anthropicResponse)) + require.Empty(t, anthropicResponse.Data) + require.Empty(t, anthropicResponse.FirstID) + require.Empty(t, anthropicResponse.LastID) +} + func TestCheckUpdatePasswordRequiresCurrentPassword(t *testing.T) { db := setupModelListControllerTestDB(t) hashedPassword, err := common.Password2Hash("CurrentPassword123") diff --git a/controller/model_owned_by_test.go b/controller/model_owned_by_test.go index bc2ef32f135c..da9bb1a0ff0e 100644 --- a/controller/model_owned_by_test.go +++ b/controller/model_owned_by_test.go @@ -1,11 +1,14 @@ package controller import ( + "fmt" "net/http/httptest" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -83,3 +86,33 @@ func TestGetModelListGroupsUsesExplicitTokenGroup(t *testing.T) { require.Equal(t, "vip", groups.tokenGroup) require.Equal(t, []string{"vip"}, groups.ownerGroups) } + +func TestGetModelListGroupsUsesFilteredTokenAutoGroupsSnapshot(t *testing.T) { + originalMax := setting.GetMaxTokenAutoGroups() + originalUsableGroups := setting.UserUsableGroups2JSONString() + originalRatios := ratio_setting.GroupRatio2JSONString() + require.NoError(t, setting.UpdateMaxTokenAutoGroups("1")) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`)) + t.Cleanup(func() { + require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMax))) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios)) + }) + + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") + common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "auto") + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip", "default"}) + + groups, err := getModelListGroups(ctx) + require.NoError(t, err) + require.Equal(t, []string{"vip"}, groups.ownerGroups) + + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"}) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`)) + groups, err = getModelListGroups(ctx) + require.NoError(t, err) + require.Empty(t, groups.ownerGroups) +} diff --git a/controller/oauth.go b/controller/oauth.go index a477f5b1d035..4d9725c1cc7a 100644 --- a/controller/oauth.go +++ b/controller/oauth.go @@ -263,25 +263,20 @@ func handleOAuthBind(c *gin.Context, provider oauth.Provider, pendingFlow *model return } - user := model.User{Id: pendingFlow.UserId} - err = user.FillUserById() - if err != nil { - common.ApiError(c, err) - return - } + userId := pendingFlow.UserId // Handle binding based on provider type if genericProvider, ok := provider.(*oauth.GenericOAuthProvider); ok { // Custom provider: use user_oauth_bindings table - err = model.UpdateUserOAuthBinding(user.Id, genericProvider.GetProviderId(), oauthUser.ProviderUserID) + err = model.UpdateUserOAuthBinding(userId, genericProvider.GetProviderId(), oauthUser.ProviderUserID) if err != nil { common.ApiError(c, err) return } } else { - // Built-in provider: update user record directly - provider.SetProviderUserID(&user, oauthUser.ProviderUserID) - err = user.Update(false) + // Built-in provider: 只更新绑定列。完整快照的 user.Update 会把读取时刻的 + // role/status/group 一并写回,覆盖并发发生的封禁、降权或分组变更。 + err = model.UpdateUserBindColumn(userId, provider.ProviderUserIDColumn(), oauthUser.ProviderUserID) if err != nil { common.ApiError(c, err) return diff --git a/controller/option.go b/controller/option.go index 940bb3069023..20ce18b2f841 100644 --- a/controller/option.go +++ b/controller/option.go @@ -3,13 +3,18 @@ package controller import ( "fmt" "net/http" + "slices" + "sort" "strconv" "strings" "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" "github.com/QuantumNous/new-api/setting/console_setting" "github.com/QuantumNous/new-api/setting/model_setting" "github.com/QuantumNous/new-api/setting/operation_setting" @@ -81,7 +86,7 @@ func GetOptions(c *gin.Context) { optionValues := make(map[string]string) common.OptionMapRWMutex.Lock() for k, v := range common.OptionMap { - if k == "theme.frontend" { + if k == "theme.frontend" || k == "billing_setting.billing_mode" || k == "billing_setting.billing_expr" { continue } value := common.Interface2String(v) @@ -97,14 +102,24 @@ func GetOptions(c *gin.Context) { Key: k, Value: value, }) - for _, optionKey := range completionRatioMetaOptionKeys { - if optionKey == k { - optionValues[k] = value - break - } + if slices.Contains(completionRatioMetaOptionKeys, k) { + optionValues[k] = value } } common.OptionMapRWMutex.Unlock() + // Display the same effective expressions used by pricing and settlement, + // including built-in defaults absent from persisted administrator options. + for key, values := range map[string]map[string]string{ + "billing_setting.billing_mode": billing_setting.GetBillingModeCopy(), + "billing_setting.billing_expr": billing_setting.GetBillingExprCopy(), + } { + encoded, err := common.Marshal(values) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"success": false, "message": err.Error()}) + return + } + options = append(options, &model.Option{Key: key, Value: string(encoded)}) + } options = append(options, &model.Option{ Key: "CompletionRatioMeta", Value: buildCompletionRatioMetaValue(optionValues), @@ -153,6 +168,12 @@ func UpdateOption(c *gin.Context) { return } } + if option.Key == "TaskPublicAddress" && option.Value.(string) != "" { + if err := service.ValidateTaskArtifactBaseURL(option.Value.(string)); err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + } switch option.Key { case "GitHubOAuthEnabled": if option.Value == "true" && common.GitHubClientId == "" { @@ -326,6 +347,36 @@ func UpdateOption(c *gin.Context) { }) return } + case "billing_setting.billing_expr": + expressions := make(map[string]string) + if err = common.UnmarshalJsonStr(option.Value.(string), &expressions); err != nil { + common.ApiErrorMsg(c, "计费表达式配置必须是模型到表达式的 JSON 对象: "+err.Error()) + return + } + models := make([]string, 0, len(expressions)) + for modelName := range expressions { + models = append(models, modelName) + } + sort.Strings(models) + generation := jsplugin.DefaultRegistry.Generation() + 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) + } + } else { + err = billing_setting.SmokeTestExpr(expression) + } + if err != nil { + common.ApiErrorMsg(c, fmt.Sprintf("模型 %s 的计费表达式无效: %v", modelName, err)) + return + } + } case "console_setting.api_info": err = console_setting.ValidateConsoleSettings(option.Value.(string), "ApiInfo") if err != nil { diff --git a/controller/plugin_endpoint_test.go b/controller/plugin_endpoint_test.go new file mode 100644 index 000000000000..5a0aa69f7341 --- /dev/null +++ b/controller/plugin_endpoint_test.go @@ -0,0 +1,47 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/pkg/jsplugin" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" +) + +func TestRelayTaskPluginEndpointPreservesUnclaimedFallback(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + fallbackCalls := 0 + + RelayTaskPluginEndpoint(c, func(c *gin.Context) { + fallbackCalls++ + c.Status(http.StatusNoContent) + c.Writer.WriteHeaderNow() + }) + + assert.Equal(t, 1, fallbackCalls) + assert.Equal(t, http.StatusNoContent, recorder.Code) +} + +func TestRelayTaskPluginEndpointNeverEntersOrdinaryRelayWhenClaimed(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{ + Generation: &jsplugin.RoutingGeneration{}, + Plugin: &jsplugin.LoadedPlugin{}, + Protocol: "openai_responses", + Operation: jsplugin.HostProtocolOperation{Name: "create"}, + }) + fallbackCalls := 0 + + RelayTaskPluginEndpoint(c, func(c *gin.Context) { + fallbackCalls++ + c.Status(http.StatusNoContent) + c.Writer.WriteHeaderNow() + }) + + assert.Zero(t, fallbackCalls) + assert.NotEqual(t, http.StatusNoContent, recorder.Code) +} diff --git a/controller/plugin_native_e2e_test.go b/controller/plugin_native_e2e_test.go new file mode 100644 index 000000000000..f0eb4eef0416 --- /dev/null +++ b/controller/plugin_native_e2e_test.go @@ -0,0 +1,244 @@ +package controller + +import ( + "bytes" + "context" + "io" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/middleware" + "github.com/QuantumNous/new-api/model" + pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin" + "github.com/QuantumNous/new-api/relay" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +type nativeRouteBilling struct { + events []string + preConsumed int + userID int + settled bool +} + +func (b *nativeRouteBilling) Settle(int) error { + b.events = append(b.events, "settle") + b.settled = true + return nil +} + +func (b *nativeRouteBilling) Refund(*gin.Context) { + b.events = append(b.events, "refund") + if !b.settled && b.preConsumed > 0 { + _ = model.IncreaseUserQuota(b.userID, b.preConsumed, true) + b.preConsumed = 0 + } +} + +func (b *nativeRouteBilling) NeedsRefund() bool { + return !b.settled && b.preConsumed > 0 +} + +func (b *nativeRouteBilling) GetPreConsumedQuota() int { + return b.preConsumed +} + +func (b *nativeRouteBilling) Reserve(quota int) error { + b.events = append(b.events, "reserve") + if err := model.DecreaseUserQuota(b.userID, quota, true); err != nil { + return err + } + b.preConsumed = quota + return nil +} + +func TestKlingNativeRouteSubmitPollSettleAndQuery(t *testing.T) { + gin.SetMode(gin.TestMode) + service.InitHttpClient() + + previousDB := model.DB + previousLogDB := model.LOG_DB + previousMemoryCache := common.MemoryCacheEnabled + previousBatchUpdate := common.BatchUpdateEnabled + previousLogConsume := common.LogConsumeEnabled + previousRedisEnabled := common.RedisEnabled + database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, database.AutoMigrate(&model.User{}, &model.Channel{}, &model.Task{}, &model.Log{})) + model.DB = database + model.LOG_DB = database + common.MemoryCacheEnabled = false + common.BatchUpdateEnabled = false + common.LogConsumeEnabled = false + common.RedisEnabled = false + previousModelRatios := ratio_setting.ModelRatio2JSONString() + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"kling-v1":1}`)) + t.Cleanup(func() { + model.DB = previousDB + model.LOG_DB = previousLogDB + common.MemoryCacheEnabled = previousMemoryCache + common.BatchUpdateEnabled = previousBatchUpdate + common.LogConsumeEnabled = previousLogConsume + common.RedisEnabled = previousRedisEnabled + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(previousModelRatios)) + }) + require.NoError(t, database.Create(&model.User{ + Id: 7, + Username: "native-route-user", + Group: "default", + Quota: 1_000_000, + }).Error) + + var submitCalls atomic.Int32 + var queryCalls atomic.Int32 + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch { + case r.Method == http.MethodPost && r.URL.Path == "/kling/v1/videos/text2video": + submitCalls.Add(1) + body, readErr := io.ReadAll(r.Body) + if !assert.NoError(t, readErr) { + http.Error(w, "read request", http.StatusInternalServerError) + return + } + assert.Contains(t, string(body), `"model_name":"kling-v1"`) + _, _ = io.WriteString(w, `{"code":0,"message":"","data":{"task_id":"kling-private-1","task_status":"submitted"}}`) + case r.Method == http.MethodGet && r.URL.Path == "/kling/v1/videos/text2video/kling-private-1": + queryCalls.Add(1) + _, _ = io.WriteString(w, `{"code":0,"message":"","data":{"task_id":"kling-private-1","task_status":"succeed","task_status_msg":"","task_result":{"videos":[{"id":"video-private","url":"https://cdn.example/video.mp4","duration":"5"}]},"final_unit_deduction":"1"}}`) + default: + http.NotFound(w, r) + } + })) + defer upstream.Close() + + channel := model.Channel{ + Type: constant.ChannelTypeKling, + Name: "kling-native-e2e", + Key: "sk-test", + BaseURL: &upstream.URL, + Status: common.ChannelStatusEnabled, + Models: "kling-v1", + Group: "default", + } + require.NoError(t, database.Create(&channel).Error) + + generation := pluginruntime.DefaultRegistry.Generation() + require.NotNil(t, generation) + submitBinding, found := generation.LookupDeclaredRoute(http.MethodPost, "/kling/v1/videos/text2video") + require.True(t, found) + require.Equal(t, "kling", submitBinding.Plugin.Meta.Key) + + submitRecorder := httptest.NewRecorder() + submitContext, _ := gin.CreateTestContext(submitRecorder) + submitContext.Request = httptest.NewRequest( + http.MethodPost, + "/kling/v1/videos/text2video", + bytes.NewBufferString(`{"model_name":"kling-v1","prompt":"a lighthouse"}`), + ) + submitContext.Request.Header.Set("Content-Type", "application/json") + submitContext.Set(pluginruntime.ContextKeyPinnedRoute, pluginruntime.PinnedRoute{ + Generation: generation, + Plugin: submitBinding.Plugin, + Route: submitBinding.Route, + }) + common.SetContextKey(submitContext, constant.ContextKeyUserId, 7) + common.SetContextKey(submitContext, constant.ContextKeyUserGroup, "default") + common.SetContextKey(submitContext, constant.ContextKeyUsingGroup, "default") + common.SetContextKey(submitContext, constant.ContextKeyTokenGroup, "default") + common.SetContextKey(submitContext, constant.ContextKeyUserQuota, 1_000_000) + + middleware.PrepareTaskPluginRoute()(submitContext) + require.False(t, submitContext.IsAborted(), submitRecorder.Body.String()) + require.Equal(t, "kling-v1", submitContext.GetString("resolved_task_model")) + require.Equal(t, "text_to_video", submitContext.GetString("task_action")) + require.Nil(t, middleware.SetupContextForSelectedChannel(submitContext, &channel, "kling-v1")) + + billing := &nativeRouteBilling{userID: 7} + relayInfo := &relaycommon.RelayInfo{ + UserId: 7, + UserGroup: "default", + UsingGroup: "default", + UserQuota: 1_000_000, + TokenGroup: "default", + OriginModelName: "kling-v1", + Billing: billing, + TaskRelayInfo: &relaycommon.TaskRelayInfo{ + Action: submitContext.GetString("task_action"), + PublicTaskID: "task_kling_public", + LockedChannel: &channel, + }, + } + + outcome, taskErr := executeTaskSubmissionWith(submitContext, relayInfo, relay.RelayTaskSubmit) + require.Nil(t, taskErr) + require.NotNil(t, outcome) + require.Equal(t, []string{"reserve", "settle"}, billing.events) + require.False(t, submitContext.Writer.Written()) + + presentTaskSubmission(submitContext, outcome) + require.Equal(t, http.StatusOK, submitRecorder.Code) + assert.Contains(t, submitRecorder.Body.String(), `"task_id":"task_kling_public"`) + assert.NotContains(t, submitRecorder.Body.String(), "kling-private-1") + assert.Equal(t, int32(1), submitCalls.Load()) + + var persisted model.Task + require.NoError(t, database.Where("task_id = ?", "task_kling_public").First(&persisted).Error) + assert.Equal(t, constant.TaskPlatform("kling"), persisted.Platform) + assert.Equal(t, "kling-private-1", persisted.PrivateData.UpstreamTaskID) + assert.Equal(t, model.TaskStatus(model.TaskStatusNotStart), persisted.Status) + + previousAdaptorFactory := service.GetTaskAdaptorFunc + service.GetTaskAdaptorFunc = func(platform constant.TaskPlatform) service.TaskPollingAdaptor { + return relay.GetTaskAdaptor(platform) + } + t.Cleanup(func() { service.GetTaskAdaptorFunc = previousAdaptorFactory }) + service.DispatchPlatformUpdate( + context.Background(), + persisted.Platform, + map[int][]string{channel.Id: {"kling-private-1"}}, + map[string]*model.Task{"kling-private-1": &persisted}, + ) + + require.NoError(t, database.Where("task_id = ?", "task_kling_public").First(&persisted).Error) + assert.Equal(t, model.TaskStatus(model.TaskStatusSuccess), persisted.Status) + assert.Equal(t, "100%", persisted.Progress) + assert.Equal(t, 1, persisted.Quota) + assert.Equal(t, int32(1), queryCalls.Load()) + var settledUser model.User + require.NoError(t, database.First(&settledUser, 7).Error) + assert.Equal(t, 999_999, settledUser.Quota) + + queryBinding, found := generation.LookupDeclaredRoute(http.MethodGet, "/kling/v1/videos/text2video/:task_id") + require.True(t, found) + queryRecorder := httptest.NewRecorder() + queryContext, _ := gin.CreateTestContext(queryRecorder) + queryContext.Request = httptest.NewRequest(http.MethodGet, "/kling/v1/videos/text2video/task_kling_public", nil) + queryContext.Params = gin.Params{{Key: "task_id", Value: "task_kling_public"}} + queryContext.Set(pluginruntime.ContextKeyPinnedRoute, pluginruntime.PinnedRoute{ + Generation: generation, + Plugin: queryBinding.Plugin, + Route: queryBinding.Route, + }) + common.SetContextKey(queryContext, constant.ContextKeyUserId, 7) + + middleware.PrepareTaskPluginRoute()(queryContext) + + require.True(t, queryContext.IsAborted()) + require.Equal(t, http.StatusOK, queryRecorder.Code) + assert.Contains(t, queryRecorder.Body.String(), `"task_id":"task_kling_public"`) + assert.Contains(t, queryRecorder.Body.String(), `"task_status":"succeed"`) + assert.NotContains(t, queryRecorder.Body.String(), "kling-private-1") + assert.NotContains(t, queryRecorder.Body.String(), upstream.URL) +} diff --git a/controller/plugin_protocol.go b/controller/plugin_protocol.go new file mode 100644 index 000000000000..7b9248aeffdb --- /dev/null +++ b/controller/plugin_protocol.go @@ -0,0 +1,1276 @@ +package controller + +import ( + "context" + "errors" + "fmt" + "hash/fnv" + "net/http" + "strconv" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin" + "github.com/QuantumNous/new-api/relay" + taskjsplugin "github.com/QuantumNous/new-api/relay/channel/task/jsplugin" + relaycommon "github.com/QuantumNous/new-api/relay/common" + relayconstant "github.com/QuantumNous/new-api/relay/constant" + "github.com/QuantumNous/new-api/relay/helper" + "github.com/QuantumNous/new-api/relaykit/types" + "github.com/QuantumNous/new-api/service" + "github.com/gin-gonic/gin" +) + +type pluginProtocolBridgeDeps struct { + submit func(*gin.Context, *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) + loadTask func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) + now func() time.Time + admissions *pluginProtocolObservationLimiter + protocolLimits relay.PluginProtocolLimits + artifactContentURL func(taskID, artifactKey string) (string, error) + submissionTimeout time.Duration + observationTimeout time.Duration + loadTimeout time.Duration + tickInterval time.Duration + tickJitter time.Duration + heartbeatInterval time.Duration + admissionTimeout time.Duration + getByTaskId func(int, string) (*model.Task, bool, error) + resolvePlugin func(constant.TaskPlatform) (*pluginruntime.LoadedPlugin, *pluginruntime.RoutingGeneration, bool) +} + +func defaultPluginProtocolBridgeDeps() pluginProtocolBridgeDeps { + timeout := time.Duration(constant.TaskPluginProtocolTimeoutSeconds) * time.Second + if timeout <= 0 { + timeout = 10 * time.Minute + } + tick := time.Duration(constant.TaskPluginProtocolTickMilliseconds) * time.Millisecond + if tick <= 0 { + tick = 2 * time.Second + } + jitter := max(time.Duration(constant.TaskPluginProtocolTickJitterMilliseconds)*time.Millisecond, 0) + heartbeat := time.Duration(constant.TaskPluginProtocolHeartbeatSeconds) * time.Second + if heartbeat <= 0 { + heartbeat = 15 * time.Second + } + loadTimeout := 5 * time.Second + if halfHeartbeat := heartbeat / 2; halfHeartbeat > 0 && halfHeartbeat < loadTimeout { + loadTimeout = halfHeartbeat + } + return pluginProtocolBridgeDeps{ + submit: executeTaskSubmission, + loadTask: model.GetTaskForProtocolObservation, + now: time.Now, + admissions: pluginProtocolObservationAdmissions, + protocolLimits: relay.DefaultPluginProtocolLimits(), + artifactContentURL: service.BuildTaskArtifactContentURL, + submissionTimeout: timeout, + observationTimeout: timeout, + loadTimeout: loadTimeout, + tickInterval: tick, + tickJitter: jitter, + heartbeatInterval: heartbeat, + admissionTimeout: pluginruntime.DefaultCallTimeout, + getByTaskId: model.GetByTaskId, + resolvePlugin: resolveTaskPluginForProtocolRetrieve, + } +} + +func (d pluginProtocolBridgeDeps) withDefaults() pluginProtocolBridgeDeps { + defaults := defaultPluginProtocolBridgeDeps() + if d.submit == nil { + d.submit = defaults.submit + } + if d.loadTask == nil { + d.loadTask = defaults.loadTask + } + if d.now == nil { + d.now = defaults.now + } + if d.admissions == nil { + d.admissions = defaults.admissions + } + if d.artifactContentURL == nil { + d.artifactContentURL = defaults.artifactContentURL + } + if d.submissionTimeout <= 0 { + d.submissionTimeout = defaults.submissionTimeout + } + if d.observationTimeout <= 0 { + d.observationTimeout = defaults.observationTimeout + } + if d.loadTimeout <= 0 { + d.loadTimeout = defaults.loadTimeout + } + if d.tickInterval <= 0 { + d.tickInterval = defaults.tickInterval + } + if d.tickJitter < 0 { + d.tickJitter = 0 + } + if d.heartbeatInterval <= 0 { + d.heartbeatInterval = defaults.heartbeatInterval + } + if halfHeartbeat := d.heartbeatInterval / 2; halfHeartbeat > 0 && d.loadTimeout > halfHeartbeat { + d.loadTimeout = halfHeartbeat + } + if d.admissionTimeout <= 0 { + d.admissionTimeout = defaults.admissionTimeout + } + if d.getByTaskId == nil { + d.getByTaskId = defaults.getByTaskId + } + if d.resolvePlugin == nil { + d.resolvePlugin = defaults.resolvePlugin + } + return d +} + +func resolveTaskPluginForProtocolRetrieve(platform constant.TaskPlatform) (*pluginruntime.LoadedPlugin, *pluginruntime.RoutingGeneration, bool) { + generation := pluginruntime.DefaultRegistry.Generation() + plugin, ok := relay.ResolveTaskPluginForPlatform(generation, platform) + return plugin, generation, ok +} + +func serveTaskPluginProtocol( + c *gin.Context, + pinned pluginruntime.PinnedEndpoint, + deps pluginProtocolBridgeDeps, +) { + deps = deps.withDefaults() + generation := uint64(0) + if pinned.Generation != nil { + generation = pinned.Generation.Number + } + pluginKey := "" + if pinned.Plugin != nil { + pluginKey = pinned.Plugin.Meta.Key + } + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=bridge_start generation=%d plugin=%q protocol=%q model=%q", + generation, + pluginKey, + pinned.Protocol, + c.GetString("resolved_task_model"), + ) + if !pluginruntime.SupportsHostProtocol(pinned.Protocol) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=bridge_rejected generation=%d plugin=%q reason=unsupported_protocol", generation, pluginKey) + respondPluginProtocolError(c, http.StatusNotImplemented, "task_protocol_not_available", "Task protocol bridge is not available") + return + } + requestValue, exists := c.Get(pluginruntime.ContextKeyProtocolRequest) + protocolRequest, ok := requestValue.(pluginruntime.ProtocolRequestContext) + if !exists || !ok || protocolRequest.Protocol != pinned.Protocol { + logger.LogDebug(c, "task_plugin subsystem=protocol event=bridge_rejected generation=%d plugin=%q reason=invalid_protocol_context", generation, pluginKey) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + if definition, known := pluginruntime.HostProtocol(pinned.Protocol); known && len(definition.DefinedModes()) > 0 && pinned.Plugin != nil { + background := false + if body, ok := protocolRequest.Body.(map[string]any); ok && body["kind"] == string(pluginruntime.BodyJSON) { + if requestBody, ok := body["value"].(map[string]any); ok { + background, _ = requestBody["background"].(bool) + } + } + missing := false + if protocolRequest.Stream && !pinned.Plugin.Meta.ProtocolSupports(pinned.Protocol, "stream") { + missing = true + } + if background && !pinned.Plugin.Meta.ProtocolSupports(pinned.Protocol, "background") { + missing = true + } + if !protocolRequest.Stream && !background && !pinned.Plugin.Meta.ProtocolSupports(pinned.Protocol, "sync") { + missing = true + } + if missing { + logger.LogError(c, "pinned task plugin does not support the requested protocol form") + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + } + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=request_ready generation=%d plugin=%q protocol=%q stream=%t", + generation, + pluginKey, + protocolRequest.Protocol, + protocolRequest.Stream, + ) + + release, admissionErr := deps.admissions.acquire( + pinned.Plugin.Meta.Key, + common.GetContextKeyInt(c, constant.ContextKeyUserId), + common.GetContextKeyInt(c, constant.ContextKeyTokenId), + ) + if admissionErr != nil { + if errors.Is(admissionErr, errPluginProtocolObservationLimitExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=admission_rejected generation=%d plugin=%q reason=observation_limit", generation, pluginKey) + respondPluginProtocolError(c, http.StatusTooManyRequests, "rate_limit_exceeded", "Too many active task observations") + return + } + logger.LogDebug(c, "task_plugin subsystem=protocol event=admission_rejected generation=%d plugin=%q reason=invalid_identity", generation, pluginKey) + respondPluginProtocolError(c, http.StatusUnauthorized, "authentication_error", "Authentication failed") + return + } + defer release() + logger.LogDebug(c, "task_plugin subsystem=protocol event=admission_acquired generation=%d plugin=%q", generation, pluginKey) + + clientRequest := c.Request + var relayInfo *relaycommon.RelayInfo + var outcome *taskSubmissionOutcome + var taskErr *dto.TaskError + var relayInfoErr error + submissionStage := "relay_info" + // A Responses client only observes an asynchronous task. Once admitted, + // disconnecting that observer must not cancel submission, persistence, or + // billing settlement; the submission keeps its own bounded lifetime. + func() { + submissionContext, cancelSubmission := context.WithTimeout( + context.WithoutCancel(clientRequest.Context()), + deps.submissionTimeout, + ) + c.Request = clientRequest.Clone(submissionContext) + defer func() { + c.Request = clientRequest + cancelSubmission() + }() + + relayInfo, relayInfoErr = relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil) + if relayInfoErr != nil { + return + } + relayInfo.RelayMode = relayconstant.RelayModeVideoSubmit + relayInfo.IsStream = false + relayInfo.OriginModelName = c.GetString("resolved_task_model") + if action := c.GetString("task_action"); action != "" { + relayInfo.Action = action + } + submissionStage = "origin_task" + if taskErr = relay.ResolveOriginTask(c, relayInfo); taskErr != nil { + return + } + if taskErr = relay.ApplyOriginTaskAffinity(c, relayInfo); taskErr != nil { + return + } + + submissionStage = "submission" + logger.LogDebug(c, "task_plugin subsystem=protocol event=submission_start generation=%d plugin=%q protocol=%q stream=%t", generation, pluginKey, protocolRequest.Protocol, protocolRequest.Stream) + outcome, taskErr = deps.submit(c, relayInfo) + }() + + if clientRequest.Context().Err() != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q stage=%s", generation, pluginKey, submissionStage) + return + } + if relayInfoErr != nil { + err := relayInfoErr + logger.LogError(c, "build task protocol relay info failed: "+err.Error()) + logger.LogDebug(c, "task_plugin subsystem=protocol event=bridge_failed generation=%d plugin=%q stage=relay_info reason=invalid_context", generation, pluginKey) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + if submissionStage == "origin_task" && taskErr != nil { + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=bridge_failed generation=%d plugin=%q stage=origin_task code=%q status=%d", + generation, + pluginKey, + taskErr.Code, + taskErr.StatusCode, + ) + respondPluginProtocolSubmissionError(c, taskErr) + return + } + + if taskErr != nil { + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=submission_failed generation=%d plugin=%q code=%q status=%d local=%t", + generation, + pluginKey, + taskErr.Code, + taskErr.StatusCode, + taskErr.LocalError, + ) + respondPluginProtocolSubmissionError(c, taskErr) + return + } + if outcome == nil || outcome.Task == nil || outcome.RelayInfo == nil || + outcome.Task.UserId != relayInfo.UserId || + outcome.Task.Platform != constant.TaskPlatform(pinned.Plugin.Meta.Key) { + logger.LogError(c, "task protocol submission returned an invalid durable outcome") + logger.LogDebug(c, "task_plugin subsystem=protocol event=submission_failed generation=%d plugin=%q reason=invalid_durable_outcome", generation, pluginKey) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=submission_durable generation=%d plugin=%q public_task_id=%q status=%q stream=%t", + generation, + pluginKey, + outcome.Task.TaskID, + taskPluginDebugStatus(string(outcome.Task.Status)), + protocolRequest.Stream, + ) + + createdAt := outcome.Task.CreatedAt + if createdAt == 0 { + createdAt = outcome.Task.SubmitTime + } + if createdAt == 0 { + createdAt = deps.now().Unix() + } + machine := relay.NewPluginResponsesMachine( + outcome.Task.TaskID, + outcome.RelayInfo.OriginModelName, + createdAt, + deps.protocolLimits, + ) + background := false + if body, ok := protocolRequest.Body.(map[string]any); ok && body["kind"] == string(pluginruntime.BodyJSON) { + if requestBody, ok := body["value"].(map[string]any); ok { + background, _ = requestBody["background"].(bool) + } + } + if background { + outcome.Task.PrivateData.ResponsesBackground = true + if outcome.Task.ID != 0 { + if err := model.DB.Model(outcome.Task).Update("private_data", outcome.Task.PrivateData).Error; err != nil { + logger.LogError(c, "persist task background flag failed: "+err.Error()) + } + } + machine.SetBackground(true) + if !protocolRequest.Stream { + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=background_return generation=%d plugin=%q public_task_id=%q status=%q", + generation, + pluginKey, + outcome.Task.TaskID, + taskPluginDebugStatus(string(outcome.Task.Status)), + ) + c.JSON(http.StatusOK, machine.PendingResponse(string(outcome.Task.Status))) + return + } + logger.LogDebug(c, "task_plugin subsystem=protocol event=background_stream generation=%d plugin=%q public_task_id=%q", generation, pluginKey, outcome.Task.TaskID) + } + if protocolRequest.Stream { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_enter generation=%d plugin=%q mode=stream public_task_id=%q", generation, pluginKey, outcome.Task.TaskID) + streamTaskPluginProtocol(c, pinned, protocolRequest, outcome.Task.TaskID, machine, deps) + return + } + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_enter generation=%d plugin=%q mode=nonstream public_task_id=%q", generation, pluginKey, outcome.Task.TaskID) + waitTaskPluginProtocol(c, pinned, protocolRequest, outcome.Task.TaskID, machine, deps) +} + +func streamTaskPluginProtocol( + c *gin.Context, + pinned pluginruntime.PinnedEndpoint, + protocolRequest pluginruntime.ProtocolRequestContext, + taskID string, + machine *relay.PluginResponsesMachine, + deps pluginProtocolBridgeDeps, +) { + generation := pinned.Generation.Number + pluginKey := pinned.Plugin.Meta.Key + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_start generation=%d plugin=%q mode=stream public_task_id=%q timeout_ms=%d tick_ms=%d heartbeat_ms=%d", + generation, + pluginKey, + taskID, + deps.observationTimeout.Milliseconds(), + deps.tickInterval.Milliseconds(), + deps.heartbeatInterval.Milliseconds(), + ) + created, err := machine.CreatedEvent() + if err != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_failed generation=%d plugin=%q mode=stream stage=created_event reason=state_machine_error", generation, pluginKey) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + helper.SetEventStreamHeaders(c) + if err = writeTaskPluginProtocolEvent(c, created); err != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_write_failed generation=%d plugin=%q mode=stream stage=created_event", generation, pluginKey) + return + } + + observationContext, cancelObservation := context.WithTimeout(c.Request.Context(), deps.observationTimeout) + defer cancelObservation() + heartbeatTicker := time.NewTicker(deps.heartbeatInterval) + defer heartbeatTicker.Stop() + + var previous relay.ProtocolState + tickNumber := uint64(0) + lastStatus := "" + for { + loadStarted := deps.now() + loadContext, cancelLoad := context.WithTimeout(observationContext, deps.loadTimeout) + task, exists, loadErr := deps.loadTask( + loadContext, + common.GetContextKeyInt(c, constant.ContextKeyUserId), + constant.TaskPlatform(pinned.Plugin.Meta.Key), + taskID, + ) + loadContextErr := loadContext.Err() + cancelLoad() + loadElapsed := deps.now().Sub(loadStarted) + if errors.Is(loadContextErr, context.DeadlineExceeded) && + observationContext.Err() == nil && + c.Request.Context().Err() == nil { + logger.LogWarn(c, fmt.Sprintf( + "task protocol database observation overloaded; plugin=%s task=%s", + pinned.Plugin.Meta.Key, + taskID, + )) + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_tick generation=%d plugin=%q mode=stream tick=%d load_ms=%d overloaded=true", + generation, + pluginKey, + tickNumber, + loadElapsed.Milliseconds(), + ) + delay := pluginProtocolTickDelay(taskID, tickNumber, deps.tickInterval, deps.tickJitter) + deps.tickInterval + tickNumber++ + if !waitForTaskPluginProtocolTick(c, observationContext, heartbeatTicker, delay) { + if errors.Is(observationContext.Err(), context.DeadlineExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_timeout generation=%d plugin=%q mode=stream last_status=%q", generation, pluginKey, taskPluginDebugStatus(lastStatus)) + writeTaskPluginProtocolTimeout(c, machine, lastStatus) + } else if c.Request.Context().Err() != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q mode=stream stage=backoff_wait", generation, pluginKey) + } else { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_write_failed generation=%d plugin=%q mode=stream stage=heartbeat", generation, pluginKey) + } + return + } + continue + } + if loadErr != nil || !exists || task == nil { + if errors.Is(observationContext.Err(), context.DeadlineExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_timeout generation=%d plugin=%q mode=stream last_status=%q", generation, pluginKey, taskPluginDebugStatus(lastStatus)) + writeTaskPluginProtocolTimeout(c, machine, lastStatus) + return + } + if loadErr != nil && !errors.Is(loadErr, context.Canceled) { + logger.LogError(c, "task protocol database observation failed") + } + if c.Request.Context().Err() == nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_failed generation=%d plugin=%q mode=stream stage=load reason=task_unavailable", generation, pluginKey) + writeTaskPluginProtocolFailure(c, machine, lastStatus) + } else { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q mode=stream stage=load", generation, pluginKey) + } + return + } + previousStatus := lastStatus + lastStatus = string(task.Status) + if lastStatus != previousStatus { + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=status_transition generation=%d plugin=%q mode=stream tick=%d previous=%q status=%q load_ms=%d", + generation, + pluginKey, + tickNumber, + taskPluginDebugStatus(previousStatus), + taskPluginDebugStatus(lastStatus), + loadElapsed.Milliseconds(), + ) + } + view, viewErr := service.BuildTaskPluginView(task) + if viewErr != nil { + logger.LogError(c, "build task protocol view failed: "+viewErr.Error()) + writeTaskPluginProtocolFailure(c, machine, lastStatus) + return + } + viewValue, viewErr := taskPluginProtocolJSONValue(view) + if viewErr != nil { + logger.LogError(c, "encode task protocol view failed: "+viewErr.Error()) + writeTaskPluginProtocolFailure(c, machine, lastStatus) + return + } + hookStarted := deps.now() + rendererContext, contextErr := taskPluginProtocolRendererContext(protocolRequest, pinned, task, deps.artifactContentURL) + if contextErr != nil { + logger.LogError(c, "build task protocol renderer context failed") + writeTaskPluginProtocolFailure(c, machine, lastStatus) + return + } + args := []any{rendererContext, viewValue} + if previous.Present { + previousValue, stateErr := previous.PluginValue() + if stateErr != nil { + logger.LogError(c, "decode task protocol state failed: "+stateErr.Error()) + writeTaskPluginProtocolFailure(c, machine, lastStatus) + return + } + args = append(args, previousValue) + } + value, callErr := pinned.Plugin.Engine.CallPathWithAdmissionTimeout(observationContext, deps.admissionTimeout, "protocols", []string{pinned.Protocol, "renderEvents"}, args...) + hookElapsed := deps.now().Sub(hookStarted) + overloaded := false + if callErr != nil { + if errors.Is(observationContext.Err(), context.DeadlineExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_timeout generation=%d plugin=%q mode=stream stage=render_events last_status=%q", generation, pluginKey, taskPluginDebugStatus(lastStatus)) + writeTaskPluginProtocolTimeout(c, machine, lastStatus) + return + } + if c.Request.Context().Err() != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q mode=stream stage=render_events", generation, pluginKey) + return + } + if errors.Is(callErr, pluginruntime.ErrCallAdmissionTimeout) { + overloaded = true + logger.LogWarn(c, fmt.Sprintf( + "task protocol render hook overloaded; plugin=%s task=%s", + pinned.Plugin.Meta.Key, + taskID, + )) + } else { + logger.LogError(c, "task protocol render hook failed") + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_failed generation=%d plugin=%q mode=stream stage=render_events reason=hook_failed elapsed_ms=%d", + generation, + pluginKey, + hookElapsed.Milliseconds(), + ) + writeTaskPluginProtocolFailure(c, machine, lastStatus) + return + } + } + if !overloaded { + result, decodeErr := relay.DecodePluginProtocolEventResult(value, deps.protocolLimits) + if decodeErr != nil { + logger.LogError(c, "task protocol render result invalid: "+decodeErr.Error()) + writeTaskPluginProtocolFailure(c, machine, lastStatus) + return + } + events, applyErr := machine.ApplyTick(result, lastStatus) + if applyErr != nil { + logger.LogError(c, "task protocol state transition failed: "+applyErr.Error()) + writeTaskPluginProtocolFailure(c, machine, lastStatus) + return + } + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=render_events generation=%d plugin=%q mode=stream tick=%d status=%q semantic_events=%d wire_events=%d done=%t state_present=%t elapsed_ms=%d", + generation, + pluginKey, + tickNumber, + taskPluginDebugStatus(lastStatus), + len(result.Events), + len(events), + result.Done, + result.State.Present, + hookElapsed.Milliseconds(), + ) + for _, event := range events { + if err = writeTaskPluginProtocolEvent(c, event); err != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_write_failed generation=%d plugin=%q mode=stream stage=event event_type=%q sequence=%d", generation, pluginKey, event.Type, event.SequenceNumber) + return + } + } + if taskPluginProtocolEventsTerminal(events) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_complete generation=%d plugin=%q mode=stream reason=terminal status=%q ticks=%d", generation, pluginKey, taskPluginDebugStatus(lastStatus), tickNumber+1) + return + } + previous = result.State + } + + delay := pluginProtocolTickDelay(taskID, tickNumber, deps.tickInterval, deps.tickJitter) + tickNumber++ + if overloaded { + delay += deps.tickInterval + } else if hookElapsed > deps.tickInterval { + delay += deps.tickInterval + logger.LogWarn(c, fmt.Sprintf( + "task protocol render hook slow; plugin=%s task=%s elapsed_ms=%d", + pinned.Plugin.Meta.Key, + taskID, + hookElapsed.Milliseconds(), + )) + } + if !waitForTaskPluginProtocolTick(c, observationContext, heartbeatTicker, delay) { + if errors.Is(observationContext.Err(), context.DeadlineExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_timeout generation=%d plugin=%q mode=stream last_status=%q", generation, pluginKey, taskPluginDebugStatus(lastStatus)) + writeTaskPluginProtocolTimeout(c, machine, lastStatus) + } else if c.Request.Context().Err() != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q mode=stream stage=tick_wait", generation, pluginKey) + } else { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_write_failed generation=%d plugin=%q mode=stream stage=heartbeat", generation, pluginKey) + } + return + } + } +} + +func waitForTaskPluginProtocolTick( + c *gin.Context, + observationContext context.Context, + heartbeatTicker *time.Ticker, + delay time.Duration, +) bool { + tickTimer := time.NewTimer(delay) + defer tickTimer.Stop() + for { + select { + case <-c.Request.Context().Done(): + return false + case <-observationContext.Done(): + return false + case <-heartbeatTicker.C: + helper.ExtendWriteDeadline(c) + if err := writeTaskPluginProtocolHeartbeat(c); err != nil { + return false + } + case <-tickTimer.C: + return true + } + } +} + +func waitTaskPluginProtocol( + c *gin.Context, + pinned pluginruntime.PinnedEndpoint, + protocolRequest pluginruntime.ProtocolRequestContext, + taskID string, + machine *relay.PluginResponsesMachine, + deps pluginProtocolBridgeDeps, +) { + generation := pinned.Generation.Number + pluginKey := pinned.Plugin.Meta.Key + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_start generation=%d plugin=%q mode=nonstream public_task_id=%q timeout_ms=%d tick_ms=%d", + generation, + pluginKey, + taskID, + deps.observationTimeout.Milliseconds(), + deps.tickInterval.Milliseconds(), + ) + observationContext, cancelObservation := context.WithTimeout(c.Request.Context(), deps.observationTimeout) + defer cancelObservation() + tickNumber := uint64(0) + lastStatus := "" + for { + loadStarted := deps.now() + loadContext, cancelLoad := context.WithTimeout(observationContext, deps.loadTimeout) + task, exists, err := deps.loadTask( + loadContext, + common.GetContextKeyInt(c, constant.ContextKeyUserId), + constant.TaskPlatform(pinned.Plugin.Meta.Key), + taskID, + ) + loadContextErr := loadContext.Err() + cancelLoad() + loadElapsed := deps.now().Sub(loadStarted) + loadOverloaded := errors.Is(loadContextErr, context.DeadlineExceeded) && + observationContext.Err() == nil && + c.Request.Context().Err() == nil + if loadOverloaded { + logger.LogWarn(c, fmt.Sprintf( + "task protocol database observation overloaded; plugin=%s task=%s", + pinned.Plugin.Meta.Key, + taskID, + )) + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_tick generation=%d plugin=%q mode=nonstream tick=%d load_ms=%d overloaded=true", + generation, + pluginKey, + tickNumber, + loadElapsed.Milliseconds(), + ) + } else if err != nil || !exists || task == nil { + if errors.Is(observationContext.Err(), context.DeadlineExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_timeout generation=%d plugin=%q mode=nonstream last_status=%q", generation, pluginKey, taskPluginDebugStatus(lastStatus)) + writeTaskPluginProtocolTimeoutResponse(c, machine, lastStatus) + return + } + if err != nil && !errors.Is(err, context.Canceled) { + logger.LogError(c, "task protocol database observation failed") + } + if c.Request.Context().Err() == nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_failed generation=%d plugin=%q mode=nonstream stage=load reason=task_unavailable", generation, pluginKey) + writeTaskPluginProtocolFailureResponse(c, machine, lastStatus) + } else { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q mode=nonstream stage=load", generation, pluginKey) + } + return + } + overloaded := loadOverloaded + if !loadOverloaded { + previousStatus := lastStatus + lastStatus = string(task.Status) + if lastStatus != previousStatus { + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=status_transition generation=%d plugin=%q mode=nonstream tick=%d previous=%q status=%q load_ms=%d", + generation, + pluginKey, + tickNumber, + taskPluginDebugStatus(previousStatus), + taskPluginDebugStatus(lastStatus), + loadElapsed.Milliseconds(), + ) + } + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_tick generation=%d plugin=%q mode=nonstream tick=%d status=%q load_ms=%d overloaded=false", + generation, + pluginKey, + tickNumber, + taskPluginDebugStatus(lastStatus), + loadElapsed.Milliseconds(), + ) + } + if !loadOverloaded && (task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure) { + if task.Status == model.TaskStatusFailure { + writeTaskPluginProtocolFailureResponse(c, machine, string(task.Status)) + return + } + response, hookElapsed, callErr := renderTaskPluginProtocolFinalResponse( + observationContext, + pinned, + protocolRequest, + task, + machine, + deps, + ) + if callErr != nil { + if errors.Is(observationContext.Err(), context.DeadlineExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_timeout generation=%d plugin=%q mode=nonstream stage=render_final last_status=%q", generation, pluginKey, taskPluginDebugStatus(lastStatus)) + writeTaskPluginProtocolTimeoutResponse(c, machine, lastStatus) + return + } + if c.Request.Context().Err() != nil { + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q mode=nonstream stage=render_final", generation, pluginKey) + return + } + if errors.Is(callErr, pluginruntime.ErrCallAdmissionTimeout) { + overloaded = true + logger.LogWarn(c, fmt.Sprintf( + "task protocol final hook overloaded; plugin=%s task=%s", + pinned.Plugin.Meta.Key, + taskID, + )) + } else { + logger.LogError(c, "task protocol final hook failed") + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_failed generation=%d plugin=%q mode=nonstream stage=render_final reason=hook_failed elapsed_ms=%d", + generation, + pluginKey, + hookElapsed.Milliseconds(), + ) + writeTaskPluginProtocolFailureResponse(c, machine, lastStatus) + return + } + } else { + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=render_final generation=%d plugin=%q mode=nonstream status=%q elapsed_ms=%d", + generation, + pluginKey, + taskPluginDebugStatus(lastStatus), + hookElapsed.Milliseconds(), + ) + c.JSON(http.StatusOK, response) + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=observation_complete generation=%d plugin=%q mode=nonstream reason=terminal status=%q ticks=%d", + generation, + pluginKey, + taskPluginDebugStatus(lastStatus), + tickNumber+1, + ) + return + } + } + + delay := pluginProtocolTickDelay(taskID, tickNumber, deps.tickInterval, deps.tickJitter) + tickNumber++ + if overloaded { + delay += deps.tickInterval + } + tickTimer := time.NewTimer(delay) + select { + case <-c.Request.Context().Done(): + if !tickTimer.Stop() { + select { + case <-tickTimer.C: + default: + } + } + logger.LogDebug(c, "task_plugin subsystem=protocol event=client_disconnected generation=%d plugin=%q mode=nonstream stage=tick_wait", generation, pluginKey) + return + case <-observationContext.Done(): + if !tickTimer.Stop() { + select { + case <-tickTimer.C: + default: + } + } + if errors.Is(observationContext.Err(), context.DeadlineExceeded) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=observation_timeout generation=%d plugin=%q mode=nonstream last_status=%q", generation, pluginKey, taskPluginDebugStatus(lastStatus)) + writeTaskPluginProtocolTimeoutResponse(c, machine, lastStatus) + } + return + case <-tickTimer.C: + } + } +} + +func renderTaskPluginProtocolFinalResponse( + ctx context.Context, + pinned pluginruntime.PinnedEndpoint, + protocolRequest pluginruntime.ProtocolRequestContext, + task *model.Task, + machine *relay.PluginResponsesMachine, + deps pluginProtocolBridgeDeps, +) (map[string]any, time.Duration, error) { + view, err := service.BuildTaskPluginView(task) + if err != nil { + return nil, 0, err + } + viewValue, err := taskPluginProtocolJSONValue(view) + if err != nil { + return nil, 0, err + } + rendererContext, err := taskPluginProtocolRendererContext( + protocolRequest, + pinned, + task, + deps.artifactContentURL, + ) + if err != nil { + return nil, 0, err + } + hookStarted := deps.now() + payload, err := pinned.Plugin.Engine.CallPathWithAdmissionTimeout( + ctx, + deps.admissionTimeout, + "protocols", + []string{pinned.Protocol, "renderFinal"}, + rendererContext, + viewValue, + ) + hookElapsed := deps.now().Sub(hookStarted) + if err != nil { + return nil, hookElapsed, err + } + response, err := machine.FinalResponse(payload, string(task.Status)) + if err != nil { + return nil, hookElapsed, err + } + return response, hookElapsed, nil +} + +func renderTaskPluginProtocolEventsResponse( + ctx context.Context, + pinned pluginruntime.PinnedEndpoint, + protocolRequest pluginruntime.ProtocolRequestContext, + task *model.Task, + machine *relay.PluginResponsesMachine, + deps pluginProtocolBridgeDeps, +) (map[string]any, time.Duration, error) { + view, err := service.BuildTaskPluginView(task) + if err != nil { + return nil, 0, err + } + viewValue, err := taskPluginProtocolJSONValue(view) + if err != nil { + return nil, 0, err + } + rendererContext, err := taskPluginProtocolRendererContext( + protocolRequest, + pinned, + task, + deps.artifactContentURL, + ) + if err != nil { + return nil, 0, err + } + hookStarted := deps.now() + value, err := pinned.Plugin.Engine.CallPathWithAdmissionTimeout( + ctx, + deps.admissionTimeout, + "protocols", + []string{pinned.Protocol, "renderEvents"}, + rendererContext, + viewValue, + ) + hookElapsed := deps.now().Sub(hookStarted) + if err != nil { + return nil, hookElapsed, err + } + result, err := relay.DecodePluginProtocolEventResult(value, deps.protocolLimits) + if err != nil { + return nil, hookElapsed, err + } + response, err := machine.FinalFromEvents(result, string(task.Status)) + if err != nil { + return nil, hookElapsed, err + } + return response, hookElapsed, nil +} + +func RetrieveTaskPluginResponse(c *gin.Context) { + retrieveTaskPluginResponse(c, defaultPluginProtocolBridgeDeps()) +} + +func retrieveTaskPluginResponse(c *gin.Context, deps pluginProtocolBridgeDeps) { + deps = deps.withDefaults() + responseID := strings.TrimSpace(c.Param("response_id")) + if !strings.HasPrefix(responseID, "resp_") { + writeTaskPluginResponseNotFound(c, responseID, "bad_prefix") + return + } + taskID := "task_" + strings.TrimPrefix(responseID, "resp_") + userID := common.GetContextKeyInt(c, constant.ContextKeyUserId) + logger.LogDebug(c, "task_plugin subsystem=protocol event=retrieve_start response_id=%q public_task_id=%q", responseID, taskID) + + task, exists, err := deps.getByTaskId(userID, taskID) + if err != nil { + logger.LogError(c, "task protocol retrieve lookup failed") + logger.LogDebug(c, "task_plugin subsystem=protocol event=retrieve_failed reason=lookup_error public_task_id=%q", taskID) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + if !exists || task == nil { + writeTaskPluginResponseNotFound(c, responseID, "missing") + return + } + + plugin, generation, ok := deps.resolvePlugin(task.Platform) + if !ok || plugin == nil { + writeTaskPluginResponseNotFound(c, responseID, "no_plugin") + return + } + claimsProtocol := false + for _, claim := range plugin.Meta.Protocols { + if claim.Name == "openai_responses" { + claimsProtocol = true + break + } + } + if !claimsProtocol { + writeTaskPluginResponseNotFound(c, responseID, "no_claim") + return + } + + generationNumber := uint64(0) + if generation != nil { + generationNumber = generation.Number + } + createdAt := task.CreatedAt + if createdAt == 0 { + createdAt = task.SubmitTime + } + if createdAt == 0 { + createdAt = deps.now().Unix() + } + machine := relay.NewPluginResponsesMachine( + task.TaskID, + task.Properties.OriginModelName, + createdAt, + deps.protocolLimits, + ) + machine.SetBackground(task.PrivateData.ResponsesBackground) + pinned := pluginruntime.PinnedEndpoint{ + Generation: generation, + Plugin: plugin, + Protocol: "openai_responses", + Model: task.Properties.OriginModelName, + } + protocolRequest := pluginruntime.ProtocolRequestContext{ + RouteRequestContext: pluginruntime.RouteRequestContext{ + Path: c.Request.URL.Path, + Method: http.MethodGet, + Params: map[string]string{"response_id": responseID}, + Query: c.Request.URL.Query(), + Body: map[string]any{"kind": string(pluginruntime.BodyNone)}, + }, + Protocol: "openai_responses", + Operation: "retrieve", + Model: task.Properties.OriginModelName, + } + + if task.Status == model.TaskStatusFailure { + logger.LogDebug(c, "task_plugin subsystem=protocol event=retrieve_final generation=%d plugin=%q public_task_id=%q status=%q", generationNumber, plugin.Meta.Key, task.TaskID, taskPluginDebugStatus(string(task.Status))) + writeTaskPluginProtocolFailureResponse(c, machine, string(task.Status)) + return + } + if task.Status != model.TaskStatusSuccess { + logger.LogDebug(c, "task_plugin subsystem=protocol event=retrieve_pending generation=%d plugin=%q public_task_id=%q status=%q", generationNumber, plugin.Meta.Key, task.TaskID, taskPluginDebugStatus(string(task.Status))) + c.JSON(http.StatusOK, machine.PendingResponse(string(task.Status))) + return + } + + var ( + response map[string]any + hookElapsed time.Duration + renderErr error + ) + if plugin.Meta.ProtocolSupports("openai_responses", "sync") || plugin.Meta.ProtocolSupports("openai_responses", "background") { + response, hookElapsed, renderErr = renderTaskPluginProtocolFinalResponse( + c.Request.Context(), + pinned, + protocolRequest, + task, + machine, + deps, + ) + } else { + response, hookElapsed, renderErr = renderTaskPluginProtocolEventsResponse( + c.Request.Context(), + pinned, + protocolRequest, + task, + machine, + deps, + ) + } + if renderErr != nil { + logger.LogError(c, "task protocol retrieve render failed") + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=retrieve_failed generation=%d plugin=%q public_task_id=%q stage=render_final elapsed_ms=%d", + generationNumber, + plugin.Meta.Key, + task.TaskID, + hookElapsed.Milliseconds(), + ) + writeTaskPluginProtocolFailureResponse(c, machine, string(task.Status)) + return + } + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=retrieve_final generation=%d plugin=%q public_task_id=%q status=%q elapsed_ms=%d", + generationNumber, + plugin.Meta.Key, + task.TaskID, + taskPluginDebugStatus(string(task.Status)), + hookElapsed.Milliseconds(), + ) + c.JSON(http.StatusOK, response) +} + +func writeTaskPluginResponseNotFound(c *gin.Context, responseID, reason string) { + logger.LogDebug(c, "task_plugin subsystem=protocol event=retrieve_not_found reason=%s response_id=%q", reason, responseID) + respondPluginProtocolError(c, http.StatusNotFound, "not_found", "No response found with id '"+responseID+"'.") +} + +func writeTaskPluginProtocolHeartbeat(c *gin.Context) error { + if _, err := c.Writer.Write([]byte(": PING\n")); err != nil { + return err + } + return helper.FlushWriter(c) +} + +func writeTaskPluginProtocolFailure( + c *gin.Context, + machine *relay.PluginResponsesMachine, + taskStatus string, +) { + failed, err := machine.FailureEvent(taskStatus) + if err != nil { + logger.LogError(c, "task protocol failure event failed: "+err.Error()) + return + } + _ = writeTaskPluginProtocolEvent(c, failed) +} + +func writeTaskPluginProtocolTimeout( + c *gin.Context, + machine *relay.PluginResponsesMachine, + taskStatus string, +) { + incomplete, err := machine.TimeoutEvent(taskStatus) + if err != nil { + logger.LogError(c, "task protocol timeout event failed: "+err.Error()) + return + } + _ = writeTaskPluginProtocolEvent(c, incomplete) +} + +func writeTaskPluginProtocolFailureResponse( + c *gin.Context, + machine *relay.PluginResponsesMachine, + taskStatus string, +) { + if taskStatus == string(model.TaskStatusFailure) { + response, err := machine.FinalResponse(nil, taskStatus) + if err != nil { + logger.LogError(c, "task protocol terminal failure response failed: "+err.Error()) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + c.JSON(http.StatusOK, response) + return + } + response, err := machine.FailureResponse(taskStatus) + if err != nil { + logger.LogError(c, "task protocol failure response failed: "+err.Error()) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + c.JSON(http.StatusOK, response) +} + +func writeTaskPluginProtocolTimeoutResponse( + c *gin.Context, + machine *relay.PluginResponsesMachine, + lastStatus string, +) { + response, err := machine.TimeoutResponse(lastStatus) + if err != nil { + logger.LogError(c, "task protocol timeout response failed: "+err.Error()) + respondPluginProtocolError(c, http.StatusInternalServerError, "task_protocol_error", "Task protocol request failed") + return + } + c.JSON(http.StatusOK, response) +} + +func taskPluginProtocolJSONValue(value any) (any, error) { + encoded, err := common.Marshal(value) + if err != nil { + return nil, err + } + var decoded any + if err = common.Unmarshal(encoded, &decoded); err != nil { + return nil, err + } + return decoded, nil +} + +func taskPluginProtocolRendererContext( + request pluginruntime.ProtocolRequestContext, + pinned pluginruntime.PinnedEndpoint, + task *model.Task, + artifactContentURL func(taskID, artifactKey string) (string, error), +) (map[string]any, error) { + rendererContext := request.JSValue() + if task == nil || task.Status != model.TaskStatusSuccess { + return rendererContext, nil + } + if pinned.Plugin == nil { + return nil, errors.New("task artifact projection is unavailable") + } + + artifacts, err := taskjsplugin.New(pinned.Plugin).ListArtifacts(task) + if err != nil { + return nil, fmt.Errorf("project task artifacts: %w", err) + } + artifacts, err = validateProjectedTaskArtifacts(artifacts) + if err != nil { + return nil, err + } + if len(artifacts) > 0 && artifactContentURL == nil { + return nil, errors.New("task artifact projection is unavailable") + } + + rendererArtifacts := make(map[string]any, len(artifacts)) + for _, artifact := range artifacts { + contentURL, buildErr := artifactContentURL(task.TaskID, artifact.Key) + if buildErr != nil { + return nil, fmt.Errorf("build task artifact content URL: %w", buildErr) + } + item := map[string]any{ + "key": artifact.Key, + "type": artifact.Type, + "url": contentURL, + } + if artifact.MimeType != "" { + item["mimeType"] = artifact.MimeType + } + rendererArtifacts[artifact.Key] = item + } + rendererContext["artifacts"] = rendererArtifacts + return rendererContext, nil +} + +func pluginProtocolTickDelay(taskID string, tick uint64, base, jitter time.Duration) time.Duration { + if jitter <= 0 { + return base + } + hash := fnv.New64a() + _, _ = hash.Write([]byte(taskID)) + _, _ = hash.Write([]byte(":")) + _, _ = hash.Write([]byte(strconv.FormatUint(tick, 10))) + return base + time.Duration(hash.Sum64()%uint64(jitter+1)) +} + +func taskPluginProtocolEventsTerminal(events []dto.PluginResponsesStreamEvent) bool { + for _, event := range events { + switch event.Type { + case "response.completed", "response.failed", "response.incomplete": + return true + } + } + return false +} + +func writeTaskPluginProtocolEvent(c *gin.Context, event dto.PluginResponsesStreamEvent) error { + encoded, err := common.Marshal(event) + if err != nil { + return err + } + helper.ExtendWriteDeadline(c) + if _, err = c.Writer.Write([]byte("event: " + event.Type + "\n")); err != nil { + return err + } + if _, err = c.Writer.Write([]byte("data: " + string(encoded) + "\n\n")); err != nil { + return err + } + if err = helper.FlushWriter(c); err != nil { + return err + } + logger.LogDebug( + c, + "task_plugin subsystem=protocol event=sse_event_sent event_type=%q sequence=%d", + event.Type, + event.SequenceNumber, + ) + return nil +} + +func respondPluginProtocolSubmissionError(c *gin.Context, taskErr *dto.TaskError) { + status := http.StatusInternalServerError + if taskErr != nil && taskErr.StatusCode >= 400 && taskErr.StatusCode <= 599 { + status = taskErr.StatusCode + } + switch status { + case http.StatusBadRequest: + message := "Invalid task protocol request" + if taskErr != nil && taskErr.Message != "" && (taskErr.Code == "invalid_request" || strings.HasPrefix(taskErr.Code, "invalid_request")) { + message = taskErr.Message + } + respondPluginProtocolError(c, status, "invalid_request_error", message) + case http.StatusUnauthorized: + respondPluginProtocolError(c, status, "authentication_error", "Authentication failed") + case http.StatusForbidden: + respondPluginProtocolError(c, status, "permission_denied", "Task protocol request was denied") + case http.StatusTooManyRequests: + respondPluginProtocolError(c, status, "rate_limit_exceeded", "Too many requests") + default: + respondPluginProtocolError(c, status, "task_protocol_error", "Task protocol request failed") + } +} + +func respondPluginProtocolError(c *gin.Context, status int, code, message string) { + c.JSON(status, gin.H{ + "error": gin.H{ + "message": message, + "type": "new_api_error", + "code": code, + }, + }) +} diff --git a/controller/plugin_protocol_limiter.go b/controller/plugin_protocol_limiter.go new file mode 100644 index 000000000000..d88d177ee8de --- /dev/null +++ b/controller/plugin_protocol_limiter.go @@ -0,0 +1,157 @@ +package controller + +import ( + "errors" + "fmt" + "strings" + "sync" +) + +var ( + errPluginProtocolObservationLimitExceeded = errors.New("plugin protocol observation limit exceeded") + errInvalidPluginProtocolObservationIdentity = errors.New("invalid plugin protocol observation identity") +) + +type pluginProtocolObservationLimits struct { + global int + perPlugin int + perUser int + perToken int +} + +var defaultPluginProtocolObservationLimits = pluginProtocolObservationLimits{ + global: 128, + perPlugin: 32, + perUser: 4, + perToken: 2, +} + +var pluginProtocolObservationAdmissions = newPluginProtocolObservationLimiter( + defaultPluginProtocolObservationLimits, +) + +type pluginProtocolObservationLimitError struct { + scope string + limit int +} + +func (e *pluginProtocolObservationLimitError) Error() string { + return fmt.Sprintf("%s: %s capacity is %d", errPluginProtocolObservationLimitExceeded, e.scope, e.limit) +} + +func (e *pluginProtocolObservationLimitError) Unwrap() error { + return errPluginProtocolObservationLimitExceeded +} + +type pluginProtocolObservationLimiter struct { + mu sync.Mutex + + limits pluginProtocolObservationLimits + global int + plugin map[string]int + user map[int]int + token map[int]int +} + +func newPluginProtocolObservationLimiter(limits pluginProtocolObservationLimits) *pluginProtocolObservationLimiter { + return &pluginProtocolObservationLimiter{ + limits: limits, + plugin: make(map[string]int), + user: make(map[int]int), + token: make(map[int]int), + } +} + +func (l *pluginProtocolObservationLimiter) acquire( + pluginKey string, + userID int, + tokenID int, +) (func(), error) { + pluginKey = strings.TrimSpace(pluginKey) + switch { + case pluginKey == "": + return nil, fmt.Errorf("%w: plugin key is required", errInvalidPluginProtocolObservationIdentity) + case userID <= 0: + return nil, fmt.Errorf("%w: user id must be positive", errInvalidPluginProtocolObservationIdentity) + case tokenID <= 0: + return nil, fmt.Errorf("%w: token id must be positive", errInvalidPluginProtocolObservationIdentity) + } + + l.mu.Lock() + if l.global >= l.limits.global { + l.mu.Unlock() + return nil, &pluginProtocolObservationLimitError{ + scope: "global", + limit: l.limits.global, + } + } + l.global++ + + if l.plugin[pluginKey] >= l.limits.perPlugin { + l.global-- + l.mu.Unlock() + return nil, &pluginProtocolObservationLimitError{ + scope: "plugin", + limit: l.limits.perPlugin, + } + } + l.plugin[pluginKey]++ + + if l.user[userID] >= l.limits.perUser { + l.global-- + l.plugin[pluginKey]-- + if l.plugin[pluginKey] == 0 { + delete(l.plugin, pluginKey) + } + l.mu.Unlock() + return nil, &pluginProtocolObservationLimitError{ + scope: "user", + limit: l.limits.perUser, + } + } + l.user[userID]++ + + if l.token[tokenID] >= l.limits.perToken { + l.global-- + l.plugin[pluginKey]-- + if l.plugin[pluginKey] == 0 { + delete(l.plugin, pluginKey) + } + l.user[userID]-- + if l.user[userID] == 0 { + delete(l.user, userID) + } + l.mu.Unlock() + return nil, &pluginProtocolObservationLimitError{ + scope: "token", + limit: l.limits.perToken, + } + } + l.token[tokenID]++ + l.mu.Unlock() + + var once sync.Once + return func() { + once.Do(func() { + l.mu.Lock() + defer l.mu.Unlock() + + l.global-- + + l.plugin[pluginKey]-- + if l.plugin[pluginKey] == 0 { + delete(l.plugin, pluginKey) + } + + l.user[userID]-- + if l.user[userID] == 0 { + delete(l.user, userID) + } + + l.token[tokenID]-- + if l.token[tokenID] == 0 { + delete(l.token, tokenID) + } + }) + }, nil +} diff --git a/controller/plugin_protocol_limiter_test.go b/controller/plugin_protocol_limiter_test.go new file mode 100644 index 000000000000..a6b48cf65d43 --- /dev/null +++ b/controller/plugin_protocol_limiter_test.go @@ -0,0 +1,248 @@ +package controller + +import ( + "errors" + "sync" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestPluginProtocolObservationLimiterCaps(t *testing.T) { + t.Run("global", func(t *testing.T) { + limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{ + global: 2, + perPlugin: 2, + perUser: 2, + perToken: 2, + }) + releaseFirst, err := limiter.acquire("first", 1, 1) + require.NoError(t, err) + defer releaseFirst() + releaseSecond, err := limiter.acquire("second", 2, 2) + require.NoError(t, err) + defer releaseSecond() + + release, err := limiter.acquire("third", 3, 3) + assert.Nil(t, release) + assertLimitError(t, err, "global", 2) + }) + + t.Run("plugin", func(t *testing.T) { + limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{ + global: 3, + perPlugin: 1, + perUser: 3, + perToken: 3, + }) + releaseFirst, err := limiter.acquire("shared", 1, 1) + require.NoError(t, err) + defer releaseFirst() + + release, err := limiter.acquire("shared", 2, 2) + assert.Nil(t, release) + assertLimitError(t, err, "plugin", 1) + + releaseOther, err := limiter.acquire("other", 2, 2) + require.NoError(t, err) + defer releaseOther() + }) + + t.Run("user", func(t *testing.T) { + limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{ + global: 3, + perPlugin: 3, + perUser: 1, + perToken: 3, + }) + releaseFirst, err := limiter.acquire("first", 1, 1) + require.NoError(t, err) + defer releaseFirst() + + release, err := limiter.acquire("second", 1, 2) + assert.Nil(t, release) + assertLimitError(t, err, "user", 1) + + releaseOther, err := limiter.acquire("second", 2, 2) + require.NoError(t, err) + defer releaseOther() + }) + + t.Run("token", func(t *testing.T) { + limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{ + global: 3, + perPlugin: 3, + perUser: 3, + perToken: 1, + }) + releaseFirst, err := limiter.acquire("first", 1, 1) + require.NoError(t, err) + defer releaseFirst() + + release, err := limiter.acquire("second", 2, 1) + assert.Nil(t, release) + assertLimitError(t, err, "token", 1) + + releaseOther, err := limiter.acquire("second", 2, 2) + require.NoError(t, err) + defer releaseOther() + }) +} + +func TestPluginProtocolObservationLimiterRollsBackFailedAdmission(t *testing.T) { + t.Run("user failure", func(t *testing.T) { + limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{ + global: 2, + perPlugin: 2, + perUser: 1, + perToken: 2, + }) + releaseHeld, err := limiter.acquire("first", 1, 1) + require.NoError(t, err) + defer releaseHeld() + + release, err := limiter.acquire("second", 1, 2) + assert.Nil(t, release) + assertLimitError(t, err, "user", 1) + + releaseReplacement, err := limiter.acquire("second", 2, 2) + require.NoError(t, err) + defer releaseReplacement() + }) + + t.Run("token failure", func(t *testing.T) { + limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{ + global: 2, + perPlugin: 2, + perUser: 2, + perToken: 1, + }) + releaseHeld, err := limiter.acquire("first", 1, 1) + require.NoError(t, err) + defer releaseHeld() + + release, err := limiter.acquire("second", 2, 1) + assert.Nil(t, release) + assertLimitError(t, err, "token", 1) + + releaseReplacement, err := limiter.acquire("second", 2, 2) + require.NoError(t, err) + defer releaseReplacement() + }) +} + +func TestPluginProtocolObservationLimiterReleaseIsIdempotent(t *testing.T) { + limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{ + global: 1, + perPlugin: 1, + perUser: 1, + perToken: 1, + }) + release, err := limiter.acquire("plugin", 1, 1) + require.NoError(t, err) + + release() + release() + + releaseAgain, err := limiter.acquire("plugin", 1, 1) + require.NoError(t, err) + releaseAgain() +} + +func TestPluginProtocolObservationLimiterRejectsMissingIdentityWithoutConsumingCapacity(t *testing.T) { + limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{ + global: 1, + perPlugin: 1, + perUser: 1, + perToken: 1, + }) + + for _, testCase := range []struct { + name string + pluginKey string + userID int + tokenID int + }{ + {name: "empty plugin", userID: 1, tokenID: 1}, + {name: "blank plugin", pluginKey: " \t", userID: 1, tokenID: 1}, + {name: "zero user", pluginKey: "plugin", tokenID: 1}, + {name: "negative user", pluginKey: "plugin", userID: -1, tokenID: 1}, + {name: "zero token", pluginKey: "plugin", userID: 1}, + {name: "negative token", pluginKey: "plugin", userID: 1, tokenID: -1}, + } { + t.Run(testCase.name, func(t *testing.T) { + release, err := limiter.acquire(testCase.pluginKey, testCase.userID, testCase.tokenID) + assert.Nil(t, release) + assert.ErrorIs(t, err, errInvalidPluginProtocolObservationIdentity) + }) + } + + release, err := limiter.acquire("plugin", 1, 1) + require.NoError(t, err) + release() +} + +func TestPluginProtocolObservationLimiterConcurrentAdmissionsRespectCap(t *testing.T) { + const ( + workerCount = 8 + globalLimit = 3 + ) + limiter := newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{ + global: globalLimit, + perPlugin: workerCount, + perUser: workerCount, + perToken: workerCount, + }) + start := make(chan struct{}) + releases := make(chan func(), workerCount) + errorsFound := make(chan error, workerCount) + + var workers sync.WaitGroup + workers.Add(workerCount) + for worker := 1; worker <= workerCount; worker++ { + go func(id int) { + defer workers.Done() + <-start + release, err := limiter.acquire("plugin", id, id) + if err != nil { + errorsFound <- err + return + } + releases <- release + }(worker) + } + close(start) + workers.Wait() + close(releases) + close(errorsFound) + + assert.Len(t, releases, globalLimit) + assert.Len(t, errorsFound, workerCount-globalLimit) + for err := range errorsFound { + assert.ErrorIs(t, err, errPluginProtocolObservationLimitExceeded) + } + for release := range releases { + release() + } + + release, err := limiter.acquire("plugin", 1, 1) + require.NoError(t, err) + release() +} + +func assertLimitError( + t *testing.T, + err error, + expectedScope string, + expectedLimit int, +) { + t.Helper() + require.Error(t, err) + assert.ErrorIs(t, err, errPluginProtocolObservationLimitExceeded) + + var limitError *pluginProtocolObservationLimitError + require.True(t, errors.As(err, &limitError)) + assert.Equal(t, expectedScope, limitError.scope) + assert.Equal(t, expectedLimit, limitError.limit) +} diff --git a/controller/plugin_protocol_sdk_test.go b/controller/plugin_protocol_sdk_test.go new file mode 100644 index 000000000000..971794c9b249 --- /dev/null +++ b/controller/plugin_protocol_sdk_test.go @@ -0,0 +1,305 @@ +package controller + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" + pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin" + builtinplugins "github.com/QuantumNous/new-api/plugins" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/gin-gonic/gin" + "github.com/openai/openai-go" + "github.com/openai/openai-go/option" + "github.com/openai/openai-go/responses" + "github.com/openai/openai-go/shared" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestTaskPluginResponsesNonStreamDecodesWithOfficialGoSDK(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "official-sdk-non-stream", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("non-stream called renderEvents"); }, + renderFinal: function() { + return { + output: [{ + type: "message", + status: "completed", + role: "assistant", + content: [{ + type: "output_text", + text: "official-sdk-final", + annotations: [], + logprobs: [] + }] + }] + }; + } + }}; + `) + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_sdk_final", nil), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + return &model.Task{ + TaskID: "task_sdk_final", + UserId: 71, + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + Status: model.TaskStatusSuccess, + }, true, nil + } + server := newPluginProtocolSDKTestServer(t, pinned, deps) + defer server.Close() + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/v1/"), + option.WithMaxRetries(0), + ) + requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + response, err := client.Responses.New(requestContext, responses.ResponseNewParams{ + Model: shared.ResponsesModel("video-model"), + Input: responses.ResponseNewParamsInputUnion{ + OfString: openai.String("create a video"), + }, + }) + + require.NoError(t, err) + assert.Equal(t, "resp_sdk_final", response.ID) + assert.Equal(t, responses.ResponseStatusCompleted, response.Status) + assert.Equal(t, "video-model", response.Model) + assert.Equal(t, "official-sdk-final", response.OutputText()) + assert.Equal(t, "task_sdk_final", response.Metadata["task_id"]) +} + +func TestTaskPluginResponsesStreamDecodesWithOfficialGoSDK(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "official-sdk-stream", ` + export const protocols = {openai_responses: { + renderEvents: function() { + return {events: [{type: "output", data: "official-sdk-stream"}], done: true}; + }, + renderFinal: function() { throw new Error("stream called renderFinal"); } + }}; + `) + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_sdk_stream", nil), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + return &model.Task{ + TaskID: "task_sdk_stream", + UserId: 71, + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + Status: model.TaskStatusSuccess, + }, true, nil + } + server := newPluginProtocolSDKTestServer(t, pinned, deps) + defer server.Close() + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/v1/"), + option.WithMaxRetries(0), + ) + requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + stream := client.Responses.NewStreaming(requestContext, responses.ResponseNewParams{ + Model: shared.ResponsesModel("video-model"), + Input: responses.ResponseNewParamsInputUnion{ + OfString: openai.String("create a video"), + }, + }) + defer stream.Close() + eventTypes := make([]string, 0, 8) + sequenceNumbers := make([]int64, 0, 8) + var completedText string + for stream.Next() { + event := stream.Current() + eventTypes = append(eventTypes, event.Type) + sequenceNumbers = append(sequenceNumbers, event.SequenceNumber) + if event.Type == "response.completed" { + completedText = event.Response.OutputText() + } + } + + require.NoError(t, stream.Err()) + assert.Equal(t, []string{ + "response.created", + "response.output_item.added", + "response.content_part.added", + "response.output_text.delta", + "response.output_text.done", + "response.content_part.done", + "response.output_item.done", + "response.completed", + }, eventTypes) + assert.Equal(t, []int64{0, 1, 2, 3, 4, 5, 6, 7}, sequenceNumbers) + assert.Equal(t, "official-sdk-stream", completedText) +} + +func TestBuiltInKlingResponsesNonStreamDecodesWithOfficialGoSDK(t *testing.T) { + pinned, deps := builtInKlingProtocolSDKFixture(t, "task_kling_sdk_final") + server := newPluginProtocolSDKTestServer(t, pinned, deps) + defer server.Close() + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/v1/"), + option.WithMaxRetries(0), + ) + requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + response, err := client.Responses.New(requestContext, responses.ResponseNewParams{ + Model: shared.ResponsesModel("kling-v2-master"), + Input: responses.ResponseNewParamsInputUnion{ + OfString: openai.String("camera orbit"), + }, + }) + + require.NoError(t, err) + assert.Equal(t, responses.ResponseStatusCompleted, response.Status) + assert.Equal(t, "kling-v2-master", response.Model) + assert.Contains(t, response.OutputText(), "https://gateway.example/v1/tasks/task_kling_sdk_final/artifacts/video/content") + assert.NotContains(t, response.OutputText(), "upstream.example") + assert.Equal(t, "kling", response.Metadata["vendor"]) + assert.Equal(t, "task_kling_sdk_final", response.Metadata["task_id"]) +} + +func TestBuiltInKlingResponsesStreamDecodesWithOfficialGoSDK(t *testing.T) { + pinned, deps := builtInKlingProtocolSDKFixture(t, "task_kling_sdk_stream") + server := newPluginProtocolSDKTestServer(t, pinned, deps) + defer server.Close() + client := openai.NewClient( + option.WithAPIKey("test-key"), + option.WithBaseURL(server.URL+"/v1/"), + option.WithMaxRetries(0), + ) + requestContext, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + stream := client.Responses.NewStreaming(requestContext, responses.ResponseNewParams{ + Model: shared.ResponsesModel("kling-v2-master"), + Input: responses.ResponseNewParamsInputUnion{ + OfString: openai.String("camera orbit"), + }, + }) + defer stream.Close() + eventTypes := make([]string, 0, 8) + sequenceNumbers := make([]int64, 0, 8) + var completedText string + for stream.Next() { + event := stream.Current() + eventTypes = append(eventTypes, event.Type) + sequenceNumbers = append(sequenceNumbers, event.SequenceNumber) + if event.Type == "response.completed" { + completedText = event.Response.OutputText() + } + } + + require.NoError(t, stream.Err()) + assert.Equal(t, []string{ + "response.created", + "response.output_item.added", + "response.content_part.added", + "response.output_text.delta", + "response.output_text.done", + "response.content_part.done", + "response.output_item.done", + "response.completed", + }, eventTypes) + assert.Equal(t, []int64{0, 1, 2, 3, 4, 5, 6, 7}, sequenceNumbers) + assert.Contains(t, completedText, "https://gateway.example/v1/tasks/task_kling_sdk_stream/artifacts/video/content") + assert.NotContains(t, completedText, "upstream.example") +} + +func builtInKlingProtocolSDKFixture(t *testing.T, taskID string) (pluginruntime.PinnedEndpoint, pluginProtocolBridgeDeps) { + t.Helper() + source, err := builtinplugins.Source("kling") + require.NoError(t, err) + registry := pluginruntime.NewRegistry() + plugin, err := registry.RegisterFactory(source, pluginruntime.Options{Key: "kling"}) + require.NoError(t, err) + binding, found := registry.Generation().LookupEndpoint(http.MethodPost, "/v1/responses", "kling-v2-master") + require.True(t, found) + pinned := pluginruntime.PinnedEndpoint{ + Generation: registry.Generation(), + Plugin: plugin, + Protocol: binding.Protocol, + Operation: binding.Operation, + Model: binding.Model, + Candidates: []pluginruntime.ProtocolBinding{binding}, + } + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, plugin.Meta.Key, taskID, nil), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + task := &model.Task{ + TaskID: taskID, Platform: constant.TaskPlatform(plugin.Meta.Key), UserId: 71, + Status: model.TaskStatusSuccess, Progress: "100%", CreatedAt: 1_710_000_000, + } + task.SetData(map[string]any{ + "code": 0, + "data": map[string]any{ + "task_id": "upstream-private", + "task_status": "succeed", + "task_result": map[string]any{ + "videos": []any{map[string]any{"url": "https://upstream.example/private-video.mp4"}}, + }, + }, + }) + return task, true, nil + } + deps.artifactContentURL = func(publicTaskID, artifactKey string) (string, error) { + return "https://gateway.example/v1/tasks/" + publicTaskID + "/artifacts/" + artifactKey + "/content", nil + } + return pinned, deps +} + +func newPluginProtocolSDKTestServer( + t *testing.T, + pinned pluginruntime.PinnedEndpoint, + deps pluginProtocolBridgeDeps, +) *httptest.Server { + t.Helper() + return httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if request.Method != http.MethodPost || request.URL.Path != "/v1/responses" { + http.NotFound(writer, request) + return + } + var requestBody map[string]any + if err := common.DecodeJson(request.Body, &requestBody); err != nil { + http.Error(writer, "invalid request", http.StatusBadRequest) + return + } + modelName, _ := requestBody["model"].(string) + stream, _ := requestBody["stream"].(bool) + c, _ := gin.CreateTestContext(writer) + c.Request = request + common.SetContextKey(c, constant.ContextKeyUserId, 71) + common.SetContextKey(c, constant.ContextKeyTokenId, 81) + common.SetContextKey(c, constant.ContextKeyUsingGroup, "default") + c.Set("resolved_task_model", modelName) + c.Set(pluginruntime.ContextKeyProtocolRequest, pluginruntime.ProtocolRequestContext{ + RouteRequestContext: pluginruntime.RouteRequestContext{ + Path: request.URL.Path, + Method: request.Method, + Params: map[string]string{}, + Query: request.URL.Query(), + RequestBody: requestBody, + }, + Protocol: pinned.Protocol, + Stream: stream, + }) + serveTaskPluginProtocol(c, pinned, deps) + })) +} diff --git a/controller/plugin_protocol_test.go b/controller/plugin_protocol_test.go new file mode 100644 index 000000000000..ad4261ead6b5 --- /dev/null +++ b/controller/plugin_protocol_test.go @@ -0,0 +1,1581 @@ +package controller + +import ( + "context" + "errors" + "io" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" + pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin" + "github.com/QuantumNous/new-api/relay" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "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" +) + +func TestServeTaskPluginProtocolWaitsForDurableSubmissionBeforeWriting(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "durable-barrier", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + c, recorder := newPluginProtocolTestContext(true, false) + submitStarted := make(chan struct{}) + releaseSubmit := make(chan struct{}) + done := make(chan struct{}) + + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + close(submitStarted) + <-releaseSubmit + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_durable", map[string]any{ + "must_not": "be_written", + }), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + return nil, false, errors.New("observation failed after durable barrier") + } + + go func() { + defer close(done) + serveTaskPluginProtocol(c, pinned, deps) + }() + select { + case <-submitStarted: + case <-time.After(2 * time.Second): + require.FailNow(t, "submission did not start") + } + + assert.Empty(t, recorder.Header().Get("Content-Type")) + assert.Empty(t, recorder.Body.String()) + assert.False(t, recorder.Flushed) + + close(releaseSubmit) + select { + case <-done: + case <-time.After(2 * time.Second): + require.FailNow(t, "protocol handler did not finish") + } + assert.Equal(t, "text/event-stream", recorder.Header().Get("Content-Type")) + assert.Contains(t, recorder.Body.String(), "event: response.created\n") + assert.NotContains(t, recorder.Body.String(), "must_not") +} + +func TestServeTaskPluginProtocolDisconnectDuringSubmissionFinishesDurableWithoutWriting(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "disconnect-during-submit", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + c, recorder := newPluginProtocolTestContext(true, true) + requestContext, cancel := context.WithCancel(c.Request.Context()) + c.Request = c.Request.WithContext(requestContext) + submitStarted := make(chan struct{}) + checkSubmissionContext := make(chan struct{}) + submissionContextActive := make(chan struct{}) + releaseSubmit := make(chan struct{}) + observationStarted := make(chan struct{}, 1) + done := make(chan struct{}) + deps := pluginProtocolTestDeps() + deps.submit = func(c *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + close(submitStarted) + <-checkSubmissionContext + select { + case <-c.Request.Context().Done(): + return nil, service.TaskErrorWrapperLocal(c.Request.Context().Err(), "request_cancelled", http.StatusRequestTimeout) + default: + close(submissionContextActive) + } + <-releaseSubmit + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_disconnect_durable", nil), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + observationStarted <- struct{}{} + return nil, false, errors.New("observation must not start after disconnect") + } + + go func() { + defer close(done) + serveTaskPluginProtocol(c, pinned, deps) + }() + select { + case <-submitStarted: + case <-time.After(2 * time.Second): + require.FailNow(t, "submission did not start") + } + cancel() + close(checkSubmissionContext) + select { + case <-submissionContextActive: + case <-time.After(2 * time.Second): + require.FailNow(t, "submission context was canceled with the client") + } + select { + case <-done: + require.FailNow(t, "protocol handler stopped before submission became durable") + default: + } + close(releaseSubmit) + select { + case <-done: + case <-time.After(2 * time.Second): + require.FailNow(t, "protocol handler did not finish after durable submission") + } + + assert.Empty(t, recorder.Header().Get("Content-Type")) + assert.Empty(t, recorder.Body.String()) + assert.False(t, recorder.Flushed) + select { + case <-observationStarted: + require.FailNow(t, "protocol observation started after client disconnect") + default: + } +} + +func TestServeTaskPluginProtocolDisconnectBeforeDurableBarrierPersistsAndSettlesWithoutRefund(t *testing.T) { + events := make([]string, 0, 3) + database := setupTaskSubmissionDatabase(t, true, &events) + previousLogConsumeEnabled := common.LogConsumeEnabled + common.LogConsumeEnabled = false + t.Cleanup(func() { common.LogConsumeEnabled = previousLogConsumeEnabled }) + + pinned := compilePluginProtocolTestEndpoint(t, "disconnect-before-durable", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + c, recorder := newPluginProtocolTestContext(true, true) + requestContext, cancel := context.WithCancel(c.Request.Context()) + c.Request = c.Request.WithContext(requestContext) + billing := &taskSubmissionTestBilling{events: &events} + submitStarted := make(chan struct{}) + releaseSubmit := make(chan struct{}) + observationStarted := make(chan struct{}, 1) + done := make(chan struct{}) + + deps := pluginProtocolTestDeps() + deps.submit = func(c *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + info.Billing = billing + info.TaskRelayInfo.PublicTaskID = "task_disconnect_persisted" + info.TaskRelayInfo.LockedChannel = &model.Channel{ + Id: 1, + Type: constant.ChannelTypeTaskPlugin, + Name: "disconnect-before-durable", + } + info.ChannelMeta = &relaycommon.ChannelMeta{ + ChannelId: 1, + ChannelType: constant.ChannelTypeTaskPlugin, + } + return executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) { + close(submitStarted) + <-releaseSubmit + return &relay.TaskSubmitResult{ + UpstreamTaskID: "upstream_disconnect_persisted", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + Quota: 7, + }, nil + }) + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + observationStarted <- struct{}{} + return nil, false, errors.New("observation must not start after disconnect") + } + + go func() { + defer close(done) + serveTaskPluginProtocol(c, pinned, deps) + }() + select { + case <-submitStarted: + case <-time.After(2 * time.Second): + require.FailNow(t, "submission did not start") + } + cancel() + close(releaseSubmit) + select { + case <-done: + case <-time.After(2 * time.Second): + require.FailNow(t, "detached submission did not finish") + } + + assert.Equal(t, []string{"reserve", "insert", "settle"}, events) + assert.Zero(t, billing.refunds) + var persisted model.Task + require.NoError(t, database.Where("task_id = ?", "task_disconnect_persisted").First(&persisted).Error) + assert.Equal(t, model.TaskStatus(model.TaskStatusNotStart), persisted.Status) + assert.Equal(t, 7, persisted.Quota) + assert.Equal(t, "upstream_disconnect_persisted", persisted.PrivateData.UpstreamTaskID) + assert.Empty(t, recorder.Header().Get("Content-Type")) + assert.Empty(t, recorder.Body.String()) + assert.False(t, recorder.Flushed) + select { + case <-observationStarted: + require.FailNow(t, "protocol observation started after client disconnect") + default: + } +} + +func TestServeTaskPluginProtocolDisconnectDuringTerminalSettlementStopsOnlyObservation(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "disconnect-terminal-settlement", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + + previousDB := model.DB + previousMemoryCache := common.MemoryCacheEnabled + database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, database.AutoMigrate(&model.Channel{}, &model.Task{})) + model.DB = database + common.MemoryCacheEnabled = false + t.Cleanup(func() { + model.DB = previousDB + common.MemoryCacheEnabled = previousMemoryCache + }) + baseURL := "https://example.com" + channel := model.Channel{ + Type: constant.ChannelTypeTaskPlugin, + Name: "terminal-settlement", + Key: "test-key", + BaseURL: &baseURL, + Status: common.ChannelStatusEnabled, + } + require.NoError(t, database.Create(&channel).Error) + task := model.Task{ + TaskID: "task_terminal_disconnect", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + ChannelId: channel.Id, + Quota: 10, + Status: model.TaskStatusSubmitted, + PrivateData: model.TaskPrivateData{ + UpstreamTaskID: "upstream-terminal", + }, + } + require.NoError(t, database.Create(&task).Error) + + c, recorder := newPluginProtocolTestContext(true, true) + requestContext, cancel := context.WithCancel(c.Request.Context()) + c.Request = c.Request.WithContext(requestContext) + billingEvents := make([]string, 0) + billing := &taskSubmissionTestBilling{events: &billingEvents} + observationStarted := make(chan struct{}) + settlementStarted := make(chan struct{}) + releaseSettlement := make(chan struct{}) + t.Cleanup(func() { + select { + case <-releaseSettlement: + default: + close(releaseSettlement) + } + }) + pollingDone := make(chan struct{}) + adaptor := &terminalSettlementPollingAdaptor{ + started: settlementStarted, + release: releaseSettlement, + } + previousAdaptorFactory := service.GetTaskAdaptorFunc + service.GetTaskAdaptorFunc = func(constant.TaskPlatform) service.TaskPollingAdaptor { + return adaptor + } + t.Cleanup(func() { service.GetTaskAdaptorFunc = previousAdaptorFactory }) + + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + info.Billing = billing + return &taskSubmissionOutcome{ + Result: &relay.TaskSubmitResult{}, + Task: &task, + RelayInfo: info, + }, nil + } + deps.loadTask = func(ctx context.Context, _ int, _ constant.TaskPlatform, _ string) (*model.Task, bool, error) { + close(observationStarted) + <-ctx.Done() + return nil, false, ctx.Err() + } + done := make(chan struct{}) + + go func() { + <-observationStarted + defer close(pollingDone) + service.DispatchPlatformUpdate( + context.Background(), + task.Platform, + map[int][]string{channel.Id: {"upstream-terminal"}}, + map[string]*model.Task{"upstream-terminal": &task}, + ) + }() + go func() { + defer close(done) + serveTaskPluginProtocol(c, pinned, deps) + }() + select { + case <-settlementStarted: + case <-time.After(2 * time.Second): + require.FailNow(t, "terminal settlement did not start") + } + cancel() + select { + case <-done: + case <-time.After(2 * time.Second): + require.FailNow(t, "protocol observation did not stop after terminal disconnect") + } + assert.Equal(t, []string{"response.created"}, pluginProtocolTestSSEEventTypes(recorder.Body.String())) + assert.Zero(t, billing.refunds) + + close(releaseSettlement) + select { + case <-pollingDone: + case <-time.After(2 * time.Second): + require.FailNow(t, "terminal settlement was canceled with the client observation") + } + + var persisted model.Task + require.NoError(t, database.Where("task_id = ?", task.TaskID).First(&persisted).Error) + assert.Equal(t, model.TaskStatus(model.TaskStatusSuccess), persisted.Status) + assert.Equal(t, "100%", persisted.Progress) + assert.Equal(t, 10, persisted.Quota) + assert.True(t, adaptor.completed) + assert.Empty(t, billingEvents) +} + +type terminalSettlementPollingAdaptor struct { + started chan struct{} + release chan struct{} + completed bool +} + +func (a *terminalSettlementPollingAdaptor) Init(*relaycommon.RelayInfo) {} + +func (a *terminalSettlementPollingAdaptor) FetchTask(string, string, *model.Task, string) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{}`)), + }, nil +} + +func (a *terminalSettlementPollingAdaptor) ParseTaskResult(*model.Task, *http.Response, []byte) (*relaycommon.TaskInfo, error) { + return &relaycommon.TaskInfo{ + Status: model.TaskStatusSuccess, + Progress: "100%", + }, nil +} + +func (a *terminalSettlementPollingAdaptor) AdjustBillingOnComplete(task *model.Task, _ *relaycommon.TaskInfo) int { + close(a.started) + <-a.release + a.completed = true + return task.Quota +} + +func TestPluginProtocolBridgeBoundsDatabaseReadBelowHeartbeat(t *testing.T) { + deps := pluginProtocolBridgeDeps{ + observationTimeout: time.Minute, + loadTimeout: 10 * time.Second, + tickInterval: time.Second, + heartbeatInterval: 4 * time.Second, + admissionTimeout: time.Second, + }.withDefaults() + + assert.Equal(t, 2*time.Second, deps.loadTimeout) +} + +func TestServeTaskPluginProtocolPostDurableObservationFailureUsesCanonicalResponse(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "observation-failure", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + c, recorder := newPluginProtocolTestContext(false, false) + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_observation_failure", nil), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + return nil, false, errors.New("database-secret https://database.invalid/?token=hidden") + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, http.StatusOK, recorder.Code) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "failed", response.Status) + require.NotNil(t, response.Error) + assert.Equal(t, "server_error", response.Error.Code) + assert.Equal(t, "The task could not be observed.", response.Error.Message) + assert.Equal(t, "queued", response.Metadata["task_status"]) + assert.Equal(t, "/v1/responses/resp_observation_failure", response.Metadata["retrieval_path"]) + assert.NotContains(t, recorder.Body.String(), "secret") + assert.NotContains(t, recorder.Body.String(), "database.invalid") +} + +func TestServeTaskPluginProtocolStreamsPinnedGenerationWithHostFraming(t *testing.T) { + oldPinned := compilePluginProtocolTestEndpoint(t, "generation-pinned", ` + export const protocols = {openai_responses: { + renderEvents: function(ctx, task, previousState) { + if (ctx.stream !== true || ctx.body.value.stream !== true) { + throw new Error("host did not preserve parsed stream mode"); + } + if (arguments.length === 2) { + return {events: [], state: null, done: false}; + } + if (arguments.length !== 3 || previousState !== null) { + throw new Error("explicit null state was not supplied on the next tick"); + } + return {events: [{type: "output", data: "old-generation"}], done: true}; + }, + renderFinal: function() { throw new Error("stream called renderFinal"); } + }}; + `) + newPinned := compilePluginProtocolTestEndpoint(t, "generation-pinned", ` + export const protocols = {openai_responses: { + renderEvents: function() { + return {events: [{type: "output", data: "new-generation"}], done: true}; + }, + renderFinal: function() { return "new-generation"; } + }}; + `) + require.NotSame(t, oldPinned.Plugin.Engine, newPinned.Plugin.Engine) + + c, recorder := newPluginProtocolTestContext(true, true) + c.Set(pluginruntime.ContextKeyRouteRequest, pluginruntime.RouteRequestContext{ + Path: "/v1/responses", + Method: http.MethodPost, + RequestBody: map[string]any{"model": "video-model", "stream": false}, + }) + loadCount := 0 + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, oldPinned.Plugin.Meta.Key, "task_generation", map[string]any{ + "client_response": "ignored", + }), nil + } + deps.loadTask = func(_ context.Context, userID int, platform constant.TaskPlatform, taskID string) (*model.Task, bool, error) { + loadCount++ + assert.Equal(t, 71, userID) + assert.Equal(t, constant.TaskPlatform(oldPinned.Plugin.Meta.Key), platform) + assert.Equal(t, "task_generation", taskID) + status := model.TaskStatus(model.TaskStatusInProgress) + if loadCount == 2 { + status = model.TaskStatus(model.TaskStatusSuccess) + } + return &model.Task{ + TaskID: taskID, + UserId: userID, + Platform: platform, + Status: status, + }, true, nil + } + + serveTaskPluginProtocol(c, oldPinned, deps) + + assert.Equal(t, 2, loadCount) + assert.True(t, recorder.Flushed) + assert.Equal(t, "text/event-stream", recorder.Header().Get("Content-Type")) + assert.Equal(t, []string{ + "response.created", + "response.output_item.added", + "response.content_part.added", + "response.output_text.delta", + "response.output_text.done", + "response.content_part.done", + "response.output_item.done", + "response.completed", + }, pluginProtocolTestSSEEventTypes(recorder.Body.String())) + assert.True(t, strings.HasPrefix(recorder.Body.String(), "event: response.created\ndata: {")) + assert.Contains(t, recorder.Body.String(), `"sequence_number":0`) + assert.Contains(t, recorder.Body.String(), `"sequence_number":7`) + assert.Contains(t, recorder.Body.String(), "old-generation") + assert.NotContains(t, recorder.Body.String(), "new-generation") +} + +func TestServeTaskPluginProtocolStreamMissingRenderEventsUsesFailureEnvelope(t *testing.T) { + tests := []struct { + name string + status model.TaskStatus + }{ + {name: "success", status: model.TaskStatusSuccess}, + {name: "failure", status: model.TaskStatusFailure}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "default-events-"+testCase.name, ` + export const protocols = {openai_responses: { + renderFinal: function() { throw new Error("stream must not call renderFinal"); } + }}; + `) + c, recorder := newPluginProtocolTestContext(true, true) + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_default_events"), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + return &model.Task{TaskID: "task_default_events", Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), UserId: 71, Status: testCase.status}, true, nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, []string{"response.created", "response.failed"}, pluginProtocolTestSSEEventTypes(recorder.Body.String())) + assert.NotContains(t, recorder.Body.String(), "stream must not call") + }) + } +} + +func TestServeTaskPluginProtocolStreamInjectsHostArtifactCapabilities(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "stream-artifacts", ` + export function listArtifacts(task) { + if (task.data.output.video_url !== "https://upstream.invalid/video.mp4?secret=hidden") { + throw new Error("listArtifacts did not receive raw Task.Data"); + } + return [{key: "video", type: "video", mimeType: "video/mp4"}]; + } + export function buildContentRequest() { + throw new Error("rendering must not resolve provider content"); + } + export const protocols = {openai_responses: { + renderEvents: function(ctx, task) { + const artifact = ctx.artifacts && ctx.artifacts.video; + if (!artifact || artifact.key !== "video" || artifact.type !== "video" || + artifact.mimeType !== "video/mp4") { + throw new Error("host artifact context is invalid"); + } + return {events: [{type: "output", data: artifact.url}], done: true}; + }, + renderFinal: function() { throw new Error("stream called renderFinal"); } + }}; + `) + c, recorder := newPluginProtocolTestContext(true, true) + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_stream_artifact", nil), nil + } + task := &model.Task{ + TaskID: "task_stream_artifact", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + } + task.SetData(map[string]any{ + "output": map[string]any{ + "video_url": "https://upstream.invalid/video.mp4?secret=hidden", + }, + }) + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + return task, true, nil + } + deps.artifactContentURL = func(taskID, artifactKey string) (string, error) { + assert.Equal(t, "task_stream_artifact", taskID) + assert.Equal(t, "video", artifactKey) + return "https://gateway.example/v1/tasks/task_stream_artifact/artifacts/video/content?access=host-capability", nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, model.TaskStatus(model.TaskStatusSuccess), task.Status) + assert.Equal(t, []string{ + "response.created", + "response.output_item.added", + "response.content_part.added", + "response.output_text.delta", + "response.output_text.done", + "response.content_part.done", + "response.output_item.done", + "response.completed", + }, pluginProtocolTestSSEEventTypes(recorder.Body.String())) + assert.Contains(t, recorder.Body.String(), "host-capability") + assert.NotContains(t, recorder.Body.String(), "upstream.invalid") + assert.NotContains(t, recorder.Body.String(), "secret") +} + +func TestTaskPluginProtocolHeartbeatDoesNotDispatchEmptySDKEvent(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil) + + require.NoError(t, writeTaskPluginProtocolHeartbeat(c)) + + assert.Equal(t, ": PING\n", recorder.Body.String()) + assert.True(t, recorder.Flushed) +} + +func TestServeTaskPluginProtocolNonStreamUsesFinalHookAndHostEnvelope(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "final-response", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("non-stream called renderEvents"); }, + renderFinal: function(ctx, task) { + if (ctx.stream !== false || ctx.body.value.stream !== false) { + throw new Error("host did not preserve parsed non-stream mode"); + } + return { + id: "plugin-controlled-id", + status: "plugin-controlled-status", + metadata: {plugin_field: "kept", task_id: "plugin-controlled-task"}, + output: [{ + id: "plugin-controlled-item", + type: "message", + status: "plugin-controlled-item-status", + role: "assistant", + content: [{ + id: "plugin-controlled-content", + type: "output_text", + text: task.data.value, + annotations: [], + logprobs: [] + }] + }], + custom_field: "kept" + }; + } + }}; + `) + c, recorder := newPluginProtocolTestContext(false, false) + c.Set(pluginruntime.ContextKeyRouteRequest, pluginruntime.RouteRequestContext{ + Path: "/v1/responses", + Method: http.MethodPost, + RequestBody: map[string]any{"model": "video-model", "stream": true}, + }) + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_final", map[string]any{ + "client_response_secret": "must-be-ignored", + }), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + task := &model.Task{ + TaskID: "task_final", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + } + task.SetData(map[string]any{"value": "plugin-semantic-result"}) + return task, true, nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, http.StatusOK, recorder.Code) + assert.NotEqual(t, "text/event-stream", recorder.Header().Get("Content-Type")) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "resp_final", response.ID) + assert.Equal(t, "response", response.Object) + assert.Equal(t, "completed", response.Status) + assert.Equal(t, "video-model", response.Model) + assert.Equal(t, "task_final", response.Metadata["task_id"]) + assert.Equal(t, "kept", response.Metadata["plugin_field"]) + require.Len(t, response.Output, 1) + assert.Equal(t, "item_task_final_0", response.Output[0].ID) + assert.Equal(t, "completed", response.Output[0].Status) + require.Len(t, response.Output[0].Content, 1) + assert.Equal(t, "content_task_final_0_0", response.Output[0].Content[0].ID) + assert.Equal(t, "plugin-semantic-result", response.Output[0].Content[0].Text) + assert.Contains(t, recorder.Body.String(), `"custom_field":"kept"`) + assert.NotContains(t, recorder.Body.String(), "plugin-controlled-id") + assert.NotContains(t, recorder.Body.String(), "client_response_secret") +} + +func TestServeTaskPluginProtocolNonStreamInjectsHostArtifactCapabilities(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "final-artifacts", ` + export function listArtifacts() { + return [{key: "video", type: "video"}]; + } + export function buildContentRequest() { + throw new Error("rendering must not resolve provider content"); + } + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("non-stream called renderEvents"); }, + renderFinal: function(ctx) { + const artifact = ctx.artifacts && ctx.artifacts.video; + if (!artifact || artifact.key !== "video" || artifact.type !== "video" || + Object.prototype.hasOwnProperty.call(artifact, "mimeType")) { + throw new Error("host artifact context is invalid"); + } + return { + output: [{ + type: "message", + status: "completed", + role: "assistant", + content: [{ + type: "output_text", + text: artifact.url, + annotations: [], + logprobs: [] + }] + }] + }; + } + }}; + `) + c, recorder := newPluginProtocolTestContext(false, false) + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_final_artifact", nil), nil + } + task := &model.Task{ + TaskID: "task_final_artifact", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + return task, true, nil + } + deps.artifactContentURL = func(taskID, artifactKey string) (string, error) { + assert.Equal(t, "task_final_artifact", taskID) + assert.Equal(t, "video", artifactKey) + return "https://gateway.example/v1/tasks/task_final_artifact/artifacts/video/content?access=host-capability", nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, model.TaskStatus(model.TaskStatusSuccess), task.Status) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "completed", response.Status) + require.Len(t, response.Output, 1) + require.Len(t, response.Output[0].Content, 1) + assert.Contains(t, response.Output[0].Content[0].Text, "host-capability") +} + +func TestServeTaskPluginProtocolArtifactURLFailureOnlyFailsCurrentRendering(t *testing.T) { + for _, stream := range []bool{false, true} { + t.Run(strconv.FormatBool(stream), func(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "artifact-url-failure-"+strconv.FormatBool(stream), ` + export function listArtifacts() { + return [{key: "video", type: "video"}]; + } + export function buildContentRequest() { + throw new Error("unused"); + } + export const protocols = {openai_responses: { + renderEvents: function() { + return {events: [{type: "output", data: "must-not-render"}], done: true}; + }, + renderFinal: function() { + return {output: []}; + } + }}; + `) + c, recorder := newPluginProtocolTestContext(stream, stream) + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_capability_failure", nil), nil + } + task := &model.Task{ + TaskID: "task_capability_failure", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + return task, true, nil + } + deps.artifactContentURL = func(string, string) (string, error) { + return "", errors.New("public address is unavailable") + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, model.TaskStatus(model.TaskStatusSuccess), task.Status) + assert.NotContains(t, recorder.Body.String(), "must-not-render") + assert.NotContains(t, recorder.Body.String(), "public address") + if stream { + assert.Equal(t, []string{"response.created", "response.failed"}, pluginProtocolTestSSEEventTypes(recorder.Body.String())) + } else { + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "failed", response.Status) + require.NotNil(t, response.Error) + assert.Equal(t, "server_error", response.Error.Code) + assert.Equal(t, "completed", response.Metadata["task_status"]) + } + }) + } +} + +func TestServeTaskPluginProtocolNonStreamTaskFailureSkipsFinalHook(t *testing.T) { + logs := make([]string, 0, 1) + pinned := compilePluginProtocolTestEndpointWithOptions(t, "failed-final", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { + console.log("renderFinal called"); + return { + output: [{id: "secret-id", content: [{text: "plugin-secret"}]}], + secret: "https://secret.invalid/" + }; + } + }}; + `, pluginruntime.Options{ + Log: func(message string) { logs = append(logs, message) }, + }) + c, recorder := newPluginProtocolTestContext(false, false) + deps := pluginProtocolTestDeps() + deps.artifactContentURL = func(string, string) (string, error) { + require.FailNow(t, "failed tasks must not project artifact URLs") + return "", nil + } + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_failed", map[string]any{ + "credential": "client-response-secret", + }), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + task := &model.Task{ + TaskID: "task_failed", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusFailure, + FailReason: "upstream credential at https://secret.invalid/", + } + task.SetData(map[string]any{"secret": "database-secret"}) + return task, true, nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, http.StatusOK, recorder.Code) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "failed", response.Status) + require.NotNil(t, response.Error) + assert.Equal(t, "server_error", response.Error.Code) + assert.Equal(t, "The task failed.", response.Error.Message) + assert.Empty(t, response.Output) + assert.NotContains(t, recorder.Body.String(), "secret") + assert.NotContains(t, recorder.Body.String(), "credential") + assert.Empty(t, logs) +} + +func TestServeTaskPluginProtocolStreamTaskFailureSuppressesPluginAndDatabaseDetails(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "failed-stream", ` + export const protocols = {openai_responses: { + renderEvents: function() { + return { + events: [{type: "output", data: "plugin-secret https://plugin.invalid/?key=hidden"}], + done: true + }; + }, + renderFinal: function() { return "unused-secret"; } + }}; + `) + c, recorder := newPluginProtocolTestContext(true, false) + deps := pluginProtocolTestDeps() + deps.artifactContentURL = func(string, string) (string, error) { + require.FailNow(t, "failed tasks must not project artifact URLs") + return "", nil + } + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_stream_failed", nil), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + task := &model.Task{ + TaskID: "task_stream_failed", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusFailure, + FailReason: "database-secret https://database.invalid/?token=hidden", + } + task.SetData(map[string]any{"secret": "private-result"}) + return task, true, nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, []string{"response.created", "response.failed"}, pluginProtocolTestSSEEventTypes(recorder.Body.String())) + assert.Contains(t, recorder.Body.String(), `"code":"server_error"`) + assert.Contains(t, recorder.Body.String(), `"message":"The task failed."`) + assert.Contains(t, recorder.Body.String(), `"task_status":"failed"`) + assert.NotContains(t, recorder.Body.String(), "secret") + assert.NotContains(t, recorder.Body.String(), "invalid") + assert.NotContains(t, recorder.Body.String(), "hidden") +} + +func TestServeTaskPluginProtocolRejectsUnsupportedProtocolBeforeSubmission(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "unsupported-protocol", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + pinned.Protocol = "unsupported" + c, recorder := newPluginProtocolTestContext(false, false) + c.Set(pluginruntime.ContextKeyProtocolRequest, pluginruntime.ProtocolRequestContext{ + RouteRequestContext: pluginruntime.RouteRequestContext{ + Path: "/v1/videos", + Method: http.MethodPost, + RequestBody: map[string]any{"model": "video-model"}, + }, + Protocol: pinned.Protocol, + }) + submitted := false + deps := pluginProtocolTestDeps() + deps.submit = func(*gin.Context, *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + submitted = true + return nil, nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.False(t, submitted) + assert.Equal(t, http.StatusNotImplemented, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"code":"task_protocol_not_available"`) +} + +func TestServeTaskPluginProtocolRejectsObservationAdmissionBeforeSubmission(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "admission-limit", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + for _, stream := range []bool{false, true} { + t.Run(map[bool]string{false: "non-stream", true: "stream"}[stream], func(t *testing.T) { + c, recorder := newPluginProtocolTestContext(stream, stream) + submitted := false + deps := pluginProtocolTestDeps() + deps.admissions = newPluginProtocolObservationLimiter(pluginProtocolObservationLimits{ + global: 0, + perPlugin: 1, + perUser: 1, + perToken: 1, + }) + deps.submit = func(*gin.Context, *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + submitted = true + return nil, nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.False(t, submitted) + assert.Equal(t, http.StatusTooManyRequests, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"code":"rate_limit_exceeded"`) + assert.NotEqual(t, "text/event-stream", recorder.Header().Get("Content-Type")) + }) + } +} + +func TestServeTaskPluginProtocolBackgroundNonStreamReturnsPendingWithoutObservation(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "background-create", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("background create called renderEvents"); }, + renderFinal: function() { throw new Error("background create called renderFinal"); } + }}; + `) + c, recorder := newPluginProtocolTestContext(false, false) + setProtocolRequestBackground(c, true) + loadCalls := 0 + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_background", nil), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + loadCalls++ + return nil, false, errors.New("observation must not start for background create") + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, 0, loadCalls) + assert.Equal(t, http.StatusOK, recorder.Code) + var response map[string]any + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "resp_background", response["id"]) + assert.Equal(t, "response", response["object"]) + assert.Equal(t, "queued", response["status"]) + assert.Equal(t, true, response["background"]) + assert.Nil(t, response["completed_at"]) + assert.Nil(t, response["error"]) + assert.Nil(t, response["usage"]) + assert.Empty(t, response["output"]) + metadata, ok := response["metadata"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "task_background", metadata["task_id"]) + assert.Equal(t, "queued", metadata["task_status"]) + assert.Equal(t, "/v1/responses/resp_background", metadata["retrieval_path"]) + assert.NotEqual(t, "text/event-stream", recorder.Header().Get("Content-Type")) +} + +func TestServeTaskPluginProtocolBackgroundStreamEntersObservation(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "background-stream", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [{type: "output", data: "streamed"}], done: true}; }, + renderFinal: function() { throw new Error("stream called renderFinal"); } + }}; + `) + c, recorder := newPluginProtocolTestContext(true, true) + setProtocolRequestBackground(c, true) + loadCalls := 0 + deps := pluginProtocolTestDeps() + deps.submit = func(_ *gin.Context, info *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return pluginProtocolTestOutcome(info, pinned.Plugin.Meta.Key, "task_background_stream", nil), nil + } + deps.loadTask = func(context.Context, int, constant.TaskPlatform, string) (*model.Task, bool, error) { + loadCalls++ + return &model.Task{ + TaskID: "task_background_stream", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + }, true, nil + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Greater(t, loadCalls, 0) + assert.Equal(t, []string{ + "response.created", + "response.output_item.added", + "response.content_part.added", + "response.output_text.delta", + "response.output_text.done", + "response.content_part.done", + "response.output_item.done", + "response.completed", + }, pluginProtocolTestSSEEventTypes(recorder.Body.String())) +} + +func TestRetrieveTaskPluginResponsePendingSkipsRenderFinal(t *testing.T) { + logs := make([]string, 0, 1) + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-pending", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("retrieve pending called renderEvents"); }, + renderFinal: function() { + console.log("renderFinal called"); + return {}; + } + }}; + `, logsAppender(&logs)) + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_pending") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_pending", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusInProgress, + PrivateData: model.TaskPrivateData{ + ResponsesBackground: true, + }, + Properties: model.Properties{OriginModelName: "video-model"}, + CreatedAt: 1_710_000_000, + }, true, nil) + + retrieveTaskPluginResponse(c, deps) + + assert.Empty(t, logs) + assert.Equal(t, http.StatusOK, recorder.Code) + var response map[string]any + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "resp_retrieve_pending", response["id"]) + assert.Equal(t, "in_progress", response["status"]) + assert.Equal(t, "video-model", response["model"]) + assert.Equal(t, true, response["background"]) + assert.Nil(t, response["completed_at"]) + assert.Empty(t, response["output"]) + metadata, ok := response["metadata"].(map[string]any) + require.True(t, ok) + assert.Equal(t, "/v1/responses/resp_retrieve_pending", metadata["retrieval_path"]) +} + +func TestRetrieveTaskPluginResponseEchoesOriginModelName(t *testing.T) { + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-alias-echo", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("pending retrieve called renderEvents"); }, + renderFinal: function() { throw new Error("pending retrieve called renderFinal"); } + }}; + `, pluginruntime.Options{}) + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_alias") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_alias", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusInProgress, + Properties: model.Properties{OriginModelName: "alias-model"}, + CreatedAt: 1_710_000_000, + }, true, nil) + + retrieveTaskPluginResponse(c, deps) + + assert.Equal(t, http.StatusOK, recorder.Code) + var response map[string]any + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "alias-model", response["model"]) +} + +func TestRetrieveTaskPluginResponseSuccessRendersFinal(t *testing.T) { + logs := make([]string, 0, 1) + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-success", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("retrieve success called renderEvents"); }, + renderFinal: function() { + console.log("renderFinal called"); + return { + output: [{ + type: "message", + status: "completed", + role: "assistant", + content: [{type: "output_text", text: "retrieved-final", annotations: [], logprobs: []}] + }] + }; + } + }}; + `, logsAppender(&logs)) + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_success") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_success", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + Properties: model.Properties{OriginModelName: "video-model"}, + CreatedAt: 1_710_000_000, + }, true, nil) + + retrieveTaskPluginResponse(c, deps) + + require.NotEmpty(t, logs) + assert.Contains(t, logs[0], "renderFinal called") + assert.Equal(t, http.StatusOK, recorder.Code) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "completed", response.Status) + assert.Equal(t, "resp_retrieve_success", response.ID) + require.Len(t, response.Output, 1) + require.Len(t, response.Output[0].Content, 1) + assert.Equal(t, "retrieved-final", response.Output[0].Content[0].Text) +} + +func TestRetrieveTaskPluginResponseStreamOnlySuccessSynthesizesFromEvents(t *testing.T) { + logs := make([]string, 0, 1) + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-stream-only", ` + export const protocols = {openai_responses: { + renderEvents: function() { + console.log("renderEvents called"); + return {events: [{type: "output", data: "synthesized-retrieve"}], done: true}; + }, + renderFinal: function() { throw new Error("stream-only retrieve called renderFinal"); } + }}; + `, logsAppender(&logs)) + pinned.Plugin.Meta.Protocols = []pluginruntime.ProtocolClaim{{Name: "openai_responses", Supports: []string{"stream"}}} + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_stream") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_stream", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + Properties: model.Properties{OriginModelName: "video-model"}, + CreatedAt: 1_710_000_000, + }, true, nil) + + retrieveTaskPluginResponse(c, deps) + + require.NotEmpty(t, logs) + assert.Contains(t, logs[0], "renderEvents called") + assert.Equal(t, http.StatusOK, recorder.Code) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "completed", response.Status) + require.NotEmpty(t, response.Output) +} + +func TestRetrieveTaskPluginResponseStreamOnlyPendingAndFailureStayHostEnvelopes(t *testing.T) { + logs := make([]string, 0, 1) + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-stream-envelope", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("envelope retrieve called renderEvents"); }, + renderFinal: function() { throw new Error("envelope retrieve called renderFinal"); } + }}; + `, logsAppender(&logs)) + pinned.Plugin.Meta.Protocols = []pluginruntime.ProtocolClaim{{Name: "openai_responses", Supports: []string{"stream"}}} + + t.Run("pending", func(t *testing.T) { + logs = logs[:0] + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_stream_pending") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_stream_pending", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusInProgress, + Properties: model.Properties{OriginModelName: "video-model"}, + CreatedAt: 1_710_000_000, + }, true, nil) + retrieveTaskPluginResponse(c, deps) + assert.Empty(t, logs) + assert.Equal(t, http.StatusOK, recorder.Code) + var response map[string]any + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "in_progress", response["status"]) + assert.Empty(t, response["output"]) + }) + + t.Run("failure", func(t *testing.T) { + logs = logs[:0] + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_stream_failure") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_stream_failure", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusFailure, + Properties: model.Properties{OriginModelName: "video-model"}, + }, true, nil) + retrieveTaskPluginResponse(c, deps) + assert.Empty(t, logs) + assert.Equal(t, http.StatusOK, recorder.Code) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "failed", response.Status) + require.NotNil(t, response.Error) + assert.Equal(t, "The task failed.", response.Error.Message) + }) +} + +func TestRetrieveTaskPluginResponseStreamOnlyRenderErrorUsesFailureEnvelope(t *testing.T) { + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-stream-throw", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("retrieve boom"); }, + renderFinal: function() { throw new Error("stream-only retrieve called renderFinal"); } + }}; + `, pluginruntime.Options{}) + pinned.Plugin.Meta.Protocols = []pluginruntime.ProtocolClaim{{Name: "openai_responses", Supports: []string{"stream"}}} + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_stream_throw") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_stream_throw", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusSuccess, + Properties: model.Properties{OriginModelName: "video-model"}, + CreatedAt: 1_710_000_000, + }, true, nil) + + retrieveTaskPluginResponse(c, deps) + + assert.Equal(t, http.StatusOK, recorder.Code) + assert.NotContains(t, recorder.Body.String(), "retrieve boom") + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "failed", response.Status) + require.NotNil(t, response.Error) + assert.Equal(t, "server_error", response.Error.Code) + assert.Equal(t, "The task could not be observed.", response.Error.Message) +} + +func TestRetrieveTaskPluginResponseFailureUsesFailedEnvelope(t *testing.T) { + logs := make([]string, 0, 1) + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-failure", ` + export const protocols = {openai_responses: { + renderEvents: function() { throw new Error("retrieve failure called renderEvents"); }, + renderFinal: function() { + console.log("renderFinal called"); + return {}; + } + }}; + `, logsAppender(&logs)) + c, recorder := newPluginProtocolRetrieveContext("resp_retrieve_failure") + deps := pluginProtocolRetrieveDeps(pinned, &model.Task{ + TaskID: "task_retrieve_failure", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusFailure, + Properties: model.Properties{OriginModelName: "video-model"}, + }, true, nil) + + retrieveTaskPluginResponse(c, deps) + + assert.Empty(t, logs) + assert.Equal(t, http.StatusOK, recorder.Code) + var response dto.PluginResponsesResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "failed", response.Status) + require.NotNil(t, response.Error) + assert.Equal(t, "server_error", response.Error.Code) + assert.Equal(t, "The task failed.", response.Error.Message) +} + +func TestRetrieveTaskPluginResponseNotFound(t *testing.T) { + pinned := compilePluginProtocolRetrieveEndpoint(t, "retrieve-404", ` + export const protocols = {openai_responses: { + renderFinal: function() { return {}; } + }}; + `, pluginruntime.Options{}) + owned := &model.Task{ + TaskID: "task_owned", + Platform: constant.TaskPlatform(pinned.Plugin.Meta.Key), + UserId: 71, + Status: model.TaskStatusInProgress, + Properties: model.Properties{OriginModelName: "video-model"}, + } + + tests := []struct { + name string + responseID string + userID int + task *model.Task + exists bool + plugin *pluginruntime.LoadedPlugin + claims []pluginruntime.ProtocolClaim + }{ + {name: "bad prefix", responseID: "task_owned", userID: 71, task: owned, exists: true, plugin: pinned.Plugin}, + {name: "missing", responseID: "resp_missing", userID: 71, exists: false, plugin: pinned.Plugin}, + {name: "other user", responseID: "resp_owned", userID: 99, task: owned, exists: false, plugin: pinned.Plugin}, + {name: "no plugin", responseID: "resp_owned", userID: 71, task: owned, exists: true}, + {name: "plugin does not claim protocol", responseID: "resp_owned", userID: 71, task: owned, exists: true, plugin: pinned.Plugin, claims: []pluginruntime.ProtocolClaim{{Name: "openai_video"}}}, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + if testCase.plugin != nil { + if testCase.claims != nil { + testCase.plugin.Meta.Protocols = testCase.claims + } else { + testCase.plugin.Meta.Protocols = []pluginruntime.ProtocolClaim{{Name: "openai_responses", Supports: []string{"stream", "sync", "background"}}} + } + } + c, recorder := newPluginProtocolRetrieveContext(testCase.responseID) + common.SetContextKey(c, constant.ContextKeyUserId, testCase.userID) + deps := pluginProtocolRetrieveDeps(pinned, testCase.task, testCase.exists, nil) + if testCase.plugin == nil { + deps.resolvePlugin = func(constant.TaskPlatform) (*pluginruntime.LoadedPlugin, *pluginruntime.RoutingGeneration, bool) { + return nil, nil, false + } + } + + retrieveTaskPluginResponse(c, deps) + + assert.Equal(t, http.StatusNotFound, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"code":"not_found"`) + assert.Contains(t, recorder.Body.String(), "No response found with id '"+testCase.responseID+"'.") + }) + } +} + +func TestRespondPluginProtocolSubmissionErrorPassesValidationMessage(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "protocol-validation-detail", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + c, recorder := newPluginProtocolTestContext(false, false) + deps := pluginProtocolTestDeps() + deps.submit = func(*gin.Context, *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return nil, &dto.TaskError{ + Code: "invalid_request", + Message: "model is required", + StatusCode: http.StatusBadRequest, + LocalError: true, + } + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), `"message":"model is required"`) + assert.Contains(t, recorder.Body.String(), `"code":"invalid_request_error"`) + assert.NotContains(t, recorder.Body.String(), "Invalid task protocol request") +} + +func TestRespondPluginProtocolSubmissionErrorKeepsGenericNonValidation400(t *testing.T) { + pinned := compilePluginProtocolTestEndpoint(t, "protocol-generic-400", ` + export const protocols = {openai_responses: { + renderEvents: function() { return {events: [], done: false}; }, + renderFinal: function() { return {}; } + }}; + `) + c, recorder := newPluginProtocolTestContext(false, false) + deps := pluginProtocolTestDeps() + deps.submit = func(*gin.Context, *relaycommon.RelayInfo) (*taskSubmissionOutcome, *dto.TaskError) { + return nil, &dto.TaskError{ + Code: "task_not_exist", + Message: "task_origin_not_exist", + StatusCode: http.StatusBadRequest, + LocalError: true, + } + } + + serveTaskPluginProtocol(c, pinned, deps) + + assert.Equal(t, http.StatusBadRequest, recorder.Code) + assert.Contains(t, recorder.Body.String(), "Invalid task protocol request") + assert.NotContains(t, recorder.Body.String(), "task_origin_not_exist") +} + +func compilePluginProtocolTestEndpoint(t *testing.T, key, source string) pluginruntime.PinnedEndpoint { + t.Helper() + return compilePluginProtocolTestEndpointWithOptions(t, key, source, pluginruntime.Options{}) +} + +func compilePluginProtocolTestEndpointWithOptions( + t *testing.T, + key string, + source string, + options pluginruntime.Options, +) pluginruntime.PinnedEndpoint { + t.Helper() + options.Key = key + options.Version = "1.0.0" + options.Concurrency = 1 + engine, err := pluginruntime.Compile(source, options) + require.NoError(t, err) + return pluginruntime.PinnedEndpoint{ + Generation: &pluginruntime.RoutingGeneration{Number: 41}, + Plugin: &pluginruntime.LoadedPlugin{ + Meta: pluginruntime.Meta{ + Key: key, + Version: "1.0.0", + Protocols: []pluginruntime.ProtocolClaim{{ + Name: "openai_responses", + Supports: []string{"stream", "sync", "background"}, + }}, + }, + Engine: engine, + }, + Protocol: "openai_responses", + Operation: pluginruntime.HostProtocolOperation{Name: "create", Methods: []string{http.MethodPost}, Path: "/v1/responses", BodyKinds: []pluginruntime.BodyKind{pluginruntime.BodyJSON}, ModelField: "model"}, + Model: "video-model", + } +} + +func newPluginProtocolTestContext(stream, requestBodyStream bool) (*gin.Context, *httptest.ResponseRecorder) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(`{}`)) + common.SetContextKey(c, constant.ContextKeyUserId, 71) + common.SetContextKey(c, constant.ContextKeyTokenId, 81) + common.SetContextKey(c, constant.ContextKeyUsingGroup, "default") + c.Set("resolved_task_model", "video-model") + c.Set(pluginruntime.ContextKeyProtocolRequest, pluginruntime.ProtocolRequestContext{ + RouteRequestContext: pluginruntime.RouteRequestContext{ + Path: "/v1/responses", + Method: http.MethodPost, + Params: map[string]string{}, + Query: map[string][]string{}, + Body: map[string]any{ + "kind": "json", + "value": map[string]any{"model": "video-model", "stream": requestBodyStream}, + }, + RequestBody: map[string]any{ + "model": "video-model", + "stream": requestBodyStream, + }, + }, + Protocol: "openai_responses", + Stream: stream, + }) + return c, recorder +} + +func pluginProtocolTestDeps() pluginProtocolBridgeDeps { + return pluginProtocolBridgeDeps{ + now: func() time.Time { return time.Unix(1_710_000_000, 0) }, + admissions: newPluginProtocolObservationLimiter(defaultPluginProtocolObservationLimits), + protocolLimits: relay.DefaultPluginProtocolLimits(), + observationTimeout: time.Hour, + tickInterval: time.Nanosecond, + tickJitter: 0, + heartbeatInterval: time.Hour, + admissionTimeout: time.Second, + } +} + +func pluginProtocolTestOutcome( + info *relaycommon.RelayInfo, + pluginKey string, + taskID string, + _ ...any, +) *taskSubmissionOutcome { + return &taskSubmissionOutcome{ + Result: &relay.TaskSubmitResult{}, + Task: &model.Task{ + TaskID: taskID, + Platform: constant.TaskPlatform(pluginKey), + UserId: info.UserId, + Status: model.TaskStatusSubmitted, + CreatedAt: 1_710_000_000, + }, + RelayInfo: info, + } +} + +func pluginProtocolTestSSEEventTypes(body string) []string { + lines := strings.Split(body, "\n") + events := make([]string, 0) + for _, line := range lines { + if after, ok := strings.CutPrefix(line, "event: "); ok { + events = append(events, after) + } + } + return events +} + +func setProtocolRequestBackground(c *gin.Context, background bool) { + request := c.MustGet(pluginruntime.ContextKeyProtocolRequest).(pluginruntime.ProtocolRequestContext) + if body, ok := request.Body.(map[string]any); ok { + if value, ok := body["value"].(map[string]any); ok { + value["background"] = background + } + } +} + +func compilePluginProtocolRetrieveEndpoint(t *testing.T, key, source string, options pluginruntime.Options) pluginruntime.PinnedEndpoint { + t.Helper() + pinned := compilePluginProtocolTestEndpointWithOptions(t, key, source, options) + pinned.Plugin.Meta.Protocols = []pluginruntime.ProtocolClaim{{Name: "openai_responses", Supports: []string{"stream", "sync", "background"}}} + return pinned +} + +func logsAppender(logs *[]string) pluginruntime.Options { + if logs == nil { + return pluginruntime.Options{} + } + return pluginruntime.Options{ + Log: func(message string) { *logs = append(*logs, message) }, + } +} + +func newPluginProtocolRetrieveContext(responseID string) (*gin.Context, *httptest.ResponseRecorder) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses/"+responseID, nil) + c.Params = gin.Params{{Key: "response_id", Value: responseID}} + common.SetContextKey(c, constant.ContextKeyUserId, 71) + common.SetContextKey(c, constant.ContextKeyTokenId, 81) + return c, recorder +} + +func pluginProtocolRetrieveDeps(pinned pluginruntime.PinnedEndpoint, task *model.Task, exists bool, err error) pluginProtocolBridgeDeps { + deps := pluginProtocolTestDeps() + deps.getByTaskId = func(userId int, taskId string) (*model.Task, bool, error) { + if !exists { + return nil, false, err + } + if task != nil && (userId != task.UserId || taskId != task.TaskID) { + return nil, false, err + } + return task, task != nil, err + } + deps.resolvePlugin = func(constant.TaskPlatform) (*pluginruntime.LoadedPlugin, *pluginruntime.RoutingGeneration, bool) { + if pinned.Plugin == nil { + return nil, nil, false + } + return pinned.Plugin, pinned.Generation, true + } + return deps +} diff --git a/controller/redemption.go b/controller/redemption.go index 838746e7dd7f..86289f8a2dbf 100644 --- a/controller/redemption.go +++ b/controller/redemption.go @@ -1,6 +1,7 @@ package controller import ( + "errors" "net/http" "strconv" "unicode/utf8" @@ -85,6 +86,14 @@ func AddRedemption(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgRedemptionCountMax) return } + if redemption.Quota <= 0 { + common.ApiError(c, errors.New("redemption quota must be positive")) + return + } + if err := common.ValidateWalletQuota(redemption.Quota); err != nil { + common.ApiError(c, err) + return + } if valid, msg := validateExpiredTime(c, redemption.ExpiredTime); !valid { c.JSON(http.StatusOK, gin.H{"success": false, "message": msg}) return @@ -153,6 +162,14 @@ func UpdateRedemption(c *gin.Context) { return } if statusOnly == "" { + if redemption.Quota <= 0 { + common.ApiError(c, errors.New("redemption quota must be positive")) + return + } + if err := common.ValidateWalletQuota(redemption.Quota); err != nil { + common.ApiError(c, err) + return + } if valid, msg := validateExpiredTime(c, redemption.ExpiredTime); !valid { c.JSON(http.StatusOK, gin.H{"success": false, "message": msg}) return diff --git a/controller/relay.go b/controller/relay.go index 7e4270867463..9da14e339f30 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -15,8 +15,10 @@ import ( "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/model" + pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin" perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics" "github.com/QuantumNous/new-api/relay" + "github.com/QuantumNous/new-api/relay/channel/task/taskcommon" relaycommon "github.com/QuantumNous/new-api/relay/common" relayconstant "github.com/QuantumNous/new-api/relay/constant" "github.com/QuantumNous/new-api/relay/helper" @@ -115,7 +117,7 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) { newAPIError = types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry()) } else { - newAPIError = types.NewError(err, types.ErrorCodeInvalidRequest) + newAPIError = types.NewError(err, types.ErrorCodeInvalidRequest, types.ErrOptionWithStatusCode(http.StatusBadRequest), types.ErrOptionWithSkipRetry()) } return } @@ -199,12 +201,12 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { newAPIError = channelErr break } + addUsedChannel(c, channel.Id) if billingErr := service.PrepareTieredBillingForSelectedGroup(c, relayInfo); billingErr != nil { newAPIError = billingErr break } - addUsedChannel(c, channel.Id) bodyStorage, bodyErr := common.GetBodyStorage(c) if bodyErr != nil { // Ensure consistent 413 for oversized bodies even when error occurs later (e.g., retry path) @@ -236,7 +238,7 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { newAPIError = service.NormalizeViolationFeeError(newAPIError) relayInfo.LastError = newAPIError - processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError) + 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()) { break @@ -255,6 +257,38 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { } } +// CountClaudeTokens implements Anthropic's token-counting utility endpoint. +// It deliberately skips upstream generation and billing; callers use this +// endpoint to size prompts before creating a Message. +func CountClaudeTokens(c *gin.Context) { + request, err := helper.GetAndValidateClaudeRequest(c) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{ + "type": "error", + "error": gin.H{ + "type": "invalid_request_error", + "message": common.MessageWithRequestId(err.Error(), c.GetString(common.RequestIdKey)), + }, + }) + return + } + + info := relaycommon.GenRelayInfoClaude(c, request) + inputTokens, err := service.CountRequestToken(c, request.GetTokenCountMeta(), info) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "type": "error", + "error": gin.H{ + "type": "api_error", + "message": common.MessageWithRequestId(err.Error(), c.GetString(common.RequestIdKey)), + }, + }) + return + } + + c.JSON(http.StatusOK, gin.H{"input_tokens": inputTokens}) +} + var upgrader = websocket.Upgrader{ Subprotocols: []string{"realtime"}, // WS 握手支持的协议,如果有使用 Sec-WebSocket-Protocol,则必须在此声明对应的 Protocol TODO add other protocol CheckOrigin: func(r *http.Request) bool { @@ -312,9 +346,6 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service }, nil } channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam) - - info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info) - if err != nil { return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败(retry): %s", selectGroup, info.OriginModelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) } @@ -322,6 +353,8 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service return nil, types.NewError(fmt.Errorf("分组 %s 下模型 %s 的可用渠道不存在(retry)", selectGroup, info.OriginModelName), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) } + info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info) + newAPIError := middleware.SetupContextForSelectedChannel(c, channel, info.OriginModelName) if newAPIError != nil { return nil, newAPIError @@ -345,7 +378,7 @@ func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) b if retryTimes <= 0 { return false } - if _, ok := c.Get("specific_channel_id"); ok { + if service.GetChannelConstraints(c).SuppressesRetry() { return false } code := openaiErr.StatusCode @@ -361,7 +394,7 @@ func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) b return operation_setting.ShouldRetryByStatusCode(code) } -func processChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError) { +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 @@ -378,32 +411,21 @@ func processChannelError(c *gin.Context, channelError types.ChannelError, err *t modelName := c.GetString("original_model") tokenId := c.GetInt("token_id") userGroup := c.GetString("group") - channelId := c.GetInt("channel_id") - other := make(map[string]interface{}) + other := model.NewLogOther() if c.Request != nil && c.Request.URL != nil { - other["request_path"] = c.Request.URL.Path - } - other["error_type"] = err.GetErrorType() - other["error_code"] = err.GetErrorCode() - other["status_code"] = err.StatusCode - other["channel_id"] = channelId - other["channel_name"] = c.GetString("channel_name") - other["channel_type"] = c.GetInt("channel_type") - adminInfo := make(map[string]interface{}) - adminInfo["use_channel"] = c.GetStringSlice("use_channel") - isMultiKey := common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey) - if isMultiKey { - adminInfo["is_multi_key"] = true - adminInfo["multi_key_index"] = common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex) - } - service.AppendChannelAffinityAdminInfo(c, adminInfo) - other["admin_info"] = adminInfo + 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()) - model.RecordErrorLog(c, userId, channelId, modelName, tokenName, err.MaskSensitiveErrorWithStatusCode(), tokenId, useTimeSeconds, common.GetContextKeyBool(c, constant.ContextKeyIsStream), userGroup, other) + model.RecordErrorLog(c, userId, channelError.ChannelId, modelName, tokenName, err.MaskSensitiveErrorWithStatusCode(), tokenId, useTimeSeconds, common.GetContextKeyBool(c, constant.ContextKeyIsStream), userGroup, other) } } @@ -475,6 +497,33 @@ func RelayNotFound(c *gin.Context) { }) } +// RelayTaskPluginEndpoint keeps unclaimed shared-endpoint traffic on its +// existing handler while claimed requests enter the generation-pinned +// host-owned protocol bridge. +func RelayTaskPluginEndpoint(c *gin.Context, fallback gin.HandlerFunc) { + pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedEndpoint) + if !exists { + fallback(c) + return + } + pinned, ok := pinnedValue.(pluginruntime.PinnedEndpoint) + if !ok || pinned.Plugin == nil || pinned.Generation == nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "error": gin.H{ + "message": "Task protocol request failed", + "type": "new_api_error", + "code": "task_protocol_error", + }, + }) + return + } + if pinned.Protocol != "openai_responses" { + fallback(c) + return + } + serveTaskPluginProtocol(c, pinned, defaultPluginProtocolBridgeDeps()) +} + func RelayTaskFetch(c *gin.Context) { relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil) if err != nil { @@ -490,29 +539,76 @@ func RelayTaskFetch(c *gin.Context) { } } +type taskSubmissionOutcome struct { + Result *relay.TaskSubmitResult + Task *model.Task + RelayInfo *relaycommon.RelayInfo +} + func RelayTask(c *gin.Context) { relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil) if err != nil { - c.JSON(http.StatusInternalServerError, &taskdto.TaskError{ + respondTaskSubmissionError(c, &taskdto.TaskError{ Code: "gen_relay_info_failed", Message: err.Error(), StatusCode: http.StatusInternalServerError, }) return } + if action := c.GetString("task_action"); action != "" { + relayInfo.Action = action + } if taskErr := relay.ResolveOriginTask(c, relayInfo); taskErr != nil { - respondTaskError(c, taskErr) + respondTaskSubmissionError(c, taskErr) return } + if taskErr := relay.ApplyOriginTaskAffinity(c, relayInfo); taskErr != nil { + respondTaskSubmissionError(c, taskErr) + return + } + + outcome, taskErr := executeTaskSubmission(c, relayInfo) + if taskErr != nil { + respondTaskSubmissionError(c, taskErr) + return + } + presentTaskSubmission(c, outcome) +} + +// executeTaskSubmission owns the retry, billing, and persistence lifecycle. +// It deliberately performs no client response writes so JSON and protocol +// presenters share the same durable task barrier. Its cancellation semantics +// come from c.Request.Context: native task endpoints use the client context, +// while the Responses bridge supplies an independently bounded context. +func executeTaskSubmission(c *gin.Context, relayInfo *relaycommon.RelayInfo) (*taskSubmissionOutcome, *taskdto.TaskError) { + return executeTaskSubmissionWith(c, relayInfo, relay.RelayTaskSubmit) +} +type taskSubmitAttempt func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *taskdto.TaskError) + +func executeTaskSubmissionWith( + c *gin.Context, + relayInfo *relaycommon.RelayInfo, + submit taskSubmitAttempt, +) (*taskSubmissionOutcome, *taskdto.TaskError) { + diagnostics := newTaskPluginSubmitDiagnostics(c) + diagnostics.start(relayInfo) var result *relay.TaskSubmitResult var taskErr *taskdto.TaskError + durable := false + stage := "start" defer func() { - if taskErr != nil && relayInfo.Billing != nil { + if !durable && relayInfo.Billing != nil { + diagnostics.refund(stage) relayInfo.Billing.Refund(c) } }() + stage = "before_attempt" + if requestErr := c.Request.Context().Err(); requestErr != nil { + diagnostics.cancelled("before_attempt", 0) + return nil, service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout) + } retryParam := &service.RetryParam{ Ctx: c, @@ -523,6 +619,12 @@ func RelayTask(c *gin.Context) { } for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() { + stage = "select_channel" + if requestErr := c.Request.Context().Err(); requestErr != nil { + diagnostics.cancelled("before_attempt", retryParam.GetRetry()+1) + taskErr = service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout) + break + } var channel *model.Channel if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil { @@ -542,10 +644,12 @@ func RelayTask(c *gin.Context) { break } } + diagnostics.attempt(retryParam.GetRetry()+1, channel, relayInfo.LockedChannel != nil) addUsedChannel(c, channel.Id) bodyStorage, bodyErr := common.GetBodyStorage(c) if bodyErr != nil { + stage = "read_body" if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) { taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusRequestEntityTooLarge) } else { @@ -555,8 +659,15 @@ func RelayTask(c *gin.Context) { } c.Request.Body = io.NopCloser(bodyStorage) - result, taskErr = relay.RelayTaskSubmit(c, relayInfo) + stage = "submit" + result, taskErr = submit(c, relayInfo) + if requestErr := c.Request.Context().Err(); requestErr != nil { + diagnostics.cancelled("after_submit", retryParam.GetRetry()+1) + taskErr = service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout) + break + } if taskErr == nil { + diagnostics.attemptSucceeded(retryParam.GetRetry()+1, result) break } @@ -564,10 +675,13 @@ func RelayTask(c *gin.Context) { 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)) + types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode), + relayInfo) } - if !shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry()) { + willRetry := shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry()) + diagnostics.attemptFailed(retryParam.GetRetry()+1, channel, taskErr, willRetry) + if !willRetry { break } } @@ -578,38 +692,160 @@ func RelayTask(c *gin.Context) { logger.LogInfo(c, retryLogStr) } - // ── 成功:结算 + 日志 + 插入任务 ── - if taskErr == nil { - if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil { - common.SysError("settle task billing error: " + settleErr.Error()) + if taskErr != nil { + diagnostics.failed(stage, "task_error", taskErr, false) + return nil, taskErr + } + if result == nil { + taskErr = service.TaskErrorWrapperLocal(errors.New("task submission returned no result"), "task_submit_failed", http.StatusInternalServerError) + diagnostics.failed("submit", "missing_result", taskErr, false) + return nil, taskErr + } + if requestErr := c.Request.Context().Err(); requestErr != nil { + diagnostics.cancelled("before_reserve", retryParam.GetRetry()+1) + return nil, service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout) + } + + // Reserve any submit-time upward billing adjustment before persistence. + // This keeps insertion failures fully refundable while ensuring settlement + // after the barrier normally has a zero positive delta. + if relayInfo.Billing != nil { + stage = "reserve" + diagnostics.reserve("reserve_start", result.Quota) + if reserveErr := relayInfo.Billing.Reserve(result.Quota); reserveErr != nil { + common.SysError("reserve adjusted task billing error: " + reserveErr.Error()) + taskErr = service.TaskErrorWrapperLocal(errors.New("insufficient quota for adjusted task cost"), string(types.ErrorCodeInsufficientUserQuota), http.StatusForbidden) + diagnostics.failed("reserve", "insufficient_quota", taskErr, false) + return nil, taskErr + } + diagnostics.reserve("reserve_complete", result.Quota) + } + if requestErr := c.Request.Context().Err(); requestErr != nil { + diagnostics.cancelled("before_insert", retryParam.GetRetry()+1) + return nil, service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout) + } + + stage = "insert" + task := model.InitTask(result.Platform, relayInfo) + task.PrivateData.Execution = service.TaskExecutionSnapshotFromContext(c) + task.PrivateData.UpstreamTaskID = result.UpstreamTaskID + task.PrivateData.BillingSource = relayInfo.BillingSource + task.PrivateData.SubscriptionId = relayInfo.SubscriptionId + task.PrivateData.TokenId = relayInfo.TokenId + task.PrivateData.NodeName = common.NodeName + task.PrivateData.BillingContext = &model.TaskBillingContext{ + ModelPrice: relayInfo.PriceData.ModelPrice, + GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio, + ModelRatio: relayInfo.PriceData.ModelRatio, + OtherRatios: relayInfo.PriceData.OtherRatios(), + OriginModelName: relayInfo.OriginModelName, + PerCallBilling: common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName) || relayInfo.PriceData.UsePrice, + TieredSnapshot: relayInfo.TieredBillingSnapshot, + } + task.Quota = result.Quota + task.Data = result.TaskData + if len(result.PluginState) > 0 { + task.PrivateData.PluginState = result.PluginState + } + task.Action = relayInfo.Action + if immediate := result.Immediate; immediate != nil { + task.Status = model.TaskStatus(immediate.Status) + task.Progress = immediate.Progress + if immediate.Status == model.TaskStatusSuccess || immediate.Status == model.TaskStatusFailure { + task.FinishTime = time.Now().Unix() } - service.LogTaskConsumption(c, relayInfo) + if immediate.Status == model.TaskStatusFailure { + task.FailReason = immediate.Reason + } + if immediate.Url != "" { + task.PrivateData.ResultURL = immediate.Url + } else if immediate.Status == model.TaskStatusSuccess { + task.PrivateData.ResultURL = taskcommon.BuildProxyURL(task.TaskID) + } + } + diagnostics.insertStart(task) + if insertErr := task.InsertWithContext(c.Request.Context()); insertErr != nil { + common.SysError("insert task error: " + insertErr.Error()) + taskErr = service.TaskErrorWrapperLocal(errors.New("failed to persist task"), "task_insert_failed", http.StatusInternalServerError) + diagnostics.failed("insert", "database_error", taskErr, false) + return nil, taskErr + } + durable = true + stage = "settle" + diagnostics.durable(task) + diagnostics.settleStart(task, result.Quota) - task := model.InitTask(result.Platform, relayInfo) - task.PrivateData.UpstreamTaskID = result.UpstreamTaskID - task.PrivateData.BillingSource = relayInfo.BillingSource - task.PrivateData.SubscriptionId = relayInfo.SubscriptionId - task.PrivateData.TokenId = relayInfo.TokenId - task.PrivateData.NodeName = common.NodeName - task.PrivateData.BillingContext = &model.TaskBillingContext{ - ModelPrice: relayInfo.PriceData.ModelPrice, - GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio, - ModelRatio: relayInfo.PriceData.ModelRatio, - OtherRatios: relayInfo.PriceData.OtherRatios(), - OriginModelName: relayInfo.OriginModelName, - PerCallBilling: common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName) || relayInfo.PriceData.UsePrice, + if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil { + common.SysError("settle task billing error: " + settleErr.Error()) + taskErr = service.TaskErrorWrapperLocal(errors.New("failed to settle task billing"), "task_billing_settlement_failed", http.StatusInternalServerError) + diagnostics.failed("settle", "billing_error", taskErr, true) + return nil, taskErr + } + service.LogTaskConsumption(c, relayInfo, task) + diagnostics.complete(task, result.Quota) + + return &taskSubmissionOutcome{Result: result, Task: task, RelayInfo: relayInfo}, nil +} + +func presentTaskSubmission(c *gin.Context, outcome *taskSubmissionOutcome) { + diagnostics := newTaskPluginSubmitDiagnostics(c) + otherRatios := outcome.RelayInfo.PriceData.OtherRatios() + if otherRatios == nil { + otherRatios = map[string]float64{} + } + if ratiosJSON, err := common.Marshal(otherRatios); err == nil { + c.Header("X-New-Api-Other-Ratios", string(ratiosJSON)) + } + if pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedRoute); exists { + if pinned, ok := pinnedValue.(pluginruntime.PinnedRoute); ok && pinned.Plugin != nil && pinned.Route.Render != "" { + view, err := service.BuildTaskPluginView(outcome.Task) + requestValue, _ := c.Get(pluginruntime.ContextKeyRouteRequest) + requestContext, _ := requestValue.(pluginruntime.RouteRequestContext) + if err == nil { + viewValue, valueErr := taskPluginProtocolJSONValue(view) + if valueErr == nil { + if body, callErr := pinned.Plugin.Engine.CallPath(c.Request.Context(), "native", []string{pinned.Route.Render}, requestContext.JSValue(), viewValue); callErr == nil { + diagnostics.present(outcome.Task, "native_presenter") + c.JSON(http.StatusOK, body) + return + } else { + logger.LogError(c, "task plugin native submit presenter failed: "+callErr.Error()) + } + } else { + logger.LogError(c, "encode task plugin native submit view failed: "+valueErr.Error()) + } + } else { + logger.LogError(c, "build task plugin native submit view failed: "+err.Error()) + } } - task.Quota = result.Quota - task.Data = result.TaskData - task.Action = relayInfo.Action - if insertErr := task.Insert(); insertErr != nil { - common.SysError("insert task error: " + insertErr.Error()) + } + if pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedEndpoint); exists { + if pinned, ok := pinnedValue.(pluginruntime.PinnedEndpoint); ok && pinned.Protocol == "openai_video" && pinned.Operation.Name == "create" { + diagnostics.present(outcome.Task, "openai_video_create") + c.JSON(http.StatusOK, outcome.Task.ToOpenAIVideo()) + return } } + createdAt := outcome.Task.CreatedAt + if createdAt == 0 { + createdAt = outcome.Task.SubmitTime + } + diagnostics.present(outcome.Task, "host_fallback") + c.JSON(http.StatusOK, map[string]any{ + "id": outcome.Task.TaskID, + "task_id": outcome.Task.TaskID, + "status": "queued", + "model": outcome.RelayInfo.OriginModelName, + "created_at": createdAt, + }) +} - if taskErr != nil { - respondTaskError(c, taskErr) +func respondTaskSubmissionError(c *gin.Context, taskErr *taskdto.TaskError) { + newTaskPluginSubmitDiagnostics(c).presentError(taskErr) + if middleware.RespondTaskPluginError(c, taskErr) { + return } + respondTaskError(c, taskErr) } // respondTaskError 统一输出 Task 错误响应(含 429 限流提示改写) @@ -630,7 +866,7 @@ func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *taskdto.TaskEr if retryTimes <= 0 { return false } - if _, ok := c.Get("specific_channel_id"); ok { + if service.GetChannelConstraints(c).SuppressesRetry() { return false } if taskErr.StatusCode == http.StatusTooManyRequests { diff --git a/controller/relay_count_tokens_test.go b/controller/relay_count_tokens_test.go new file mode 100644 index 000000000000..2889bb76a78c --- /dev/null +++ b/controller/relay_count_tokens_test.go @@ -0,0 +1,73 @@ +package controller + +import ( + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestCountClaudeTokensReturnsInputTokensWhenRelayCountingDisabled(t *testing.T) { + gin.SetMode(gin.TestMode) + originalCountToken := constant.CountToken + constant.CountToken = false + t.Cleanup(func() { + constant.CountToken = originalCountToken + }) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest( + http.MethodPost, + "/v1/messages/count_tokens?beta=true", + strings.NewReader(`{ + "model":"gemini-3.6-flash", + "messages":[{"role":"user","content":"count this prompt"}], + "tools":[{"name":"lookup","description":"Look up a value","input_schema":{"type":"object","properties":{"query":{"type":"string"}}}}] + }`), + ) + ctx.Request.Header.Set("Content-Type", "application/json") + common.SetContextKey(ctx, constant.ContextKeyOriginalModel, "gemini-3.6-flash") + + CountClaudeTokens(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + var response struct { + InputTokens int `json:"input_tokens"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Positive(t, response.InputTokens) +} + +func TestCountClaudeTokensRejectsMissingMessages(t *testing.T) { + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest( + http.MethodPost, + "/v1/messages/count_tokens", + strings.NewReader(`{"model":"gemini-3.6-flash"}`), + ) + ctx.Request.Header.Set("Content-Type", "application/json") + + CountClaudeTokens(ctx) + + require.Equal(t, http.StatusBadRequest, recorder.Code) + var response struct { + Type string `json:"type"` + Error struct { + Type string `json:"type"` + Message string `json:"message"` + } `json:"error"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, "error", response.Type) + assert.Equal(t, "invalid_request_error", response.Error.Type) + assert.Contains(t, response.Error.Message, "messages") +} diff --git a/controller/relay_error_log_test.go b/controller/relay_error_log_test.go new file mode 100644 index 000000000000..737c05805959 --- /dev/null +++ b/controller/relay_error_log_test.go @@ -0,0 +1,99 @@ +package controller + +import ( + "errors" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/relaykit/types" + + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestProcessChannelErrorUsesSnapshotWithoutLeakingChannelMetadata(t *testing.T) { + gin.SetMode(gin.TestMode) + previousDB, previousLogDB := model.DB, model.LOG_DB + previousRedisEnabled := common.RedisEnabled + previousMainDatabaseType := common.MainDatabaseType() + previousLogDatabaseType := common.LogDatabaseType() + previousErrorLogEnabled := constant.ErrorLogEnabled + + database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + sqlDB, err := database.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + require.NoError(t, database.AutoMigrate(&model.User{}, &model.Log{})) + model.DB, model.LOG_DB = database, database + common.RedisEnabled = false + common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite) + constant.ErrorLogEnabled = true + t.Cleanup(func() { + model.DB, model.LOG_DB = previousDB, previousLogDB + common.RedisEnabled = previousRedisEnabled + common.SetDatabaseTypes(previousMainDatabaseType, previousLogDatabaseType) + constant.ErrorLogEnabled = previousErrorLogEnabled + require.NoError(t, sqlDB.Close()) + }) + + require.NoError(t, database.Create(&model.User{Id: 7, Username: "log-owner", Group: "default"}).Error) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + ctx.Set("id", 7) + ctx.Set("username", "log-owner") + ctx.Set("token_name", "test-token") + ctx.Set("token_id", 11) + ctx.Set("original_model", "gpt-test") + ctx.Set("group", "default") + ctx.Set("channel_id", 202) + ctx.Set("channel_name", "mutable-context-channel") + ctx.Set("channel_type", 9) + ctx.Set("use_channel", []string{"101"}) + common.SetContextKey(ctx, constant.ContextKeyRequestStartTime, time.Now().Add(-time.Second)) + + channelSnapshot := types.ChannelError{ + ChannelId: 101, + ChannelType: 1, + ChannelName: "snapshot-channel", + AutoBan: false, + } + apiErr := types.NewOpenAIError(errors.New("upstream failed"), types.ErrorCodeBadResponseStatusCode, http.StatusBadGateway) + + processChannelError(ctx, channelSnapshot, apiErr, nil) + + var stored model.Log + require.NoError(t, database.First(&stored).Error) + assert.Equal(t, channelSnapshot.ChannelId, stored.ChannelId) + storedOther, err := common.StrToMap(stored.Other) + require.NoError(t, err) + assert.Equal(t, float64(http.StatusBadGateway), storedOther["status_code"]) + for _, key := range []string{"channel_id", "channel_name", "channel_type"} { + assert.NotContains(t, storedOther, key) + } + adminInfo, ok := storedOther["admin_info"].(map[string]interface{}) + require.True(t, ok) + assert.Equal(t, []interface{}{"101"}, adminInfo["use_channel"]) + + logs, total, err := model.GetUserLogs(7, model.LogTypeError, 0, 0, "", "", 0, 10, "", "", "") + require.NoError(t, err) + require.Equal(t, int64(1), total) + require.Len(t, logs, 1) + assert.Equal(t, channelSnapshot.ChannelId, logs[0].ChannelId) + assert.Empty(t, logs[0].ChannelName) + userOther, err := common.StrToMap(logs[0].Other) + require.NoError(t, err) + assert.NotContains(t, userOther, "admin_info") + for _, key := range []string{"channel_id", "channel_name", "channel_type"} { + assert.NotContains(t, userOther, key) + } +} diff --git a/controller/relay_task_plugin_test.go b/controller/relay_task_plugin_test.go new file mode 100644 index 000000000000..e9405e2e6280 --- /dev/null +++ b/controller/relay_task_plugin_test.go @@ -0,0 +1,389 @@ +package controller + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" + pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin" + "github.com/QuantumNous/new-api/relay" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +type taskSubmissionTestBilling struct { + events *[]string + settleErr error + onSettle func() + refunds int +} + +func (b *taskSubmissionTestBilling) Settle(int) error { + *b.events = append(*b.events, "settle") + if b.onSettle != nil { + b.onSettle() + } + return b.settleErr +} + +func (b *taskSubmissionTestBilling) Refund(*gin.Context) { + *b.events = append(*b.events, "refund") + b.refunds++ +} + +func (b *taskSubmissionTestBilling) NeedsRefund() bool { return b.refunds == 0 } +func (b *taskSubmissionTestBilling) GetPreConsumedQuota() int { return 0 } +func (b *taskSubmissionTestBilling) Reserve(int) error { + *b.events = append(*b.events, "reserve") + return nil +} + +func TestPresentTaskSubmissionUsesNativePresenterAfterPersistence(t *testing.T) { + plugin, err := pluginruntime.CompilePlugin(` +export const meta = {apiVersion:1,key:"presenter-test",name:"Presenter",version:"1.0.0",author:{name:"Test"},models:["model"],fetchMode:"per_task",routes:[{method:"POST",path:"/vendor/jobs",type:"submit",decode:"decode",render:"created"}]}; +export const native = {decode:function(ctx){return {kind:"submit",model:"model",requestBody:ctx.body.value};},created:function(ctx,task){return {data:{task_id:task.task_id},upstream:task.data};}}; +export function buildSubmitRequest(){return {}} export function parseSubmitResponse(){return {taskId:"upstream"}} export function buildQueryRequest(){return {}} export function parseTaskResult(){return {status:"SUCCESS"}} +`, pluginruntime.Options{}) + require.NoError(t, err) + priceData := types.PriceData{} + priceData.AddOtherRatio("seconds", 5) + task := &model.Task{TaskID: "task_public", SubmitTime: 123} + task.SetData(map[string]any{"task_id": "upstream_private"}) + outcome := &taskSubmissionOutcome{ + Result: &relay.TaskSubmitResult{}, + Task: task, + RelayInfo: &relaycommon.RelayInfo{PriceData: priceData}, + } + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/vendor/jobs", strings.NewReader(`{"model":"model"}`)) + c.Set(pluginruntime.ContextKeyPinnedRoute, pluginruntime.PinnedRoute{Plugin: plugin, Route: plugin.Meta.Routes[0]}) + c.Set(pluginruntime.ContextKeyRouteRequest, pluginruntime.RouteRequestContext{Path: "/vendor/jobs", Method: http.MethodPost, Body: map[string]any{"kind": "json", "value": map[string]any{"model": "model"}}}) + + presentTaskSubmission(c, outcome) + + assert.JSONEq(t, `{ + "data":{"task_id":"task_public"}, + "upstream":{"task_id":"upstream_private"} + }`, recorder.Body.String()) + assert.JSONEq(t, `{"seconds":5}`, recorder.Header().Get("X-New-Api-Other-Ratios")) +} + +func TestPresentTaskSubmissionFallbackUsesPersistedPublicID(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + outcome := &taskSubmissionOutcome{ + Result: &relay.TaskSubmitResult{}, + Task: &model.Task{TaskID: "task_persisted", SubmitTime: 456}, + RelayInfo: &relaycommon.RelayInfo{OriginModelName: "video-model"}, + } + + presentTaskSubmission(c, outcome) + + assert.JSONEq(t, `{ + "id":"task_persisted", + "task_id":"task_persisted", + "status":"queued", + "model":"video-model", + "created_at":456 + }`, recorder.Body.String()) +} + +func TestPresentTaskSubmissionUsesHostOpenAIVideoCreateReceipt(t *testing.T) { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Set(pluginruntime.ContextKeyPinnedEndpoint, pluginruntime.PinnedEndpoint{ + Protocol: "openai_video", + Operation: pluginruntime.HostProtocolOperation{Name: "create"}, + }) + task := &model.Task{ + TaskID: "task_public", + Status: model.TaskStatusSubmitted, + Progress: "0%", + CreatedAt: 456, + Properties: model.Properties{OriginModelName: "video-model"}, + } + outcome := &taskSubmissionOutcome{Result: &relay.TaskSubmitResult{}, Task: task, RelayInfo: &relaycommon.RelayInfo{}} + + presentTaskSubmission(c, outcome) + + assert.JSONEq(t, `{"id":"task_public","object":"video","model":"video-model","status":"queued","progress":0,"created_at":456}`, recorder.Body.String()) + assert.NotContains(t, recorder.Body.String(), "task_id") +} + +func TestExecuteTaskSubmissionRefundsWhenInsertFails(t *testing.T) { + events := make([]string, 0, 3) + database := setupTaskSubmissionDatabase(t, false, &events) + _ = database + billing := &taskSubmissionTestBilling{events: &events} + c := taskSubmissionTestContext() + info := taskSubmissionRelayInfo(billing) + + outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) { + return &relay.TaskSubmitResult{ + UpstreamTaskID: "upstream_private", + Platform: constant.TaskPlatform("plugin"), + }, nil + }) + + assert.Nil(t, outcome) + require.NotNil(t, taskErr) + assert.Equal(t, "task_insert_failed", taskErr.Code) + assert.Equal(t, []string{"reserve", "insert", "refund"}, events) + assert.Equal(t, 1, billing.refunds) + assert.False(t, c.Writer.Written()) +} + +func TestExecuteTaskSubmissionSettlementFailureStaysDurableAndWritesNothing(t *testing.T) { + events := make([]string, 0, 3) + database := setupTaskSubmissionDatabase(t, true, &events) + billing := &taskSubmissionTestBilling{events: &events, settleErr: errors.New("settlement failed")} + c := taskSubmissionTestContext() + info := taskSubmissionRelayInfo(billing) + + outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) { + return &relay.TaskSubmitResult{ + UpstreamTaskID: "upstream_private", + Platform: constant.TaskPlatform("plugin"), + }, nil + }) + + assert.Nil(t, outcome) + require.NotNil(t, taskErr) + assert.Equal(t, "task_billing_settlement_failed", taskErr.Code) + assert.Equal(t, []string{"reserve", "insert", "settle"}, events) + assert.Zero(t, billing.refunds) + var count int64 + require.NoError(t, database.Model(&model.Task{}).Where("task_id = ?", "task_public").Count(&count).Error) + assert.Equal(t, int64(1), count) + assert.False(t, c.Writer.Written()) +} + +func TestExecuteTaskSubmissionPersistsPinnedPluginProvenance(t *testing.T) { + events := make([]string, 0, 3) + database := setupTaskSubmissionDatabase(t, true, &events) + previousLogConsumeEnabled := common.LogConsumeEnabled + common.LogConsumeEnabled = false + t.Cleanup(func() { common.LogConsumeEnabled = previousLogConsumeEnabled }) + + c := taskSubmissionTestContext() + c.Set(common.RequestIdKey, "request-public") + c.Set(pluginruntime.ContextKeyPinnedPlugin, pluginruntime.PinnedPlugin{ + Generation: &pluginruntime.RoutingGeneration{Number: 42}, + Plugin: &pluginruntime.LoadedPlugin{Meta: pluginruntime.Meta{ + Key: "document-parser", + Name: "Document Parser", + Version: "1.2.3", + APIVersion: 1, + Author: pluginruntime.AuthorMeta{ + Name: "Community Author", + URL: "https://plugins.example/author", + }, + }}, + }) + billing := &taskSubmissionTestBilling{events: &events} + info := taskSubmissionRelayInfo(billing) + + outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) { + return &relay.TaskSubmitResult{ + UpstreamTaskID: "upstream-private", + Platform: constant.TaskPlatform("document-parser"), + }, nil + }) + + require.Nil(t, taskErr) + require.NotNil(t, outcome) + require.NotNil(t, outcome.Task.PrivateData.Execution) + require.NotNil(t, outcome.Task.PrivateData.Execution.TaskPlugin) + assert.Equal(t, "request-public", outcome.Task.PrivateData.Execution.RequestID) + assert.Equal(t, "/plugin/submit", outcome.Task.PrivateData.Execution.RequestPath) + assert.Equal(t, "1.2.3", outcome.Task.PrivateData.Execution.TaskPlugin.Version) + assert.Equal(t, uint64(42), outcome.Task.PrivateData.Execution.TaskPlugin.Generation) + require.NotNil(t, outcome.Task.PrivateData.Execution.TaskPlugin.Author) + assert.Equal(t, "Community Author", outcome.Task.PrivateData.Execution.TaskPlugin.Author.Name) + assert.Equal(t, "https://plugins.example/author", outcome.Task.PrivateData.Execution.TaskPlugin.Author.URL) + + var stored model.Task + require.NoError(t, database.Where("task_id = ?", "task_public").First(&stored).Error) + require.NotNil(t, stored.PrivateData.Execution) + require.NotNil(t, stored.PrivateData.Execution.TaskPlugin) + assert.Equal(t, "document-parser", stored.PrivateData.Execution.TaskPlugin.Key) + require.NotNil(t, stored.PrivateData.Execution.TaskPlugin.Author) + assert.Equal(t, "Community Author", stored.PrivateData.Execution.TaskPlugin.Author.Name) + assert.Equal(t, "upstream-private", stored.PrivateData.UpstreamTaskID) +} + +func TestExecuteTaskSubmissionRefundsCancellationBeforeDurableBarrier(t *testing.T) { + events := make([]string, 0, 2) + setupTaskSubmissionDatabase(t, true, &events) + billing := &taskSubmissionTestBilling{events: &events} + c := taskSubmissionTestContext() + requestContext, cancel := context.WithCancel(c.Request.Context()) + c.Request = c.Request.WithContext(requestContext) + info := taskSubmissionRelayInfo(billing) + + outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) { + cancel() + return &relay.TaskSubmitResult{ + UpstreamTaskID: "upstream_private", + Platform: constant.TaskPlatform("plugin"), + }, nil + }) + + assert.Nil(t, outcome) + require.NotNil(t, taskErr) + assert.Equal(t, "request_cancelled", taskErr.Code) + assert.Equal(t, []string{"refund"}, events) + assert.Equal(t, 1, billing.refunds) + assert.False(t, c.Writer.Written()) +} + +func TestExecuteTaskSubmissionDisconnectBeforeUpstreamAcceptanceSkipsSubmitAndRefunds(t *testing.T) { + events := make([]string, 0, 1) + setupTaskSubmissionDatabase(t, true, &events) + billing := &taskSubmissionTestBilling{events: &events} + c := taskSubmissionTestContext() + requestContext, cancel := context.WithCancel(c.Request.Context()) + cancel() + c.Request = c.Request.WithContext(requestContext) + info := taskSubmissionRelayInfo(billing) + submitted := false + + outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) { + submitted = true + return nil, nil + }) + + assert.Nil(t, outcome) + require.NotNil(t, taskErr) + assert.Equal(t, "request_cancelled", taskErr.Code) + assert.False(t, submitted) + assert.Equal(t, []string{"refund"}, events) + assert.Equal(t, 1, billing.refunds) + assert.False(t, c.Writer.Written()) +} + +func TestExecuteTaskSubmissionCallerCancellationDuringSubmitRefundsBeforeDurableBarrier(t *testing.T) { + events := make([]string, 0, 1) + setupTaskSubmissionDatabase(t, true, &events) + billing := &taskSubmissionTestBilling{events: &events} + c := taskSubmissionTestContext() + requestContext, cancel := context.WithCancel(c.Request.Context()) + c.Request = c.Request.WithContext(requestContext) + info := taskSubmissionRelayInfo(billing) + submitStarted := make(chan struct{}) + done := make(chan struct{}) + var outcome *taskSubmissionOutcome + var taskErr *dto.TaskError + + go func() { + defer close(done) + outcome, taskErr = executeTaskSubmissionWith(c, info, func(c *gin.Context, _ *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) { + close(submitStarted) + <-c.Request.Context().Done() + return nil, service.TaskErrorWrapperLocal(c.Request.Context().Err(), "do_request_failed", http.StatusInternalServerError) + }) + }() + select { + case <-submitStarted: + case <-time.After(2 * time.Second): + require.FailNow(t, "submission did not start") + } + cancel() + select { + case <-done: + case <-time.After(2 * time.Second): + require.FailNow(t, "submission did not stop after disconnect") + } + + assert.Nil(t, outcome) + require.NotNil(t, taskErr) + assert.Equal(t, "request_cancelled", taskErr.Code) + assert.Equal(t, []string{"refund"}, events) + assert.Equal(t, 1, billing.refunds) + assert.False(t, c.Writer.Written()) +} + +func TestExecuteTaskSubmissionDisconnectAfterDurableInsertDoesNotRefund(t *testing.T) { + events := make([]string, 0, 3) + database := setupTaskSubmissionDatabase(t, true, &events) + previousLogConsumeEnabled := common.LogConsumeEnabled + common.LogConsumeEnabled = false + t.Cleanup(func() { common.LogConsumeEnabled = previousLogConsumeEnabled }) + c := taskSubmissionTestContext() + requestContext, cancel := context.WithCancel(c.Request.Context()) + c.Request = c.Request.WithContext(requestContext) + billing := &taskSubmissionTestBilling{ + events: &events, + onSettle: cancel, + } + info := taskSubmissionRelayInfo(billing) + + outcome, taskErr := executeTaskSubmissionWith(c, info, func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *dto.TaskError) { + return &relay.TaskSubmitResult{ + UpstreamTaskID: "upstream_private", + Platform: constant.TaskPlatform("plugin"), + }, nil + }) + + require.Nil(t, taskErr) + require.NotNil(t, outcome) + assert.Equal(t, "task_public", outcome.Task.TaskID) + assert.Equal(t, []string{"reserve", "insert", "settle"}, events) + assert.Zero(t, billing.refunds) + var count int64 + require.NoError(t, database.Model(&model.Task{}).Where("task_id = ?", "task_public").Count(&count).Error) + assert.Equal(t, int64(1), count) + assert.False(t, c.Writer.Written()) +} + +func setupTaskSubmissionDatabase(t *testing.T, migrate bool, events *[]string) *gorm.DB { + t.Helper() + previousDB := model.DB + database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, database.Callback().Create().Before("gorm:create").Register("test:task-submit-order", func(*gorm.DB) { + *events = append(*events, "insert") + })) + if migrate { + require.NoError(t, database.AutoMigrate(&model.Task{})) + } + model.DB = database + t.Cleanup(func() { model.DB = previousDB }) + return database +} + +func taskSubmissionTestContext() *gin.Context { + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodPost, "/plugin/submit", strings.NewReader(`{}`)) + return c +} + +func taskSubmissionRelayInfo(billing relaycommon.BillingSettler) *relaycommon.RelayInfo { + return &relaycommon.RelayInfo{ + UserId: 1, + UsingGroup: "default", + OriginModelName: "plugin-model", + Billing: billing, + TaskRelayInfo: &relaycommon.TaskRelayInfo{ + PublicTaskID: "task_public", + LockedChannel: &model.Channel{Id: 1, Type: constant.ChannelTypeTaskPlugin, Name: "plugin"}, + }, + ChannelMeta: &relaycommon.ChannelMeta{ChannelId: 1, ChannelType: constant.ChannelTypeTaskPlugin}, + } +} diff --git a/controller/revalidated_response.go b/controller/revalidated_response.go new file mode 100644 index 000000000000..f17a2dfd559b --- /dev/null +++ b/controller/revalidated_response.go @@ -0,0 +1,54 @@ +package controller + +import ( + "net/http" + + "github.com/QuantumNous/new-api/common" + "github.com/gin-gonic/gin" +) + +// etagVersionPublicContent namespaces the public-content ETag; bump it when +// the JSON envelope served by serveRevalidatedJSON changes shape. +const etagVersionPublicContent = "public-content:v1" + +type publicContentResponse struct { + Success bool `json:"success"` + Message string `json:"message"` + Data string `json:"data"` +} + +// serveRevalidatedJSON writes public content as JSON with a weak +// content-derived ETag and answers conditional requests with 304 Not +// Modified. The ETag is a weak validator derived from the content, so it is +// stable across replicas and JSON encodings, and a new one is issued when +// the content changes. Cache-Control: no-cache forces revalidation before +// reuse, so an admin edit takes effect on the next request; Vary: +// Accept-Encoding keeps the gzip and identity encodings apart in shared +// caches. +func serveRevalidatedJSON(c *gin.Context, content string) { + body, err := common.Marshal(publicContentResponse{ + Success: true, + Message: "", + Data: content, + }) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + + etag := common.ETagFor(etagVersionPublicContent, content) + + c.Header("ETag", etag) + c.Header("Cache-Control", "no-cache") + c.Header("Vary", "Accept-Encoding") + + if common.ETagMatches(c.GetHeader("If-None-Match"), etag) { + c.Status(http.StatusNotModified) + return + } + + c.Data(http.StatusOK, "application/json; charset=utf-8", body) +} diff --git a/controller/task.go b/controller/task.go index a80f1a687aab..c514bdb325b5 100644 --- a/controller/task.go +++ b/controller/task.go @@ -1,88 +1,478 @@ package controller import ( + "errors" + "fmt" + "net/http" + "regexp" "strconv" + "strings" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/middleware" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/relay" + relaychannel "github.com/QuantumNous/new-api/relay/channel" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/types" - "github.com/gin-gonic/gin" ) +type taskArtifactResponse struct { + Key string `json:"key"` + Type string `json:"type"` + MimeType string `json:"mime_type,omitempty"` + ContentURL string `json:"content_url"` +} + +var ( + taskArtifactKeyPattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._~-]{0,127}$`) + errTaskArtifactPluginUnavailable = errors.New("task artifact plugin unavailable") + errTaskArtifactPlugin = errors.New("task artifact plugin error") +) + +func GetTask(c *gin.Context) { + task, exists, err := model.GetByTaskId(c.GetInt("id"), c.Param("key")) + if err != nil { + videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to query task") + return + } + if !exists { + videoProxyError(c, http.StatusNotFound, "invalid_request_error", "Task not found") + return + } + createdAt := task.CreatedAt + if createdAt == 0 { + createdAt = task.SubmitTime + } + failReason := task.FailReason + if task.Status == model.TaskStatusSuccess && taskFailReasonIsLegacyResultURL(task.FailReason) { + failReason = "" + } + c.JSON(http.StatusOK, gin.H{ + "task_id": task.TaskID, + "platform": task.Platform, + "status": task.Status, + "progress": task.Progress, + "fail_reason": failReason, + "created_at": createdAt, + "finished_at": task.FinishTime, + }) +} + +func GetTaskArtifacts(c *gin.Context) { + task, exists, err := model.GetByTaskId(c.GetInt("id"), c.Param("key")) + if err != nil { + writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_internal_error", "Failed to query task") + return + } + if !exists || task == nil { + writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found") + return + } + writeTaskArtifacts(c, task, false) +} + +func GetDashboardTaskArtifacts(c *gin.Context) { + task, exists, err := getTaskForArtifactRequest(c, c.Param("task_id")) + if err != nil { + writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_internal_error", "Failed to query task") + return + } + if !exists || task == nil { + writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found") + return + } + writeTaskArtifacts(c, task, true) +} + +func writeTaskArtifacts(c *gin.Context, task *model.Task, dashboard bool) { + c.Header("Cache-Control", "private, no-store") + artifacts, err := projectTaskArtifacts(task) + if err != nil { + writeTaskArtifactProjectionError(c, err) + return + } + items := make([]taskArtifactResponse, 0, len(artifacts)) + for _, artifact := range artifacts { + contentURL, buildErr := service.BuildTaskArtifactContentURL(task.TaskID, artifact.Key) + if buildErr != nil { + writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_url_error", "Failed to build artifact content URL") + return + } + items = append(items, taskArtifactResponse{ + Key: artifact.Key, + Type: artifact.Type, + MimeType: artifact.MimeType, + ContentURL: contentURL, + }) + } + response := gin.H{"task_id": task.TaskID, "artifacts": items} + if legacyVideoAvailable(task) { + legacyContentURL, buildErr := service.BuildTaskArtifactContentURL(task.TaskID, "video") + if buildErr != nil { + writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_url_error", "Failed to build artifact content URL") + return + } + response["legacy_content_url"] = legacyContentURL + } + if dashboard { + common.ApiSuccess(c, response) + return + } + c.JSON(http.StatusOK, response) +} + +func projectTaskArtifacts(task *model.Task) ([]relaychannel.TaskArtifact, error) { + if task == nil || task.Status != model.TaskStatusSuccess || !taskHasPluginExecution(task) { + return []relaychannel.TaskArtifact{}, nil + } + adaptor := relay.GetTaskAdaptor(task.Platform) + if adaptor == nil { + return nil, errTaskArtifactPluginUnavailable + } + provider, ok := adaptor.(relaychannel.TaskArtifactProvider) + if !ok { + return []relaychannel.TaskArtifact{}, nil + } + artifacts, err := provider.ListArtifacts(task) + if err != nil { + return nil, fmt.Errorf("%w: %v", errTaskArtifactPlugin, err) + } + return validateProjectedTaskArtifacts(artifacts) +} + +func validateProjectedTaskArtifacts(artifacts []relaychannel.TaskArtifact) ([]relaychannel.TaskArtifact, error) { + if len(artifacts) > 64 { + return nil, fmt.Errorf("%w: too many artifacts", errTaskArtifactPlugin) + } + seen := make(map[string]struct{}, len(artifacts)) + for i := range artifacts { + if artifacts[i].Key != strings.TrimSpace(artifacts[i].Key) || + artifacts[i].Type != strings.TrimSpace(artifacts[i].Type) { + return nil, fmt.Errorf("%w: invalid artifact identity", errTaskArtifactPlugin) + } + if !taskArtifactKeyPattern.MatchString(artifacts[i].Key) { + return nil, fmt.Errorf("%w: invalid artifact key", errTaskArtifactPlugin) + } + if _, exists := seen[artifacts[i].Key]; exists { + return nil, fmt.Errorf("%w: duplicate artifact key", errTaskArtifactPlugin) + } + seen[artifacts[i].Key] = struct{}{} + switch artifacts[i].Type { + case "video", "audio", "image", "file": + default: + return nil, fmt.Errorf("%w: invalid artifact type", errTaskArtifactPlugin) + } + if len(artifacts[i].MimeType) > 255 || strings.ContainsAny(artifacts[i].MimeType, "\r\n") { + return nil, fmt.Errorf("%w: invalid artifact mime type", errTaskArtifactPlugin) + } + } + return artifacts, nil +} + +func initTaskArtifactAdaptor(task *model.Task) (relaychannel.TaskAdaptor, error) { + if task == nil || !taskHasPluginExecution(task) { + return nil, errTaskArtifactPluginUnavailable + } + channelModel, err := model.CacheGetChannel(task.ChannelId) + if err != nil { + return nil, fmt.Errorf("%w: channel unavailable", errTaskArtifactPluginUnavailable) + } + adaptor := relay.GetTaskAdaptor(task.Platform) + if adaptor == nil { + return nil, errTaskArtifactPluginUnavailable + } + pluginKey := task.PrivateData.Key + if pluginKey == "" { + pluginKey = channelModel.Key + } + baseURL := channelModel.GetBaseURL() + if baseURL == "" { + baseURL = constant.GetChannelBaseURL(channelModel.Type) + } + adaptor.Init(&relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelType: channelModel.Type, + ChannelBaseUrl: baseURL, + ApiKey: pluginKey, + ChannelSetting: channelModel.GetSetting(), + }, + }) + return adaptor, nil +} + +func taskHasPluginExecution(task *model.Task) bool { + return task != nil && + task.PrivateData.Execution != nil && + task.PrivateData.Execution.TaskPlugin != nil && + strings.TrimSpace(task.PrivateData.Execution.TaskPlugin.Key) != "" +} + +func legacyVideoAvailable(task *model.Task) bool { + if task == nil || task.Status != model.TaskStatusSuccess || + taskHasPluginExecution(task) || task.Platform == constant.TaskPlatformSuno || + strings.TrimSpace(task.GetResultURL()) == "" { + return false + } + switch constant.NormalizeTaskAction(task.Action) { + case constant.TaskActionImageToVideo, + constant.TaskActionTextToVideo, + constant.TaskActionFirstTailToVideo, + constant.TaskActionReferenceToVideo, + constant.TaskActionRemix: + return true + default: + return false + } +} + +func getTaskForArtifactRequest(c *gin.Context, taskID string) (*model.Task, bool, error) { + if middleware.IsTaskArtifactAccess(c) { + task, exists, err := model.GetUniqueByOnlyTaskId(taskID) + if err != nil || !exists || task == nil { + return task, exists, err + } + owner, err := model.GetUserCache(task.UserId) + if err != nil || owner == nil || owner.Status != common.UserStatusEnabled { + return nil, false, err + } + return task, true, nil + } + if c.GetInt("token_id") == 0 && c.GetInt("role") >= common.RoleAdminUser { + return model.GetByOnlyTaskId(taskID) + } + return model.GetByTaskId(c.GetInt("id"), taskID) +} + +func writeTaskArtifactProjectionError(c *gin.Context, err error) { + if errors.Is(err, errTaskArtifactPluginUnavailable) { + writeTaskArtifactError(c, http.StatusServiceUnavailable, "artifact_plugin_unavailable", "Artifact preview plugin is unavailable") + return + } + writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_plugin_error", "Artifact preview plugin failed") +} + +func writeTaskArtifactError(c *gin.Context, status int, code, message string) { + c.Header("Cache-Control", "private, no-store") + if middleware.IsTaskArtifactAccess(c) { + status = http.StatusNotFound + code = "artifact_not_found" + message = "Task or artifact not found" + } + if strings.HasPrefix(c.Request.URL.Path, "/api/") { + c.JSON(status, gin.H{"success": false, "code": code, "message": message}) + return + } + c.JSON(status, gin.H{ + "error": gin.H{ + "message": message, + "type": code, + "code": code, + }, + }) +} + +func TaskArtifactContent(c *gin.Context) { + task, exists, err := getTaskForArtifactRequest(c, c.Param("key")) + if err != nil { + writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_internal_error", "Failed to query task") + return + } + if !exists || task == nil { + writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found") + return + } + artifactKey := strings.TrimSpace(c.Param("artifact_key")) + if !taskArtifactKeyPattern.MatchString(artifactKey) { + writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found") + return + } + if task.Status != model.TaskStatusSuccess { + writeTaskArtifactError(c, http.StatusConflict, "artifact_not_ready", "Task artifacts are not ready") + return + } + if !taskHasPluginExecution(task) { + if artifactKey != "video" || !legacyVideoAvailable(task) { + writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found") + return + } + descriptor := &relaychannel.TaskContentRequest{ + URL: task.GetResultURL(), + Method: c.Request.Method, + Credentialless: true, + } + if err := proxyTaskMedia(c, task, descriptor); err != nil { + writeTaskMediaProxyError(c, err) + } + return + } + artifacts, err := projectTaskArtifacts(task) + if err != nil { + writeTaskArtifactProjectionError(c, err) + return + } + found := false + for _, artifact := range artifacts { + if artifact.Key == artifactKey { + found = true + break + } + } + if !found { + writeTaskArtifactError(c, http.StatusNotFound, "artifact_not_found", "Task or artifact not found") + return + } + artifactStore := service.GetTaskArtifactStore() + if ref, resolveErr := artifactStore.Resolve(task, artifactKey); resolveErr == nil && ref != nil { + _ = artifactStore.Serve(c, task, ref) + return + } + + adaptor, err := initTaskArtifactAdaptor(task) + if err != nil { + writeTaskArtifactProjectionError(c, err) + return + } + provider, ok := adaptor.(relaychannel.TaskContentRequestProvider) + if !ok { + writeTaskArtifactError(c, http.StatusServiceUnavailable, "artifact_plugin_unavailable", "Artifact content plugin is unavailable") + return + } + clientRequest := relaychannel.TaskArtifactClientRequest{ + Method: c.Request.Method, + Headers: taskArtifactClientHeaders(c.Request.Header), + } + descriptor, err := provider.BuildContentRequest(task, artifactKey, clientRequest) + if err != nil || descriptor == nil { + writeTaskArtifactError(c, http.StatusInternalServerError, "artifact_plugin_error", "Artifact content plugin failed") + return + } + if err := proxyTaskMedia(c, task, descriptor); err != nil { + writeTaskMediaProxyError(c, err) + } +} + +func taskArtifactClientHeaders(headers http.Header) map[string]string { + result := make(map[string]string, 4) + for _, name := range []string{"Range", "If-Range", "If-None-Match", "If-Modified-Since"} { + if value := strings.TrimSpace(headers.Get(name)); value != "" { + result[name] = value + } + } + return result +} + +/* + The task list handlers below deliberately do not call projectTaskArtifacts. + Artifact projection is confined to the explicit endpoints above. +*/ + func GetAllTask(c *gin.Context) { pageInfo := common.GetPageQuery(c) - startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64) endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64) - // 解析其他查询参数 - queryParams := model.SyncTaskQueryParams{ - Platform: constant.TaskPlatform(c.Query("platform")), - TaskID: c.Query("task_id"), - Status: c.Query("status"), - Action: c.Query("action"), - StartTimestamp: startTimestamp, - EndTimestamp: endTimestamp, - ChannelID: c.Query("channel_id"), - } - + queryParams := model.SyncTaskQueryParams{Platform: constant.TaskPlatform(c.Query("platform")), TaskID: c.Query("task_id"), Status: c.Query("status"), Action: c.Query("action"), StartTimestamp: startTimestamp, EndTimestamp: endTimestamp, ChannelID: c.Query("channel_id")} items := model.TaskGetAllTasks(pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams) - total := model.TaskCountAllTasks(queryParams) - pageInfo.SetTotal(int(total)) - pageInfo.SetItems(tasksToDto(items, true)) + pageInfo.SetTotal(int(model.TaskCountAllTasks(queryParams))) + pageInfo.SetItems(tasksToDto(items, true, c.GetInt("role"))) common.ApiSuccess(c, pageInfo) } func GetUserTask(c *gin.Context) { pageInfo := common.GetPageQuery(c) - - userId := c.GetInt("id") - + userID := c.GetInt("id") startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64) endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64) - - queryParams := model.SyncTaskQueryParams{ - Platform: constant.TaskPlatform(c.Query("platform")), - TaskID: c.Query("task_id"), - Status: c.Query("status"), - Action: c.Query("action"), - StartTimestamp: startTimestamp, - EndTimestamp: endTimestamp, - } - - items := model.TaskGetAllUserTask(userId, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams) - total := model.TaskCountAllUserTask(userId, queryParams) - pageInfo.SetTotal(int(total)) - pageInfo.SetItems(tasksToDto(items, false)) + queryParams := model.SyncTaskQueryParams{Platform: constant.TaskPlatform(c.Query("platform")), TaskID: c.Query("task_id"), Status: c.Query("status"), Action: c.Query("action"), StartTimestamp: startTimestamp, EndTimestamp: endTimestamp} + items := model.TaskGetAllUserTask(userID, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams) + pageInfo.SetTotal(int(model.TaskCountAllUserTask(userID, queryParams))) + pageInfo.SetItems(tasksToDto(items, false, common.RoleCommonUser)) common.ApiSuccess(c, pageInfo) } -func tasksToDto(tasks []*model.Task, fillUser bool) []*dto.TaskDto { - var userIdMap map[int]*model.UserBase +func tasksToDto(tasks []*model.Task, fillUser bool, viewerRole int) []*dto.TaskDto { + var userIDMap map[int]*model.UserBase if fillUser { - userIdMap = make(map[int]*model.UserBase) - userIds := types.NewSet[int]() + userIDMap = make(map[int]*model.UserBase) + userIDs := types.NewSet[int]() for _, task := range tasks { - userIds.Add(task.UserId) + userIDs.Add(task.UserId) } - for _, userId := range userIds.Items() { - cacheUser, err := model.GetUserCache(userId) - if err == nil { - userIdMap[userId] = cacheUser + for _, userID := range userIDs.Items() { + if cacheUser, err := model.GetUserCache(userID); err == nil { + userIDMap[userID] = cacheUser } } } result := make([]*dto.TaskDto, len(tasks)) for i, task := range tasks { if fillUser { - if user, ok := userIdMap[task.UserId]; ok { + if user, ok := userIDMap[task.UserId]; ok { task.Username = user.Username } } - result[i] = relay.TaskModel2Dto(task) + item := relay.TaskModel2Dto(task) + item.LegacyVideoAvailable = legacyVideoAvailable(task) + if task.Status == model.TaskStatusSuccess { + item.ResultURL = "" + if taskFailReasonIsLegacyResultURL(task.FailReason) { + item.FailReason = "" + } + } + if viewerRole >= common.RoleAdminUser { + adminInfo := &dto.TaskAdminInfo{} + if execution := task.PrivateData.Execution; execution != nil { + adminInfo.RequestID = execution.RequestID + adminInfo.RequestPath = execution.RequestPath + if snapshot := execution.TaskPlugin; snapshot != nil { + adminInfo.TaskPlugin = &dto.TaskPluginInfo{ + Key: snapshot.Key, + Name: snapshot.Name, + Version: snapshot.Version, + } + if snapshot.Author != nil { + adminInfo.TaskPlugin.Author = &dto.TaskPluginAuthorInfo{ + Name: snapshot.Author.Name, + URL: snapshot.Author.URL, + } + } + } + } + if adminInfo.RequestID != "" || adminInfo.RequestPath != "" || adminInfo.TaskPlugin != nil { + item.AdminInfo = adminInfo + } + } + if viewerRole >= common.RoleRootUser { + rootInfo := &dto.TaskRootInfo{ + UpstreamTaskID: task.PrivateData.UpstreamTaskID, + NodeName: task.PrivateData.NodeName, + } + if execution := task.PrivateData.Execution; execution != nil { + if snapshot := execution.TaskPlugin; snapshot != nil { + rootInfo.TaskPlugin = &dto.TaskPluginRuntimeInfo{ + Key: snapshot.Key, + Version: snapshot.Version, + APIVersion: snapshot.APIVersion, + Generation: snapshot.Generation, + } + } + } + if rootInfo.TaskPlugin != nil || rootInfo.UpstreamTaskID != "" || rootInfo.NodeName != "" { + item.RootInfo = rootInfo + } + } + result[i] = item } return result } + +func taskFailReasonIsLegacyResultURL(value string) bool { + value = strings.TrimSpace(value) + return len(value) >= len("https://") && strings.EqualFold(value[:len("https://")], "https://") || + len(value) >= len("http://") && strings.EqualFold(value[:len("http://")], "http://") || + len(value) >= len("data:") && strings.EqualFold(value[:len("data:")], "data:") +} diff --git a/controller/task_generic_test.go b/controller/task_generic_test.go new file mode 100644 index 000000000000..a3160b62ceb9 --- /dev/null +++ b/controller/task_generic_test.go @@ -0,0 +1,590 @@ +package controller + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "net/url" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/middleware" + "github.com/QuantumNous/new-api/model" + relaychannel "github.com/QuantumNous/new-api/relay/channel" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting/system_setting" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupGenericTaskTest(t *testing.T) *model.Task { + t.Helper() + originalDB := model.DB + previousRedisEnabled := common.RedisEnabled + common.RedisEnabled = false + database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, database.AutoMigrate(&model.Task{}, &model.Channel{}, &model.User{})) + model.DB = database + t.Cleanup(func() { + model.DB = originalDB + common.RedisEnabled = previousRedisEnabled + }) + + require.NoError(t, database.Create(&model.User{ + Id: 7, Username: "artifact-owner", Status: common.UserStatusEnabled, + Role: common.RoleCommonUser, Group: "default", + }).Error) + baseURL := "https://example.com" + require.NoError(t, database.Create(&model.Channel{ + Id: 1, Name: "artifact", Key: "key", BaseURL: &baseURL, Status: common.ChannelStatusEnabled, + }).Error) + task := &model.Task{ + TaskID: "task_generic", Platform: "document", UserId: 7, ChannelId: 1, + Status: model.TaskStatusSuccess, Progress: "100%", SubmitTime: 10, FinishTime: 20, + } + require.NoError(t, database.Create(task).Error) + return task +} + +func allowPrivateTaskMediaTest(t *testing.T) { + t.Helper() + originalFetchSetting := *system_setting.GetFetchSetting() + system_setting.GetFetchSetting().EnableSSRFProtection = true + system_setting.GetFetchSetting().AllowPrivateIp = true + system_setting.GetFetchSetting().AllowedPorts = []string{"1-65535"} + t.Cleanup(func() { *system_setting.GetFetchSetting() = originalFetchSetting }) + service.InitHttpClient() +} + +func TestGetTaskDoesNotProjectArtifacts(t *testing.T) { + task := setupGenericTaskTest(t) + task.FailReason = "https://stale-upstream.invalid/video.mp4" + task.PrivateData = model.TaskPrivateData{ + ResultURL: "https://private-upstream.invalid/video.mp4", + Execution: &model.TaskExecutionSnapshot{ + TaskPlugin: &model.TaskPluginSnapshot{Key: "missing-plugin", Name: "Missing"}, + }, + } + require.NoError(t, model.DB.Save(task).Error) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Set("id", 7) + c.Params = gin.Params{{Key: "key", Value: task.TaskID}} + c.Request = httptest.NewRequest(http.MethodGet, "/v1/tasks/"+task.TaskID, nil) + + GetTask(c) + + assert.Equal(t, http.StatusOK, recorder.Code) + var response map[string]any + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, task.TaskID, response["task_id"]) + assert.NotContains(t, response, "artifacts") + assert.NotContains(t, recorder.Body.String(), "upstream.invalid") +} + +func TestGetTaskArtifactsReturnsEmptyForLegacyTask(t *testing.T) { + task := setupGenericTaskTest(t) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Set("id", task.UserId) + c.Params = gin.Params{{Key: "key", Value: task.TaskID}} + c.Request = httptest.NewRequest(http.MethodGet, "/v1/tasks/"+task.TaskID+"/artifacts", nil) + + GetTaskArtifacts(c) + + assert.Equal(t, http.StatusOK, recorder.Code) + assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control")) + var response struct { + TaskID string `json:"task_id"` + Artifacts []taskArtifactResponse `json:"artifacts"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.Equal(t, task.TaskID, response.TaskID) + assert.Empty(t, response.Artifacts) +} + +func TestTaskArtifactAuthorizationKeepsForeignTasksHidden(t *testing.T) { + task := setupGenericTaskTest(t) + + commonUser, _ := gin.CreateTestContext(httptest.NewRecorder()) + commonUser.Set("id", 8) + commonUser.Set("role", common.RoleCommonUser) + _, exists, err := getTaskForArtifactRequest(commonUser, task.TaskID) + require.NoError(t, err) + assert.False(t, exists) + + admin, _ := gin.CreateTestContext(httptest.NewRecorder()) + admin.Set("id", 8) + admin.Set("role", common.RoleAdminUser) + found, exists, err := getTaskForArtifactRequest(admin, task.TaskID) + require.NoError(t, err) + require.True(t, exists) + assert.Equal(t, task.TaskID, found.TaskID) + + apiToken, _ := gin.CreateTestContext(httptest.NewRecorder()) + apiToken.Set("id", 8) + apiToken.Set("role", common.RoleRootUser) + apiToken.Set("token_id", 99) + _, exists, err = getTaskForArtifactRequest(apiToken, task.TaskID) + require.NoError(t, err) + assert.False(t, exists) +} + +func TestDashboardTaskArtifactsReturnsLegacyCapabilityWithoutUpstreamURL(t *testing.T) { + task := setupGenericTaskTest(t) + previousSecret := common.CryptoSecret + previousPublicAddress := system_setting.TaskPublicAddress + common.CryptoSecret = "controller-task-artifact-access-secret" + system_setting.TaskPublicAddress = "https://gateway.example/prefix" + t.Cleanup(func() { + common.CryptoSecret = previousSecret + system_setting.TaskPublicAddress = previousPublicAddress + }) + task.Action = constant.TaskActionTextToVideo + task.FailReason = "https://upstream.invalid/private-video.mp4?signature=secret" + require.NoError(t, model.DB.Save(task).Error) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Set("id", task.UserId) + c.Set("role", common.RoleCommonUser) + c.Params = gin.Params{{Key: "task_id", Value: task.TaskID}} + c.Request = httptest.NewRequest(http.MethodGet, "/api/task/"+task.TaskID+"/artifacts", nil) + + GetDashboardTaskArtifacts(c) + + assert.Equal(t, http.StatusOK, recorder.Code) + assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control")) + var response struct { + Success bool `json:"success"` + Data struct { + Artifacts []taskArtifactResponse `json:"artifacts"` + LegacyContentURL string `json:"legacy_content_url"` + } `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.True(t, response.Success) + assert.Empty(t, response.Data.Artifacts) + contentURL, err := url.Parse(response.Data.LegacyContentURL) + require.NoError(t, err) + assert.Equal(t, "/prefix/v1/tasks/"+task.TaskID+"/artifacts/video/content", contentURL.Path) + assert.True(t, service.VerifyTaskArtifactAccess( + contentURL.Query().Get(service.TaskArtifactAccessQueryParameter), + task.TaskID, + "video", + )) + assert.NotContains(t, recorder.Body.String(), "upstream.invalid") + assert.NotContains(t, recorder.Body.String(), "signature=secret") +} + +func TestTaskArtifactAccessRequiresActiveOwner(t *testing.T) { + task := setupGenericTaskTest(t) + task.Action = constant.TaskActionTextToVideo + task.FailReason = "https://upstream.invalid/private-video.mp4" + require.NoError(t, model.DB.Save(task).Error) + require.NoError(t, model.DB.Model(&model.User{}). + Where("id = ?", task.UserId). + Update("status", common.UserStatusDisabled).Error) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Set(middleware.TaskArtifactAccessContextKey, true) + c.Params = gin.Params{ + {Key: "key", Value: task.TaskID}, + {Key: "artifact_key", Value: "video"}, + } + c.Request = httptest.NewRequest( + http.MethodGet, + "/v1/tasks/"+task.TaskID+"/artifacts/video/content", + nil, + ) + + TaskArtifactContent(c) + + assert.Equal(t, http.StatusNotFound, recorder.Code) + assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control")) +} + +func TestTaskArtifactAccessRejectsAmbiguousHistoricalTaskID(t *testing.T) { + task := setupGenericTaskTest(t) + task.Action = constant.TaskActionTextToVideo + task.FailReason = "https://first-upstream.invalid/video.mp4" + require.NoError(t, model.DB.Save(task).Error) + require.NoError(t, model.DB.Create(&model.User{ + Id: 8, Username: "other-artifact-owner", Status: common.UserStatusEnabled, + Role: common.RoleCommonUser, Group: "default", AffCode: "artifact-owner-8", + }).Error) + require.NoError(t, model.DB.Create(&model.Task{ + TaskID: task.TaskID, Platform: task.Platform, UserId: 8, ChannelId: task.ChannelId, + Action: constant.TaskActionTextToVideo, Status: model.TaskStatusSuccess, + FailReason: "https://second-upstream.invalid/video.mp4", + }).Error) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Set(middleware.TaskArtifactAccessContextKey, true) + c.Params = gin.Params{ + {Key: "key", Value: task.TaskID}, + {Key: "artifact_key", Value: "video"}, + } + c.Request = httptest.NewRequest( + http.MethodGet, + "/v1/tasks/"+task.TaskID+"/artifacts/video/content", + nil, + ) + + TaskArtifactContent(c) + + assert.Equal(t, http.StatusNotFound, recorder.Code) + assert.NotContains(t, recorder.Body.String(), "upstream.invalid") +} + +func TestLegacyVideoArtifactContentUsesGetResultURL(t *testing.T) { + task := setupGenericTaskTest(t) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "bytes=0-3", r.Header.Get("Range")) + w.Header().Set("Content-Type", "video/mp4") + w.Header().Set("Content-Range", "bytes 0-3/4") + w.WriteHeader(http.StatusPartialContent) + _, _ = w.Write([]byte("data")) + })) + defer upstream.Close() + allowPrivateTaskMediaTest(t) + + task.Action = constant.TaskActionTextToVideo + task.PrivateData.ResultURL = upstream.URL + task.FailReason = "https://stale.invalid/legacy-fallback.mp4" + require.NoError(t, model.DB.Save(task).Error) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Set(middleware.TaskArtifactAccessContextKey, true) + c.Params = gin.Params{ + {Key: "key", Value: task.TaskID}, + {Key: "artifact_key", Value: "video"}, + } + c.Request = httptest.NewRequest( + http.MethodGet, + "/v1/tasks/"+task.TaskID+"/artifacts/video/content", + nil, + ) + c.Request.Header.Set("Range", "bytes=0-3") + + TaskArtifactContent(c) + + assert.Equal(t, http.StatusPartialContent, recorder.Code) + assert.Equal(t, "data", recorder.Body.String()) + assert.Equal(t, "bytes 0-3/4", recorder.Header().Get("Content-Range")) +} + +func TestDisabledArtifactStorePreservesPluginUpstreamContent(t *testing.T) { + task := setupGenericTaskTest(t) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, "provider-key", r.Header.Get("x-goog-api-key")) + assert.Equal(t, "bytes=0-13", r.Header.Get("Range")) + w.Header().Set("Content-Type", "video/mp4") + w.Header().Set("Content-Range", "bytes 0-13/14") + w.WriteHeader(http.StatusPartialContent) + _, _ = w.Write([]byte("artifact-bytes")) + })) + defer upstream.Close() + allowPrivateTaskMediaTest(t) + previousMemoryCache := common.MemoryCacheEnabled + common.MemoryCacheEnabled = false + t.Cleanup(func() { common.MemoryCacheEnabled = previousMemoryCache }) + + require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", task.ChannelId).Updates(map[string]any{ + "type": constant.ChannelTypeGemini, + "key": "provider-key", + "base_url": upstream.URL, + }).Error) + task.Platform = constant.TaskPlatform("google") + task.PrivateData.Execution = &model.TaskExecutionSnapshot{TaskPlugin: &model.TaskPluginSnapshot{ + Key: "google", Name: "Google Veo (Gemini API)", Version: "1.0.0", APIVersion: 1, + }} + task.SetData(map[string]any{"response": map[string]any{ + "generateVideoResponse": map[string]any{ + "generatedVideos": []any{map[string]any{"video": map[string]any{"uri": upstream.URL}}}, + }, + }}) + require.NoError(t, model.DB.Save(task).Error) + require.False(t, service.GetTaskArtifactStore().Enabled()) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Set(middleware.TaskArtifactAccessContextKey, true) + c.Params = gin.Params{ + {Key: "key", Value: task.TaskID}, + {Key: "artifact_key", Value: "video"}, + } + c.Request = httptest.NewRequest(http.MethodGet, "/v1/tasks/"+task.TaskID+"/artifacts/video/content", nil) + c.Request.Header.Set("Range", "bytes=0-13") + + TaskArtifactContent(c) + + assert.Equal(t, http.StatusPartialContent, recorder.Code) + assert.Equal(t, "artifact-bytes", recorder.Body.String()) + assert.Equal(t, "video/mp4", recorder.Header().Get("Content-Type")) + assert.Equal(t, "bytes 0-13/14", recorder.Header().Get("Content-Range")) +} + +func TestProjectedTaskArtifactValidationRejectsAmbiguousIdentity(t *testing.T) { + validated, err := validateProjectedTaskArtifacts([]relaychannel.TaskArtifact{ + {Key: "video-main", Type: "video", MimeType: "video/mp4"}, + {Key: "cover.main", Type: "image", MimeType: "image/png"}, + }) + require.NoError(t, err) + require.Len(t, validated, 2) + assert.Equal(t, "video-main", validated[0].Key) + + for _, artifacts := range [][]relaychannel.TaskArtifact{ + {{Key: "../video", Type: "video"}}, + {{Key: "video/0", Type: "video"}}, + {{Key: "video-main", Type: "video"}, {Key: "video-main", Type: "image"}}, + {{Key: "video-main", Type: "unknown"}}, + {{Key: "video-main", Type: "video", MimeType: "video/mp4\r\nX-Test: injected"}}, + } { + _, err := validateProjectedTaskArtifacts(artifacts) + assert.ErrorIs(t, err, errTaskArtifactPlugin) + } +} + +func TestProxyTaskMediaForwardsRangeAndFiltersResponseHeaders(t *testing.T) { + task := setupGenericTaskTest(t) + var receivedRange, receivedAuthorization string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + receivedRange = r.Header.Get("Range") + receivedAuthorization = r.Header.Get("Authorization") + w.Header().Set("Content-Type", "video/mp4") + w.Header().Set("Content-Range", "bytes 0-3/10") + w.Header().Set("Accept-Ranges", "bytes") + w.Header().Set("Set-Cookie", "provider=secret") + w.Header().Set("WWW-Authenticate", "Bearer provider") + w.Header().Set("X-Provider-Secret", "hidden") + w.Header().Set("Cache-Control", "public, max-age=86400") + w.WriteHeader(http.StatusPartialContent) + _, _ = w.Write([]byte("data")) + })) + defer upstream.Close() + + allowPrivateTaskMediaTest(t) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/tasks/task_generic/artifacts/video-main/content", nil) + c.Request.Header.Set("Range", "bytes=0-3") + + err := proxyTaskMedia(c, task, &relaychannel.TaskContentRequest{ + URL: upstream.URL, + Method: http.MethodGet, + Headers: map[string]string{"Authorization": "Bearer provider-secret"}, + }) + + require.NoError(t, err) + assert.Equal(t, http.StatusPartialContent, recorder.Code) + assert.Equal(t, "data", recorder.Body.String()) + assert.Equal(t, "bytes=0-3", receivedRange) + assert.Equal(t, "Bearer provider-secret", receivedAuthorization) + assert.Equal(t, "bytes 0-3/10", recorder.Header().Get("Content-Range")) + assert.Equal(t, "bytes", recorder.Header().Get("Accept-Ranges")) + assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control")) + assert.Equal(t, "sandbox; default-src 'none'", recorder.Header().Get("Content-Security-Policy")) + assert.Equal(t, "no-referrer", recorder.Header().Get("Referrer-Policy")) + assert.Equal(t, "nosniff", recorder.Header().Get("X-Content-Type-Options")) + assert.Empty(t, recorder.Header().Get("Set-Cookie")) + assert.Empty(t, recorder.Header().Get("WWW-Authenticate")) + assert.Empty(t, recorder.Header().Get("X-Provider-Secret")) +} + +func TestProxyTaskMediaPassesThroughUnsatisfiedRange(t *testing.T) { + task := setupGenericTaskTest(t) + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Range", "bytes */10") + w.WriteHeader(http.StatusRequestedRangeNotSatisfiable) + })) + defer upstream.Close() + + allowPrivateTaskMediaTest(t) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/content", nil) + + require.NoError(t, proxyTaskMedia(c, task, &relaychannel.TaskContentRequest{ + URL: upstream.URL, Method: http.MethodGet, + })) + assert.Equal(t, http.StatusRequestedRangeNotSatisfiable, recorder.Code) + assert.Equal(t, "bytes */10", recorder.Header().Get("Content-Range")) + assert.Equal(t, "private, no-store", recorder.Header().Get("Cache-Control")) +} + +func TestTaskMediaResponseHeaderTimeoutDoesNotTruncateBody(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "video/mp4") + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + time.Sleep(75 * time.Millisecond) + _, _ = w.Write([]byte("complete-body")) + })) + defer upstream.Close() + + request, err := http.NewRequest(http.MethodGet, upstream.URL, nil) + require.NoError(t, err) + response, err := doTaskMediaRequest(upstream.Client(), request, 20*time.Millisecond) + require.NoError(t, err) + defer response.Body.Close() + body, err := io.ReadAll(response.Body) + require.NoError(t, err) + assert.Equal(t, "complete-body", string(body)) +} + +func TestTaskMediaResponseHeaderTimeoutCancelsBeforeHeaders(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + time.Sleep(75 * time.Millisecond) + _, _ = w.Write([]byte("late")) + })) + defer upstream.Close() + + request, err := http.NewRequest(http.MethodGet, upstream.URL, nil) + require.NoError(t, err) + _, err = doTaskMediaRequest(upstream.Client(), request, 10*time.Millisecond) + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +func TestWriteVideoDataURLStreamsAndSupportsHead(t *testing.T) { + const dataURL = "data:video/mp4;base64,Y29tcGxldGUtYm9keQ==" + + getRecorder := httptest.NewRecorder() + getContext, _ := gin.CreateTestContext(getRecorder) + getContext.Request = httptest.NewRequest(http.MethodGet, "/content", nil) + require.NoError(t, writeVideoDataURL(getContext, dataURL)) + assert.Equal(t, http.StatusOK, getRecorder.Code) + assert.Equal(t, "complete-body", getRecorder.Body.String()) + assert.Equal(t, "13", getRecorder.Header().Get("Content-Length")) + + headRecorder := httptest.NewRecorder() + headContext, _ := gin.CreateTestContext(headRecorder) + headContext.Request = httptest.NewRequest(http.MethodHead, "/content", nil) + require.NoError(t, writeVideoDataURL(headContext, dataURL)) + assert.Equal(t, http.StatusOK, headRecorder.Code) + assert.Empty(t, headRecorder.Body.String()) + assert.Equal(t, "13", headRecorder.Header().Get("Content-Length")) +} + +func TestWriteVideoDataURLRejectsOversizedPayloadBeforeDecode(t *testing.T) { + previousLimit := taskMediaDataURLMaxEncodedBytes + taskMediaDataURLMaxEncodedBytes = 32 + t.Cleanup(func() { taskMediaDataURLMaxEncodedBytes = previousLimit }) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/content", nil) + + err := writeVideoDataURL(c, "data:video/mp4;base64,"+strings.Repeat("A", 64)) + + assert.ErrorIs(t, err, errTaskMediaRequestRejected) + assert.Empty(t, recorder.Header().Get("Content-Type")) +} + +func TestProxyTaskMediaAllowsOnlyCredentiallessCrossOriginRedirect(t *testing.T) { + task := setupGenericTaskTest(t) + var destinationAuthorization, destinationRange string + destination := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + destinationAuthorization = r.Header.Get("Authorization") + destinationRange = r.Header.Get("Range") + _, _ = w.Write([]byte("redirected")) + })) + defer destination.Close() + source := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + http.Redirect(w, r, destination.URL, http.StatusFound) + })) + defer source.Close() + allowPrivateTaskMediaTest(t) + + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/content", nil) + c.Request.Header.Set("Range", "bytes=0-3") + + err := proxyTaskMedia(c, task, &relaychannel.TaskContentRequest{ + URL: source.URL, Method: http.MethodGet, Credentialless: true, + }) + require.NoError(t, err) + assert.Equal(t, http.StatusOK, recorder.Code) + assert.Equal(t, "redirected", recorder.Body.String()) + assert.Empty(t, destinationAuthorization) + assert.Equal(t, "bytes=0-3", destinationRange) + + destinationRange = "" + rejectedRecorder := httptest.NewRecorder() + rejectedContext, _ := gin.CreateTestContext(rejectedRecorder) + rejectedContext.Request = httptest.NewRequest(http.MethodGet, "/content", nil) + err = proxyTaskMedia(rejectedContext, task, &relaychannel.TaskContentRequest{ + URL: source.URL, Method: http.MethodGet, + Headers: map[string]string{"Authorization": "Bearer provider-secret"}, + }) + var proxyErr *taskMediaProxyError + require.ErrorAs(t, err, &proxyErr) + assert.Equal(t, "artifact_request_rejected", proxyErr.code) + assert.Empty(t, destinationRange) +} + +func TestTaskMediaRequestHeaderPolicy(t *testing.T) { + header := http.Header{} + require.NoError(t, applyTaskMediaRequestHeaders(header, map[string]string{ + "Authorization": "Bearer provider-secret", + "X-Signature": "signed", + })) + assert.Equal(t, "Bearer provider-secret", header.Get("Authorization")) + assert.Equal(t, "signed", header.Get("X-Signature")) + + for _, name := range []string{"Host", "Content-Length", "Accept-Encoding", "Connection", "Proxy-Authorization", "Transfer-Encoding"} { + t.Run(name, func(t *testing.T) { + assert.ErrorIs(t, applyTaskMediaRequestHeaders(http.Header{}, map[string]string{name: "bad"}), errTaskMediaRequestRejected) + }) + } + assert.ErrorIs(t, applyTaskMediaRequestHeaders(http.Header{}, map[string]string{"X-Test": "bad\r\ninjected"}), errTaskMediaRequestRejected) +} + +func TestCredentiallessTaskMediaDescriptorRejectsCredentialsAndBody(t *testing.T) { + task := setupGenericTaskTest(t) + recorder := httptest.NewRecorder() + c, _ := gin.CreateTestContext(recorder) + c.Request = httptest.NewRequest(http.MethodGet, "/content", nil) + + for _, descriptor := range []*relaychannel.TaskContentRequest{ + {URL: "https://example.com/video", Method: http.MethodPost, Credentialless: true}, + {URL: "https://example.com/video", Method: http.MethodGet, Body: []byte("secret"), Credentialless: true}, + {URL: "https://example.com/video", Method: http.MethodGet, Headers: map[string]string{"X-Key": "secret"}, Credentialless: true}, + } { + err := proxyTaskMedia(c, task, descriptor) + var proxyErr *taskMediaProxyError + require.ErrorAs(t, err, &proxyErr) + assert.Equal(t, "artifact_request_rejected", proxyErr.code) + } +} + +func TestSelfTaskMediaURLGuard(t *testing.T) { + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodGet, "https://gateway.example/v1/videos/task-1/content", nil) + c.Request.Host = "gateway.example" + + selfURL, err := url.Parse("https://gateway.example/v1/videos/task-1/content") + require.NoError(t, err) + assert.True(t, isSelfTaskMediaURL(c, selfURL)) + + remoteURL, err := url.Parse("https://cdn.example/v1/videos/task-1/content") + require.NoError(t, err) + assert.False(t, isSelfTaskMediaURL(c, remoteURL)) + assert.True(t, isTaskMediaFallbackLoop(remoteURL.String(), "task-1")) + assert.False(t, isTaskMediaFallbackLoop(remoteURL.String(), "task-2")) +} diff --git a/controller/task_log_view_test.go b/controller/task_log_view_test.go new file mode 100644 index 000000000000..b9ab2480c2ea --- /dev/null +++ b/controller/task_log_view_test.go @@ -0,0 +1,137 @@ +package controller + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestTaskLogDTOSeparatesUserAdminAndRootDetails(t *testing.T) { + task := &model.Task{ + TaskID: "task_public", + Platform: "document-parser", + PrivateData: model.TaskPrivateData{ + Key: "channel-secret-canary", + UpstreamTaskID: "upstream-private", + NodeName: "node-a", + Execution: &model.TaskExecutionSnapshot{ + RequestID: "request-public", + RequestPath: "/v1/documents", + TaskPlugin: &model.TaskPluginSnapshot{ + Key: "document-parser", + Name: "Document Parser", + Version: "1.2.3", + Author: &model.TaskPluginAuthorSnapshot{ + Name: "Community Author", + URL: "https://plugins.example/author", + }, + APIVersion: 1, + Generation: 42, + }, + }, + }, + } + + userView := tasksToDto([]*model.Task{task}, false, common.RoleCommonUser)[0] + assert.Nil(t, userView.AdminInfo) + assert.Nil(t, userView.RootInfo) + + adminView := tasksToDto([]*model.Task{task}, false, common.RoleAdminUser)[0] + require.NotNil(t, adminView.AdminInfo) + require.NotNil(t, adminView.AdminInfo.TaskPlugin) + assert.Equal(t, "document-parser", adminView.AdminInfo.TaskPlugin.Key) + assert.Equal(t, "Document Parser", adminView.AdminInfo.TaskPlugin.Name) + assert.Equal(t, "1.2.3", adminView.AdminInfo.TaskPlugin.Version) + require.NotNil(t, adminView.AdminInfo.TaskPlugin.Author) + assert.Equal(t, "Community Author", adminView.AdminInfo.TaskPlugin.Author.Name) + assert.Equal(t, "https://plugins.example/author", adminView.AdminInfo.TaskPlugin.Author.URL) + assert.Equal(t, "request-public", adminView.AdminInfo.RequestID) + assert.Equal(t, "/v1/documents", adminView.AdminInfo.RequestPath) + assert.Nil(t, adminView.RootInfo) + + rootView := tasksToDto([]*model.Task{task}, false, common.RoleRootUser)[0] + require.NotNil(t, rootView.AdminInfo) + require.NotNil(t, rootView.RootInfo) + require.NotNil(t, rootView.RootInfo.TaskPlugin) + assert.Equal(t, 1, rootView.RootInfo.TaskPlugin.APIVersion) + assert.Equal(t, uint64(42), rootView.RootInfo.TaskPlugin.Generation) + assert.Equal(t, "upstream-private", rootView.RootInfo.UpstreamTaskID) + assert.Equal(t, "node-a", rootView.RootInfo.NodeName) + + adminJSON, err := common.Marshal(adminView) + require.NoError(t, err) + assert.NotContains(t, string(adminJSON), "channel-secret-canary") + assert.NotContains(t, string(adminJSON), "upstream-private") + + rootJSON, err := common.Marshal(rootView) + require.NoError(t, err) + assert.NotContains(t, string(rootJSON), "channel-secret-canary") + assert.Contains(t, string(rootJSON), "upstream-private") +} + +func TestTaskLogDTODoesNotInventHistoricalPluginProvenance(t *testing.T) { + task := &model.Task{ + TaskID: "task_without_snapshot", + Platform: "document-parser", + } + + adminView := tasksToDto([]*model.Task{task}, false, common.RoleAdminUser)[0] + + assert.Nil(t, adminView.AdminInfo) + assert.Nil(t, adminView.RootInfo) +} + +func TestTaskLogDTOReplacesLegacyVideoURLWithAvailabilityFlag(t *testing.T) { + task := &model.Task{ + TaskID: "task_legacy_video", + Platform: "jimeng", + Action: constant.TaskActionTextToVideo, + Status: model.TaskStatusSuccess, + FailReason: "https://private-upstream.invalid/video.mp4?signature=secret", + } + + view := tasksToDto([]*model.Task{task}, false, common.RoleCommonUser)[0] + assert.True(t, view.LegacyVideoAvailable) + assert.Empty(t, view.ResultURL) + assert.Empty(t, view.FailReason) + encoded, err := common.Marshal(view) + require.NoError(t, err) + assert.NotContains(t, string(encoded), "private-upstream.invalid") + assert.NotContains(t, string(encoded), "result_url") + assert.Contains(t, string(encoded), "legacy_video_available") +} + +func TestTaskLogDTOKeepsFailureReasonAndDoesNotMarkPluginTaskLegacy(t *testing.T) { + failed := &model.Task{ + TaskID: "task_failed", + Platform: "jimeng", + Action: constant.TaskActionTextToVideo, + Status: model.TaskStatusFailure, + FailReason: "provider rejected the request", + } + failedView := tasksToDto([]*model.Task{failed}, false, common.RoleCommonUser)[0] + assert.Equal(t, "provider rejected the request", failedView.FailReason) + assert.False(t, failedView.LegacyVideoAvailable) + + pluginTask := &model.Task{ + TaskID: "task_plugin_video", + Platform: "community-video", + Action: constant.TaskActionTextToVideo, + Status: model.TaskStatusSuccess, + FailReason: "https://stale-upstream.invalid/plugin-video.mp4", + PrivateData: model.TaskPrivateData{ + ResultURL: "https://private-upstream.invalid/plugin-video.mp4", + Execution: &model.TaskExecutionSnapshot{ + TaskPlugin: &model.TaskPluginSnapshot{Key: "community-video"}, + }, + }, + } + pluginView := tasksToDto([]*model.Task{pluginTask}, false, common.RoleCommonUser)[0] + assert.False(t, pluginView.LegacyVideoAvailable) + assert.Empty(t, pluginView.ResultURL) + assert.Empty(t, pluginView.FailReason) +} diff --git a/controller/task_plugin.go b/controller/task_plugin.go new file mode 100644 index 000000000000..d88f44114bef --- /dev/null +++ b/controller/task_plugin.go @@ -0,0 +1,761 @@ +package controller + +import ( + "context" + "crypto/sha256" + "encoding/json" + "errors" + "fmt" + "net/url" + "sort" + "strings" + "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/model" + "github.com/QuantumNous/new-api/pkg/jsplugin" + "github.com/QuantumNous/new-api/plugins" + "github.com/QuantumNous/new-api/setting" + "github.com/gin-gonic/gin" + "gorm.io/gorm" +) + +const maxTaskPluginSourceBytes = 1024 * 1024 + +type taskPluginUploadRequest struct { + Source string `json:"source" binding:"required"` + Enabled *bool `json:"enabled"` + Remark string `json:"remark"` + Force bool `json:"force"` + SourceSha256 string `json:"sourceSha256"` +} + +func UploadTaskPlugin(c *gin.Context) { + var request taskPluginUploadRequest + if err := c.ShouldBindJSON(&request); err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + if len(request.Source) > maxTaskPluginSourceBytes { + common.ApiErrorMsg(c, "plugin source exceeds 1 MiB") + return + } + if expected := strings.TrimSpace(request.SourceSha256); expected != "" { + actual := fmt.Sprintf("%x", sha256.Sum256([]byte(request.Source))) + if !strings.EqualFold(actual, expected) { + common.ApiErrorMsg(c, "plugin source sha256 mismatch") + return + } + } + temporary := jsplugin.NewRegistry() + loaded, err := temporary.Register(request.Source, jsplugin.Options{}) + if err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + if err = jsplugin.ValidateV1Meta(loaded.Meta); err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + enabled := true + if request.Enabled != nil { + enabled = *request.Enabled + } + if enabled && !request.Force { + if err = jsplugin.PreflightRoutingConflict(jsplugin.DefaultRegistry.Generation(), loaded); err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + } + plugin := model.TaskPlugin{ + Key: loaded.Meta.Key, APIVersion: loaded.Meta.APIVersion, Version: loaded.Meta.Version, + Source: request.Source, SourceHash: fmt.Sprintf("%x", sha256.Sum256([]byte(request.Source))), + Enabled: enabled, Remark: request.Remark, + } + if err = model.SaveTaskPlugin(&plugin); err != nil { + common.ApiError(c, err) + return + } + if err = syncTaskPluginsOnceContext(c.Request.Context()); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, taskPluginDetail{Plugin: &plugin, Meta: loaded.Meta, Source: plugin.Source, Layer: "override"}) +} + +func GetTaskPluginVersions(c *gin.Context) { + plugins, err := model.ListTaskPluginVersions(c.Param("key")) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, plugins) +} + +type taskPluginListItem struct { + Meta jsplugin.Meta `json:"meta"` + Source string `json:"source"` + Enabled bool `json:"enabled"` + Active bool `json:"active"` + SourceHash string `json:"source_hash"` + Remark string `json:"remark"` + RuntimeStatus string `json:"runtime_status"` + RuntimeError string `json:"runtime_error,omitempty"` + FactoryMeta *jsplugin.Meta `json:"factory_meta,omitempty"` + ChannelCount int `json:"channel_count"` + InFlightCount int64 `json:"in_flight_count"` +} + +type taskPluginRebuildOutcome struct { + Status string `json:"status"` + AttemptedAt time.Time `json:"attempted_at"` + Generation uint64 `json:"generation"` + DatabaseRevision string `json:"database_revision,omitempty"` + PluginErrorCount int `json:"plugin_error_count"` + Error string `json:"error,omitempty"` +} + +type taskPluginRuntimeStatus struct { + CurrentGeneration uint64 `json:"current_generation"` + GenerationPublishedAt time.Time `json:"generation_published_at"` + DatabaseRevision string `json:"database_revision"` + DatabaseError string `json:"database_error,omitempty"` + LastRebuild taskPluginRebuildOutcome `json:"last_rebuild"` + PluginErrors map[string]string `json:"plugin_errors"` +} + +func ListTaskPlugins(c *gin.Context) { + databasePlugins, err := model.ListTaskPlugins() + if err != nil { + common.ApiError(c, err) + return + } + snapshot := jsplugin.DefaultRegistry.Snapshot() + factory := make(map[string]jsplugin.Meta, len(snapshot.Factory)) + override := make(map[string]jsplugin.Meta, len(snapshot.Override)) + for _, meta := range snapshot.Factory { + factory[meta.Key] = meta + } + for _, meta := range snapshot.Override { + override[meta.Key] = meta + } + activeRows := make(map[string]model.TaskPlugin) + keys := make(map[string]struct{}, len(factory)+len(databasePlugins)) + for key := range factory { + keys[key] = struct{}{} + } + for _, plugin := range databasePlugins { + keys[plugin.Key] = struct{}{} + if plugin.Active { + activeRows[plugin.Key] = plugin + } + } + + runtimeErrors := jsplugin.DefaultRegistry.RoutingErrors() + taskPluginSyncState.Lock() + for key, message := range taskPluginSyncState.errors { + runtimeErrors[key] = message + } + taskPluginSyncState.Unlock() + + items := make([]taskPluginListItem, 0, len(keys)) + for key := range keys { + factoryMeta, hasFactory := factory[key] + row, hasOverride := activeRows[key] + item := taskPluginListItem{Enabled: true, Active: true, RuntimeStatus: "registered"} + if hasOverride { + item.Source = "override" + if hasFactory { + item.Source = "override_over_factory" + factoryCopy := factoryMeta + item.FactoryMeta = &factoryCopy + } + item.Meta = jsplugin.Meta{Key: row.Key, Version: row.Version, APIVersion: row.APIVersion} + if compiled, compileErr := jsplugin.NewRegistry().Register(row.Source, jsplugin.Options{Key: row.Key, Version: row.Version}); compileErr == nil { + item.Meta = compiled.Meta + } + item.Enabled = row.Enabled + item.Active = row.Active + item.SourceHash = row.SourceHash + item.Remark = row.Remark + if !constant.TaskPluginOverrideEnabled { + item.RuntimeStatus = "disabled_fallback" + } else if message := runtimeErrors[key]; message != "" { + item.RuntimeStatus = "compile_failed" + item.RuntimeError = message + } else if runtimeMeta, ok := override[key]; ok { + item.Meta = runtimeMeta + } else if !row.Enabled { + item.RuntimeStatus = "disabled_fallback" + } else { + item.RuntimeStatus = "not_registered" + } + } else { + item.Source = "factory" + item.Meta = factoryMeta + item.Enabled = !setting.IsTaskPluginFactoryDisabled(key) + source, sourceErr := plugins.Source(key) + if sourceErr == nil { + item.SourceHash = fmt.Sprintf("%x", sha256.Sum256([]byte(source))) + } + if !item.Enabled { + item.RuntimeStatus = "disabled" + } else if message := runtimeErrors[key]; message != "" { + item.RuntimeStatus = "compile_failed" + item.RuntimeError = message + } + } + if !hasFactory { + channels, inFlight, usageErr := model.GetTaskPluginUsage(key) + if usageErr != nil { + common.ApiError(c, usageErr) + return + } + item.ChannelCount = len(channels) + item.InFlightCount = inFlight + } + items = append(items, item) + } + sort.Slice(items, func(i, j int) bool { return items[i].Meta.Key < items[j].Meta.Key }) + common.ApiSuccess(c, items) +} + +func GetTaskPluginRuntime(c *gin.Context) { + routingStatus := jsplugin.DefaultRegistry.RoutingStatus() + pluginErrors := routingStatus.Errors + + taskPluginSyncState.Lock() + for key, message := range taskPluginSyncState.errors { + pluginErrors[key] = message + } + lastRebuild := taskPluginSyncState.lastRebuild + lastDatabaseRevision := lastRebuild.DatabaseRevision + taskPluginSyncState.Unlock() + + registryRebuild := routingStatus.LastRebuild + if lastRebuild.AttemptedAt.Before(registryRebuild.AttemptedAt) { + lastRebuild = taskPluginRebuildOutcome{ + Status: registryRebuild.Status, + AttemptedAt: registryRebuild.AttemptedAt, + Generation: registryRebuild.Generation, + Error: registryRebuild.Error, + } + } + if lastRebuild.Status == "" { + lastRebuild.Status = "never" + } + lastRebuild.PluginErrorCount = len(pluginErrors) + if lastRebuild.Status == "success" && len(pluginErrors) > 0 { + lastRebuild.Status = "partial" + } + + status := taskPluginRuntimeStatus{ + DatabaseRevision: lastDatabaseRevision, + LastRebuild: lastRebuild, + PluginErrors: pluginErrors, + } + databaseSnapshot, err := model.GetTaskPluginSyncSnapshot() + if err != nil { + status.DatabaseError = "database snapshot unavailable" + } else { + status.DatabaseRevision = databaseSnapshot.Revision + } + if routingStatus.Generation != nil { + status.CurrentGeneration = routingStatus.Generation.Number + status.GenerationPublishedAt = routingStatus.Generation.PublishedAt + } + common.ApiSuccess(c, status) +} + +type taskPluginDetail struct { + Plugin *model.TaskPlugin `json:"plugin,omitempty"` + Meta jsplugin.Meta `json:"meta"` + Source string `json:"source"` + Layer string `json:"layer"` +} + +func GetTaskPlugin(c *gin.Context) { + key := c.Param("key") + version := c.Query("version") + plugin, err := model.GetTaskPluginVersion(key, version) + 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()) + return + } + common.ApiSuccess(c, taskPluginDetail{Plugin: plugin, Meta: loaded.Meta, Source: plugin.Source, Layer: "override"}) + return + } + if !errors.Is(err, gorm.ErrRecordNotFound) || version != "" { + common.ApiError(c, err) + return + } + source, err := plugins.Source(key) + if err != nil { + common.ApiErrorMsg(c, "task plugin not found") + return + } + loaded, err := jsplugin.NewRegistry().RegisterFactory(source, jsplugin.Options{Key: key}) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, taskPluginDetail{Meta: loaded.Meta, Source: source, Layer: "factory"}) +} + +type taskPluginDryRunRequest struct { + Hook string `json:"hook" binding:"required"` + Member string `json:"member"` + Args []json.RawMessage `json:"args"` +} + +func DryRunTaskPlugin(c *gin.Context) { + var request taskPluginDryRunRequest + if err := c.ShouldBindJSON(&request); err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + detailSource := "" + plugin, err := model.GetTaskPluginVersion(c.Param("key"), "") + if err == nil { + detailSource = plugin.Source + } else if errors.Is(err, gorm.ErrRecordNotFound) { + detailSource, err = plugins.Source(c.Param("key")) + } + if err != nil { + common.ApiErrorMsg(c, "task plugin not found") + return + } + loaded, err := jsplugin.NewRegistry().Register(detailSource, jsplugin.Options{Key: c.Param("key")}) + if err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + args := make([]any, len(request.Args)) + for index, raw := range request.Args { + if err = common.Unmarshal(raw, &args[index]); err != nil { + common.ApiErrorMsg(c, fmt.Sprintf("invalid argument %d: %v", index+1, err)) + return + } + } + var output any + if request.Member == "" { + output, err = loaded.Engine.Call(context.Background(), request.Hook, args...) + } else { + output, err = loaded.Engine.CallMember(context.Background(), request.Hook, request.Member, args...) + } + if err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + common.ApiSuccess(c, output) +} + +func DeleteTaskPluginVersion(c *gin.Context) { + key := c.Param("key") + version := c.Param("version") + plugin, lookupErr := model.GetTaskPluginVersion(key, version) + if lookupErr != nil { + if errors.Is(lookupErr, gorm.ErrRecordNotFound) { + common.ApiErrorMsg(c, "override plugin version not found; factory plugins cannot be deleted") + return + } + common.ApiError(c, lookupErr) + return + } + if plugin.Active && !taskPluginHasFactory(key) { + channels, inFlight, usageErr := model.GetTaskPluginUsage(key) + if usageErr != nil { + common.ApiError(c, usageErr) + return + } + if (len(channels) > 0 || inFlight > 0) && c.Query("force") != "true" { + c.JSON(200, gin.H{"success": false, "message": "task plugin is still in use", "data": gin.H{"channels": channels, "in_flight_count": inFlight}}) + return + } + } + _, err := model.DeleteTaskPluginVersion(key, version) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + common.ApiErrorMsg(c, "override plugin version not found; factory plugins cannot be deleted") + return + } + common.ApiError(c, err) + return + } + if err = syncTaskPluginsOnceContext(c.Request.Context()); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, nil) +} + +type taskPluginActivateRequest struct { + Version string `json:"version" binding:"required"` +} + +func ActivateTaskPlugin(c *gin.Context) { + var request taskPluginActivateRequest + if err := c.ShouldBindJSON(&request); err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + versions, err := model.ListTaskPluginVersions(c.Param("key")) + if err != nil { + common.ApiError(c, err) + return + } + var target *model.TaskPlugin + for i := range versions { + if versions[i].Version == request.Version { + target = &versions[i] + break + } + } + if target == nil { + common.ApiErrorMsg(c, "plugin version not found") + return + } + if _, err = jsplugin.NewRegistry().Register(target.Source, jsplugin.Options{Key: target.Key, Version: target.Version}); err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + if err = model.ActivateTaskPlugin(target.Key, target.Version); err != nil { + common.ApiError(c, err) + return + } + if err = syncTaskPluginsOnceContext(c.Request.Context()); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, nil) +} + +type taskPluginStatusRequest struct { + Enabled *bool `json:"enabled" binding:"required"` +} + +func SetTaskPluginStatus(c *gin.Context) { + var request taskPluginStatusRequest + if err := c.ShouldBindJSON(&request); err != nil || request.Enabled == nil { + common.ApiErrorMsg(c, "enabled is required") + return + } + key := c.Param("key") + disabledChannels := 0 + if !*request.Enabled { + channels, inFlight, usageErr := model.GetTaskPluginUsage(key) + if usageErr != nil { + common.ApiError(c, usageErr) + return + } + cascade := c.Query("cascade") == "true" + force := c.Query("force") == "true" + if (len(channels) > 0 && !cascade) || (inFlight > 0 && !force) { + c.JSON(200, gin.H{"success": false, "message": "task plugin is still in use", "data": gin.H{"channels": channels, "in_flight_count": inFlight}}) + return + } + if cascade { + for _, channel := range channels { + if model.UpdateChannelStatus(channel.Id, "", common.ChannelStatusManuallyDisabled, "task plugin disabled") { + disabledChannels++ + } + } + } + } + _, lookupErr := model.GetTaskPluginVersion(key, "") + hasActiveOverride := lookupErr == nil + if lookupErr != nil && !errors.Is(lookupErr, gorm.ErrRecordNotFound) { + common.ApiError(c, lookupErr) + return + } + // The disabled set suppresses only the factory fallback layer. An enabled + // override for the same key keeps serving and is toggled independently. + if taskPluginHasFactory(key) && !hasActiveOverride { + keys := setting.GetTaskPluginDisabledFactoryKeys() + if *request.Enabled { + next := make([]string, 0, len(keys)) + for _, item := range keys { + if item != key { + next = append(next, item) + } + } + keys = next + } else { + keys = append(append([]string{}, keys...), key) + } + if err := setting.SetTaskPluginDisabledFactoryKeysOption(keys); err != nil { + common.ApiError(c, err) + return + } + encoded, err := common.Marshal(setting.GetTaskPluginDisabledFactoryKeys()) + if err != nil { + common.ApiError(c, err) + return + } + if err = model.UpdateOption(setting.TaskPluginDisabledFactoryKeysKey, string(encoded)); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"plugin_enabled": *request.Enabled, "disabled_channels": disabledChannels}) + return + } + if err := model.SetTaskPluginEnabled(key, *request.Enabled); err != nil { + common.ApiError(c, err) + return + } + if err := syncTaskPluginsOnceContext(c.Request.Context()); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"plugin_enabled": *request.Enabled, "disabled_channels": disabledChannels}) +} + +func taskPluginHasFactory(key string) bool { + for _, meta := range jsplugin.DefaultRegistry.Snapshot().Factory { + if meta.Key == key { + return true + } + } + return false +} + +func GetTaskPluginMarketplaceSources(c *gin.Context) { + common.ApiSuccess(c, setting.GetTaskPluginMarketplaceSources()) +} + +func UpdateTaskPluginMarketplaceSources(c *gin.Context) { + var sources []setting.TaskPluginMarketplaceSource + if err := c.ShouldBindJSON(&sources); err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + if sources == nil { + sources = []setting.TaskPluginMarketplaceSource{} + } + for i := range sources { + name := strings.TrimSpace(sources[i].Name) + indexURL := strings.TrimSpace(sources[i].IndexURL) + if name == "" { + common.ApiErrorMsg(c, "marketplace source name is required") + return + } + parsed, err := url.Parse(indexURL) + if err != nil || !parsed.IsAbs() || parsed.Host == "" || (!strings.EqualFold(parsed.Scheme, "http") && !strings.EqualFold(parsed.Scheme, "https")) { + common.ApiErrorMsg(c, "marketplace source index_url must be an absolute http(s) URL") + return + } + sources[i].Name = name + sources[i].IndexURL = indexURL + } + encoded, err := common.Marshal(sources) + if err != nil { + common.ApiError(c, err) + return + } + if err = model.UpdateOption(setting.TaskPluginMarketplaceSourcesKey, string(encoded)); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, sources) +} + +func GetTaskPluginOptions(c *gin.Context) { + snapshot := jsplugin.DefaultRegistry.Snapshot() + seen := make(map[string]bool) + options := make([]gin.H, 0, len(snapshot.Factory)+len(snapshot.Override)) + for layer, metas := range [][]jsplugin.Meta{snapshot.Override, snapshot.Factory} { + for _, meta := range metas { + if seen[meta.Key] { + continue + } + // Disabled factory keys are omitted from bind options. The disabled + // set suppresses only the factory fallback; an enabled override for + // the same key is listed in the override pass and still appears. + if layer == 1 && setting.IsTaskPluginFactoryDisabled(meta.Key) { + continue + } + if _, ok := jsplugin.DefaultRegistry.Get(meta.Key); !ok { + continue + } + seen[meta.Key] = true + options = append(options, gin.H{ + "key": meta.Key, + "name": meta.Name, + "models": meta.Models, + "usageSchema": meta.UsageSchema, + }) + } + } + sort.Slice(options, func(i, j int) bool { return options[i]["key"].(string) < options[j]["key"].(string) }) + common.ApiSuccess(c, options) +} + +var taskPluginSyncState = struct { + sync.Mutex + hashes map[string]string + errors map[string]string + lastRebuild taskPluginRebuildOutcome +}{hashes: map[string]string{}, errors: map[string]string{}} + +func syncTaskPluginsOnce() error { + return syncTaskPluginsOnceContext(context.Background()) +} + +func syncTaskPluginsOnceContext(ctx context.Context) error { + started := time.Now() + taskPluginSyncState.Lock() + defer taskPluginSyncState.Unlock() + databaseSnapshot, err := model.GetTaskPluginSyncSnapshot() + if err != nil { + syncErr := fmt.Errorf("sync task plugins: %w", err) + taskPluginSyncState.lastRebuild = taskPluginRebuildOutcome{ + Status: "failed", + AttemptedAt: time.Now(), + Generation: jsplugin.DefaultRegistry.Generation().Number, + DatabaseRevision: taskPluginSyncState.lastRebuild.DatabaseRevision, + Error: syncErr.Error(), + } + logger.LogDebug( + ctx, + "task_plugin subsystem=sync event=failed stage=database_snapshot retained_generation=%d elapsed_ms=%d", + jsplugin.DefaultRegistry.Generation().Number, + time.Since(started).Milliseconds(), + ) + return syncErr + } + databasePlugins := databaseSnapshot.Plugins + sort.Slice(databasePlugins, func(i, j int) bool { return databasePlugins[i].Key < databasePlugins[j].Key }) + currentOverrides := jsplugin.DefaultRegistry.OverridePlugins() + generationBefore := jsplugin.DefaultRegistry.Generation().Number + logger.LogDebug( + ctx, + "task_plugin subsystem=sync event=start database_revision=%q generation=%d desired_plugins=%d current_overrides=%d", + databaseSnapshot.Revision, + generationBefore, + len(databasePlugins), + len(currentOverrides), + ) + nextOverrides := make([]*jsplugin.LoadedPlugin, 0, len(databasePlugins)) + nextHashes := make(map[string]string, len(databasePlugins)) + seen := make(map[string]bool, len(databasePlugins)) + for _, plugin := range databasePlugins { + seen[plugin.Key] = true + if current := currentOverrides[plugin.Key]; current != nil && taskPluginSyncState.hashes[plugin.Key] == plugin.SourceHash { + nextOverrides = append(nextOverrides, current) + nextHashes[plugin.Key] = plugin.SourceHash + logger.LogDebug( + ctx, + "task_plugin subsystem=sync event=plugin plugin=%q version=%q action=reuse", + plugin.Key, + plugin.Version, + ) + continue + } + logger.LogDebug( + ctx, + "task_plugin subsystem=sync event=plugin plugin=%q version=%q action=compile_start", + plugin.Key, + plugin.Version, + ) + compiled, compileErr := jsplugin.CompilePlugin(plugin.Source, jsplugin.Options{Key: plugin.Key, Version: plugin.Version}) + if compileErr != nil { + retainedIncumbent := false + if current := currentOverrides[plugin.Key]; current != nil { + nextOverrides = append(nextOverrides, current) + retainedIncumbent = true + if currentHash := taskPluginSyncState.hashes[plugin.Key]; currentHash != "" { + nextHashes[plugin.Key] = currentHash + } + } + taskPluginSyncState.errors[plugin.Key] = compileErr.Error() + common.SysError(fmt.Sprintf("compile task plugin %s@%s: %v", plugin.Key, plugin.Version, compileErr)) + logger.LogDebug( + ctx, + "task_plugin subsystem=sync event=plugin plugin=%q version=%q action=compile_failed retained_incumbent=%t", + plugin.Key, + plugin.Version, + retainedIncumbent, + ) + continue + } + nextOverrides = append(nextOverrides, compiled) + nextHashes[plugin.Key] = plugin.SourceHash + delete(taskPluginSyncState.errors, plugin.Key) + logger.LogDebug( + ctx, + "task_plugin subsystem=sync event=plugin plugin=%q version=%q action=compile_success", + plugin.Key, + plugin.Version, + ) + } + if err = jsplugin.DefaultRegistry.ReplaceOverrides(nextOverrides); err != nil { + syncErr := fmt.Errorf("publish task plugin generation: %w", err) + taskPluginSyncState.lastRebuild = taskPluginRebuildOutcome{ + Status: "failed", + AttemptedAt: time.Now(), + Generation: jsplugin.DefaultRegistry.Generation().Number, + DatabaseRevision: databaseSnapshot.Revision, + Error: syncErr.Error(), + } + logger.LogDebug( + ctx, + "task_plugin subsystem=sync event=failed stage=publish retained_generation=%d retained_generation_active=true database_revision=%q elapsed_ms=%d", + jsplugin.DefaultRegistry.Generation().Number, + databaseSnapshot.Revision, + time.Since(started).Milliseconds(), + ) + return syncErr + } + taskPluginSyncState.hashes = nextHashes + for key := range taskPluginSyncState.errors { + if !seen[key] { + delete(taskPluginSyncState.errors, key) + } + } + pluginErrors := jsplugin.DefaultRegistry.RoutingErrors() + for key, message := range taskPluginSyncState.errors { + pluginErrors[key] = message + } + pluginErrorCount := len(pluginErrors) + status := "success" + if pluginErrorCount > 0 { + status = "partial" + } + taskPluginSyncState.lastRebuild = taskPluginRebuildOutcome{ + Status: status, + AttemptedAt: time.Now(), + Generation: jsplugin.DefaultRegistry.Generation().Number, + DatabaseRevision: databaseSnapshot.Revision, + PluginErrorCount: pluginErrorCount, + } + logger.LogDebug( + ctx, + "task_plugin subsystem=sync event=complete database_revision=%q previous_generation=%d generation=%d status=%q active_overrides=%d plugin_errors=%d elapsed_ms=%d", + databaseSnapshot.Revision, + generationBefore, + jsplugin.DefaultRegistry.Generation().Number, + status, + len(jsplugin.DefaultRegistry.ActiveOverridePlugins()), + pluginErrorCount, + time.Since(started).Milliseconds(), + ) + return nil +} + +func SyncTaskPluginsOnce() { + if err := syncTaskPluginsOnce(); err != nil { + common.SysError(err.Error()) + } +} + +func SyncTaskPlugins() { + SyncTaskPluginsOnce() + for range time.NewTicker(30 * time.Second).C { + SyncTaskPluginsOnce() + } +} diff --git a/controller/task_plugin_debug.go b/controller/task_plugin_debug.go new file mode 100644 index 000000000000..450f6b4b9972 --- /dev/null +++ b/controller/task_plugin_debug.go @@ -0,0 +1,269 @@ +package controller + +import ( + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin" + "github.com/QuantumNous/new-api/relay" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/gin-gonic/gin" +) + +// taskPluginSubmitDiagnostics keeps plugin-only lifecycle logging out of the +// ordinary task path. An empty plugin key makes every method a no-op. +type taskPluginSubmitDiagnostics struct { + context *gin.Context + pluginKey string + generation uint64 +} + +func newTaskPluginSubmitDiagnostics(c *gin.Context) taskPluginSubmitDiagnostics { + diagnostics := taskPluginSubmitDiagnostics{ + context: c, + pluginKey: c.GetString("expected_task_plugin_key"), + } + if diagnostics.pluginKey == "" { + return diagnostics + } + if pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedPlugin); exists { + if pinned, ok := pinnedValue.(pluginruntime.PinnedPlugin); ok && pinned.Generation != nil { + diagnostics.generation = pinned.Generation.Number + } + } + return diagnostics +} + +func (d taskPluginSubmitDiagnostics) start(info *relaycommon.RelayInfo) { + if d.pluginKey == "" { + return + } + logger.LogDebug( + d.context, + "task_plugin subsystem=submit event=start generation=%d plugin=%q model=%q action_present=%t", + d.generation, + d.pluginKey, + info.OriginModelName, + info.Action != "", + ) +} + +func (d taskPluginSubmitDiagnostics) refund(stage string) { + if d.pluginKey == "" { + return + } + logger.LogDebug( + d.context, + "task_plugin subsystem=submit event=refund_invoked generation=%d plugin=%q stage=%q durable=false", + d.generation, + d.pluginKey, + stage, + ) +} + +func (d taskPluginSubmitDiagnostics) cancelled(stage string, attempt int) { + if d.pluginKey == "" { + return + } + logger.LogDebug( + d.context, + "task_plugin subsystem=submit event=cancelled generation=%d plugin=%q stage=%q attempt=%d", + d.generation, + d.pluginKey, + stage, + attempt, + ) +} + +func (d taskPluginSubmitDiagnostics) attempt(attempt int, channel *model.Channel, locked bool) { + if d.pluginKey == "" || channel == nil { + return + } + logger.LogDebug( + d.context, + "task_plugin subsystem=submit event=attempt generation=%d plugin=%q attempt=%d channel_id=%d channel_type=%d locked=%t", + d.generation, + d.pluginKey, + attempt, + channel.Id, + channel.Type, + locked, + ) +} + +func (d taskPluginSubmitDiagnostics) attemptSucceeded(attempt int, result *relay.TaskSubmitResult) { + if d.pluginKey == "" || result == nil { + return + } + logger.LogDebug( + d.context, + "task_plugin subsystem=submit event=attempt_succeeded generation=%d plugin=%q attempt=%d platform=%q quota=%d task_data_bytes=%d client_response=%t immediate=%t", + d.generation, + d.pluginKey, + attempt, + result.Platform, + result.Quota, + len(result.TaskData), + result.ClientResponse != nil, + result.Immediate != nil, + ) +} + +func (d taskPluginSubmitDiagnostics) attemptFailed(attempt int, channel *model.Channel, taskErr *dto.TaskError, willRetry bool) { + if d.pluginKey == "" || channel == nil || taskErr == nil { + return + } + logger.LogDebug( + d.context, + "task_plugin subsystem=submit event=attempt_failed generation=%d plugin=%q attempt=%d channel_id=%d channel_type=%d code=%q status=%d local=%t will_retry=%t", + d.generation, + d.pluginKey, + attempt, + channel.Id, + channel.Type, + taskErr.Code, + taskErr.StatusCode, + taskErr.LocalError, + willRetry, + ) +} + +func (d taskPluginSubmitDiagnostics) failed(stage, reason string, taskErr *dto.TaskError, durable bool) { + if d.pluginKey == "" { + return + } + code := "" + status := 0 + local := true + if taskErr != nil { + code = taskErr.Code + status = taskErr.StatusCode + local = taskErr.LocalError + } + logger.LogDebug( + d.context, + "task_plugin subsystem=submit event=failed generation=%d plugin=%q stage=%q reason=%q code=%q status=%d local=%t durable=%t", + d.generation, + d.pluginKey, + stage, + reason, + code, + status, + local, + durable, + ) +} + +func (d taskPluginSubmitDiagnostics) reserve(event string, quota int) { + if d.pluginKey == "" { + return + } + logger.LogDebug( + d.context, + "task_plugin subsystem=submit event=%s generation=%d plugin=%q quota=%d", + event, + d.generation, + d.pluginKey, + quota, + ) +} + +func (d taskPluginSubmitDiagnostics) insertStart(task *model.Task) { + if d.pluginKey == "" || task == nil { + return + } + logger.LogDebug( + d.context, + "task_plugin subsystem=submit event=insert_start generation=%d plugin=%q public_task_id=%q platform=%q channel_id=%d quota=%d", + d.generation, + d.pluginKey, + task.TaskID, + task.Platform, + task.ChannelId, + task.Quota, + ) +} + +func (d taskPluginSubmitDiagnostics) durable(task *model.Task) { + if d.pluginKey == "" || task == nil { + return + } + logger.LogDebug( + d.context, + "task_plugin subsystem=submit event=durable generation=%d plugin=%q public_task_id=%q status=%q durable=true", + d.generation, + d.pluginKey, + task.TaskID, + taskPluginDebugStatus(string(task.Status)), + ) +} + +func (d taskPluginSubmitDiagnostics) settleStart(task *model.Task, quota int) { + if d.pluginKey == "" || task == nil { + return + } + logger.LogDebug( + d.context, + "task_plugin subsystem=submit event=settle_start generation=%d plugin=%q public_task_id=%q quota=%d durable=true", + d.generation, + d.pluginKey, + task.TaskID, + quota, + ) +} + +func (d taskPluginSubmitDiagnostics) complete(task *model.Task, quota int) { + if d.pluginKey == "" || task == nil { + return + } + logger.LogDebug( + d.context, + "task_plugin subsystem=submit event=complete generation=%d plugin=%q public_task_id=%q quota=%d durable=true", + d.generation, + d.pluginKey, + task.TaskID, + quota, + ) +} + +func (d taskPluginSubmitDiagnostics) present(task *model.Task, presenter string) { + if d.pluginKey == "" || task == nil { + return + } + logger.LogDebug( + d.context, + "task_plugin subsystem=submit event=present generation=%d plugin=%q public_task_id=%q presenter=%q durable=true", + d.generation, + d.pluginKey, + task.TaskID, + presenter, + ) +} + +func (d taskPluginSubmitDiagnostics) presentError(taskErr *dto.TaskError) { + if d.pluginKey == "" || taskErr == nil { + return + } + logger.LogDebug( + d.context, + "task_plugin subsystem=submit event=present_error generation=%d plugin=%q code=%q status=%d local=%t", + d.generation, + d.pluginKey, + taskErr.Code, + taskErr.StatusCode, + taskErr.LocalError, + ) +} + +func taskPluginDebugStatus(status string) string { + switch model.TaskStatus(status) { + case model.TaskStatusSubmitted, + model.TaskStatusQueued, + model.TaskStatusInProgress, + model.TaskStatusSuccess, + model.TaskStatusFailure: + return status + default: + return "unknown" + } +} diff --git a/controller/task_plugin_debug_test.go b/controller/task_plugin_debug_test.go new file mode 100644 index 000000000000..ff6cc56dd97a --- /dev/null +++ b/controller/task_plugin_debug_test.go @@ -0,0 +1,88 @@ +package controller + +import ( + "bytes" + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/relay" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestTaskPluginSubmitDiagnosticsArePluginOnlyAndDoNotLogPayloads(t *testing.T) { + previousDebug := common.DebugEnabled + common.DebugEnabled = true + t.Cleanup(func() { common.DebugEnabled = previousDebug }) + + 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() + }) + + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/tasks/debug-plugin", nil) + c.Set(common.RequestIdKey, "plugin-submit-request") + info := &relaycommon.RelayInfo{ + OriginModelName: "safe-model", + TaskRelayInfo: &relaycommon.TaskRelayInfo{Action: "https://private-action.invalid/?key=hidden"}, + } + + newTaskPluginSubmitDiagnostics(c).start(info) + assert.Empty(t, output.String()) + + c.Set("expected_task_plugin_key", "debug-plugin") + diagnostics := newTaskPluginSubmitDiagnostics(c) + diagnostics.start(info) + diagnostics.attemptSucceeded(1, &relay.TaskSubmitResult{ + UpstreamTaskID: "private-upstream-canary", + TaskData: []byte("private-task-data-canary"), + ClientResponse: map[string]any{"secret": "private-client-response-canary"}, + Platform: constant.TaskPlatform("debug-plugin"), + Quota: 12, + }) + task := &model.Task{ + TaskID: "public-task-id", + Platform: constant.TaskPlatform("debug-plugin"), + Status: model.TaskStatus("https://private-status.invalid/?key=hidden"), + PrivateData: model.TaskPrivateData{ + UpstreamTaskID: "private-task-record-canary", + ResultURL: "https://private-url.invalid/result", + }, + } + diagnostics.insertStart(task) + diagnostics.durable(task) + diagnostics.complete(task, 12) + + logOutput := output.String() + require.Contains(t, logOutput, "plugin-submit-request") + assert.Contains(t, logOutput, `plugin="debug-plugin"`) + assert.Contains(t, logOutput, `public_task_id="public-task-id"`) + assert.Contains(t, logOutput, "task_data_bytes=24") + assert.Contains(t, logOutput, "action_present=true") + assert.Contains(t, logOutput, `status="unknown"`) + for _, secret := range []string{ + "private-upstream-canary", + "private-task-data-canary", + "private-client-response-canary", + "private-task-record-canary", + "private-url.invalid", + "private-action.invalid", + "private-status.invalid", + "key=hidden", + } { + assert.NotContains(t, logOutput, secret) + } +} diff --git a/controller/task_plugin_test.go b/controller/task_plugin_test.go new file mode 100644 index 000000000000..a5f39b0b0cb1 --- /dev/null +++ b/controller/task_plugin_test.go @@ -0,0 +1,1166 @@ +package controller + +import ( + "crypto/sha256" + "fmt" + "net/http" + "net/http/httptest" + "regexp" + "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/pkg/jsplugin" + "github.com/QuantumNous/new-api/plugins" + "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/gorm" +) + +func setupTaskPluginControllerTest(t *testing.T) { + t.Helper() + originalDB := model.DB + database, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, database.AutoMigrate(&model.TaskPlugin{}, &model.Channel{}, &model.Ability{}, &model.Task{}, &model.Option{})) + model.DB = database + t.Cleanup(func() { model.DB = originalDB }) +} + +const lifecyclePluginSource = ` +export const meta = {apiVersion: 1, key: "lifecycle-only", name: "Lifecycle", version: "1.0.0", author: {name: "Test"}, models: ["doc"], fetchMode: "per_task"}; +export function buildSubmitRequest() { return {}; } +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +` + +func taskPluginControllerTestSource(key, version string) string { + return fmt.Sprintf(` +export const meta = {apiVersion: 1, key: %q, name: "Test", version: %q, author: {name: "Test"}, models: ["doc-1"], fetchMode: "per_task"}; +export function buildSubmitRequest() { return {}; } +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +`, key, version) +} + +func taskPluginControllerChannelSource(key, version string, channelType int) string { + return fmt.Sprintf(` +export const meta = {apiVersion: 1, key: %q, name: "Test", version: %q, author: {name: "Test"}, channelTypes: [%d], models: ["doc-1"], fetchMode: "per_task"}; +export function buildSubmitRequest() { return {}; } +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +`, key, version, channelType) +} + +func cleanupTaskPluginControllerRuntime(t *testing.T, key string) { + t.Helper() + t.Cleanup(func() { + jsplugin.DefaultRegistry.Unregister(key) + taskPluginSyncState.Lock() + delete(taskPluginSyncState.hashes, key) + delete(taskPluginSyncState.errors, key) + taskPluginSyncState.Unlock() + }) +} + +func TestDeleteThirdPartyPluginReportsAssociatedChannelsAndInFlightTasks(t *testing.T) { + setupTaskPluginControllerTest(t) + loaded, err := jsplugin.DefaultRegistry.Register(lifecyclePluginSource, jsplugin.Options{}) + require.NoError(t, err) + t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("lifecycle-only") }) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{Key: loaded.Meta.Key, APIVersion: 1, Version: "1", Source: lifecyclePluginSource, SourceHash: "hash", Enabled: true})) + baseURL := "https://example.com" + setting := `{"task_plugin_key":"lifecycle-only"}` + channel := model.Channel{Type: constant.ChannelTypeTaskPlugin, Status: common.ChannelStatusEnabled, Name: "linked", Models: "doc", Group: "default", BaseURL: &baseURL, Setting: &setting} + require.NoError(t, channel.Insert()) + require.NoError(t, model.DB.Create(&model.Task{Platform: "lifecycle-only", Status: model.TaskStatusInProgress}).Error) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: "lifecycle-only"}, {Key: "version", Value: "1"}} + context.Request = httptest.NewRequest(http.MethodDelete, "/api/plugin/task/lifecycle-only/versions/1", nil) + DeleteTaskPluginVersion(context) + + assert.Contains(t, recorder.Body.String(), `"name":"linked"`) + assert.Contains(t, recorder.Body.String(), `"in_flight_count":1`) +} + +func TestDisableThirdPartyPluginSupportsCascadeAndForce(t *testing.T) { + setupTaskPluginControllerTest(t) + loaded, err := jsplugin.DefaultRegistry.Register(lifecyclePluginSource, jsplugin.Options{}) + require.NoError(t, err) + t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("lifecycle-only") }) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{Key: loaded.Meta.Key, APIVersion: 1, Version: "1", Source: lifecyclePluginSource, SourceHash: "hash", Enabled: true})) + baseURL := "https://example.com" + setting := `{"task_plugin_key":"lifecycle-only"}` + channel := model.Channel{Type: constant.ChannelTypeTaskPlugin, Status: common.ChannelStatusEnabled, Name: "linked", Models: "doc", Group: "default", BaseURL: &baseURL, Setting: &setting} + require.NoError(t, channel.Insert()) + require.NoError(t, model.DB.Create(&model.Task{Platform: "lifecycle-only", Status: model.TaskStatusSubmitted}).Error) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: "lifecycle-only"}} + context.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task/lifecycle-only/status?cascade=true&force=true", strings.NewReader(`{"enabled":false}`)) + context.Request.Header.Set("Content-Type", "application/json") + SetTaskPluginStatus(context) + + assert.Contains(t, recorder.Body.String(), `"success":true`) + updated, err := model.GetChannelById(channel.Id, true) + require.NoError(t, err) + assert.Equal(t, common.ChannelStatusManuallyDisabled, updated.Status) +} + +// klingFactoryVersion returns the version declared in the embedded kling factory +// manifest so tests do not hardcode a value that moves with every plugin release. +func klingFactoryVersion(t *testing.T) string { + t.Helper() + factorySource, err := plugins.Source("kling") + require.NoError(t, err) + match := regexp.MustCompile(`version:\s*"([^"]+)"`).FindStringSubmatch(factorySource) + require.Len(t, match, 2, "kling factory manifest must declare a version") + return match[1] +} + +func setupTaskPluginFactoryDisableTest(t *testing.T) { + t.Helper() + setupTaskPluginControllerTest(t) + originalMap := common.OptionMap + common.OptionMapRWMutex.Lock() + common.OptionMap = map[string]string{} + common.OptionMapRWMutex.Unlock() + t.Cleanup(func() { + jsplugin.DefaultRegistry.SetDisabledFactoryKeys(nil) + common.OptionMapRWMutex.Lock() + common.OptionMap = originalMap + common.OptionMapRWMutex.Unlock() + }) +} + +func postTaskPluginStatus(t *testing.T, key, query, body string) *httptest.ResponseRecorder { + t.Helper() + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: key}} + path := "/api/plugin/task/" + key + "/status" + if query != "" { + path += "?" + query + } + context.Request = httptest.NewRequest(http.MethodPost, path, strings.NewReader(body)) + context.Request.Header.Set("Content-Type", "application/json") + SetTaskPluginStatus(context) + return recorder +} + +func listTaskPluginItem(t *testing.T, key string) taskPluginListItem { + t.Helper() + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task", nil) + ListTaskPlugins(context) + var response struct { + Success bool `json:"success"` + Data []taskPluginListItem `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + for _, item := range response.Data { + if item.Meta.Key == key { + return item + } + } + t.Fatalf("task plugin %q not found", key) + return taskPluginListItem{} +} + +func taskPluginOptionsHasKey(t *testing.T, key string) bool { + t.Helper() + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/task_plugin_options", nil) + GetTaskPluginOptions(context) + var response struct { + Success bool `json:"success"` + Data []struct { + Key string `json:"key"` + } `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + for _, option := range response.Data { + if option.Key == key { + return true + } + } + return false +} + +func TestDisableFactoryPluginPersistsOptionAndHidesFromBindOptions(t *testing.T) { + setupTaskPluginFactoryDisableTest(t) + const key = "kling" + + recorder := postTaskPluginStatus(t, key, "", `{"enabled":false}`) + assert.Contains(t, recorder.Body.String(), `"success":true`) + assert.Equal(t, []string{key}, setting.GetTaskPluginDisabledFactoryKeys()) + var stored model.Option + require.NoError(t, model.DB.Where("key = ?", setting.TaskPluginDisabledFactoryKeysKey).First(&stored).Error) + assert.Equal(t, `["kling"]`, stored.Value) + + item := listTaskPluginItem(t, key) + assert.Equal(t, "factory", item.Source) + assert.False(t, item.Enabled) + assert.Equal(t, "disabled", item.RuntimeStatus) + assert.False(t, taskPluginOptionsHasKey(t, key)) + _, ok := jsplugin.DefaultRegistry.Get(key) + assert.False(t, ok) + + recorder = postTaskPluginStatus(t, key, "", `{"enabled":true}`) + assert.Contains(t, recorder.Body.String(), `"success":true`) + assert.Empty(t, setting.GetTaskPluginDisabledFactoryKeys()) + item = listTaskPluginItem(t, key) + assert.True(t, item.Enabled) + assert.Equal(t, "registered", item.RuntimeStatus) + assert.True(t, taskPluginOptionsHasKey(t, key)) + _, ok = jsplugin.DefaultRegistry.Get(key) + assert.True(t, ok) +} + +func TestDisableFactoryPluginRespectsInUseGuard(t *testing.T) { + setupTaskPluginFactoryDisableTest(t) + const key = "kling" + baseURL := "https://example.com" + channelSetting := `{"task_plugin_key":"kling"}` + channel := model.Channel{Type: constant.ChannelTypeTaskPlugin, Status: common.ChannelStatusEnabled, Name: "linked-factory", Models: "doc", Group: "default", BaseURL: &baseURL, Setting: &channelSetting} + require.NoError(t, channel.Insert()) + + recorder := postTaskPluginStatus(t, key, "", `{"enabled":false}`) + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), "task plugin is still in use") + assert.Empty(t, setting.GetTaskPluginDisabledFactoryKeys()) + _, ok := jsplugin.DefaultRegistry.Get(key) + assert.True(t, ok) +} + +func TestDisableFactoryOverrideRowKeepsEnabledFlagPath(t *testing.T) { + setupTaskPluginFactoryDisableTest(t) + factorySource, err := plugins.Source("kling") + require.NoError(t, err) + factoryVersion := klingFactoryVersion(t) + overrideSource := strings.Replace(factorySource, `version: "`+factoryVersion+`"`, `version: "`+factoryVersion+`-test-factory-status"`, 1) + require.NotEqual(t, factorySource, overrideSource, "factory version marker must be found in kling source") + loaded, err := jsplugin.DefaultRegistry.Register(overrideSource, jsplugin.Options{}) + require.NoError(t, err) + t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("kling") }) + plugin := model.TaskPlugin{ + Key: "kling", APIVersion: loaded.Meta.APIVersion, Version: loaded.Meta.Version, + Source: overrideSource, SourceHash: "test-hash", Enabled: true, + } + require.NoError(t, model.SaveTaskPlugin(&plugin)) + require.NoError(t, syncTaskPluginsOnce()) + + recorder := postTaskPluginStatus(t, "kling", "", `{"enabled":false}`) + assert.Contains(t, recorder.Body.String(), `"success":true`) + assert.Empty(t, setting.GetTaskPluginDisabledFactoryKeys()) + + row, err := model.GetTaskPluginVersion("kling", "") + require.NoError(t, err) + assert.False(t, row.Enabled) + + item := listTaskPluginItem(t, "kling") + assert.Equal(t, "override_over_factory", item.Source) + assert.False(t, item.Enabled) + assert.Equal(t, "disabled_fallback", item.RuntimeStatus) + assert.True(t, taskPluginOptionsHasKey(t, "kling")) + got, ok := jsplugin.DefaultRegistry.Get("kling") + require.True(t, ok) + assert.Equal(t, factoryVersion, got.Meta.Version) +} + +func TestListTaskPluginsIncludesFactoryWithoutDatabaseRows(t *testing.T) { + setupTaskPluginControllerTest(t) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task", nil) + + ListTaskPlugins(context) + + var response struct { + Success bool `json:"success"` + Data []taskPluginListItem `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + var factoryItem *taskPluginListItem + for i := range response.Data { + if response.Data[i].Meta.Key == "kling" { + factoryItem = &response.Data[i] + break + } + } + require.NotNil(t, factoryItem) + assert.Equal(t, "factory", factoryItem.Source) + assert.Equal(t, "registered", factoryItem.RuntimeStatus) + assert.NotEmpty(t, factoryItem.SourceHash) +} + +func TestMasterSwitchEmptiesOptionsAndKeepsList(t *testing.T) { + setupTaskPluginControllerTest(t) + originalEnabled := constant.TaskPluginEnabled + jsplugin.DefaultRegistry.SetEnabled(false) + t.Cleanup(func() { + constant.TaskPluginEnabled = originalEnabled + jsplugin.DefaultRegistry.SetEnabled(originalEnabled) + }) + + assert.False(t, taskPluginOptionsHasKey(t, "kling")) + item := listTaskPluginItem(t, "kling") + assert.Equal(t, "factory", item.Source) + assert.Equal(t, "kling", item.Meta.Key) +} + +func TestGetTaskPluginOptionsIncludesUsageSchema(t *testing.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"}, + models: ["usage-options-model"], fetchMode: "per_task", + usageSchema: {seconds: {type: "number", unit: "second", description: "Generated media duration."}} +}; +export function buildSubmitRequest() { return {}; } +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +` + _, err := jsplugin.DefaultRegistry.Register(source, jsplugin.Options{}) + require.NoError(t, err) + t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister(key) }) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/task_plugin_options", nil) + + GetTaskPluginOptions(context) + + var response struct { + Success bool `json:"success"` + Data []struct { + Key string `json:"key"` + UsageSchema map[string]jsplugin.UsageFieldSchema `json:"usageSchema"` + } `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + for _, option := range response.Data { + if option.Key != key { + continue + } + assert.Equal(t, "second", option.UsageSchema["seconds"].Unit) + assert.Equal(t, "Generated media duration.", option.UsageSchema["seconds"].Description["en"]) + return + } + t.Fatal("task plugin option not found") +} + +func TestListTaskPluginsShowsDisabledFallbackWhenOverridesAreDisabled(t *testing.T) { + setupTaskPluginControllerTest(t) + factorySource, err := plugins.Source("kling") + require.NoError(t, err) + factoryVersion := klingFactoryVersion(t) + overrideSource := strings.Replace(factorySource, `version: "`+factoryVersion+`"`, `version: "`+factoryVersion+`-test-disabled-override"`, 1) + require.NotEqual(t, factorySource, overrideSource, "factory version marker must be found in kling source") + loaded, err := jsplugin.DefaultRegistry.Register(overrideSource, jsplugin.Options{}) + require.NoError(t, err) + plugin := model.TaskPlugin{ + Key: "kling", APIVersion: loaded.Meta.APIVersion, Version: loaded.Meta.Version, + Source: overrideSource, SourceHash: "test-hash", Enabled: true, + } + require.NoError(t, model.SaveTaskPlugin(&plugin)) + originalEnabled := constant.TaskPluginOverrideEnabled + constant.TaskPluginOverrideEnabled = false + jsplugin.DefaultRegistry.SetOverrideEnabled(false) + t.Cleanup(func() { + constant.TaskPluginOverrideEnabled = originalEnabled + jsplugin.DefaultRegistry.SetOverrideEnabled(originalEnabled) + jsplugin.DefaultRegistry.Unregister("kling") + }) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task", nil) + ListTaskPlugins(context) + + var response struct { + Success bool `json:"success"` + Data []taskPluginListItem `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + for _, item := range response.Data { + if item.Meta.Key == "kling" { + assert.Equal(t, "disabled_fallback", item.RuntimeStatus) + return + } + } + t.Fatal("kling plugin not found") +} + +func TestDeleteActiveOverrideFallsBackToFactoryAndDeletesRecord(t *testing.T) { + setupTaskPluginControllerTest(t) + factorySource, err := plugins.Source("kling") + require.NoError(t, err) + factoryVersion := klingFactoryVersion(t) + overrideSource := strings.Replace(factorySource, `version: "`+factoryVersion+`"`, `version: "`+factoryVersion+`-test-override"`, 1) + require.NotEqual(t, factorySource, overrideSource, "factory version marker must be found in kling source") + loaded, err := jsplugin.DefaultRegistry.Register(overrideSource, jsplugin.Options{Key: "kling", Version: "test-override"}) + require.NoError(t, err) + t.Cleanup(func() { jsplugin.DefaultRegistry.Unregister("kling") }) + plugin := model.TaskPlugin{ + Key: "kling", APIVersion: loaded.Meta.APIVersion, Version: loaded.Meta.Version, + Source: overrideSource, SourceHash: "test-hash", Enabled: true, + } + require.NoError(t, model.SaveTaskPlugin(&plugin)) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: "kling"}, {Key: "version", Value: loaded.Meta.Version}} + context.Request = httptest.NewRequest(http.MethodDelete, "/api/plugin/task/kling/versions/"+loaded.Meta.Version, nil) + + DeleteTaskPluginVersion(context) + + assert.Contains(t, recorder.Body.String(), `"success":true`) + versions, err := model.ListTaskPluginVersions("kling") + require.NoError(t, err) + assert.Empty(t, versions) + runtimePlugin, ok := jsplugin.DefaultRegistry.Get("kling") + require.True(t, ok) + assert.NotEqual(t, loaded.Meta.Version, runtimePlugin.Meta.Version) +} + +func TestDeleteActiveTaskPluginPromotesEnabledVersionInRuntime(t *testing.T) { + setupTaskPluginControllerTest(t) + key := "delete-promote-probe" + cleanupTaskPluginControllerRuntime(t, key) + v1Source := taskPluginControllerTestSource(key, "1.0.0") + v2Source := taskPluginControllerTestSource(key, "2.0.0") + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "1.0.0", Source: v1Source, SourceHash: "hash-v1", Enabled: true, + })) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "2.0.0", Source: v2Source, SourceHash: "hash-v2", Enabled: true, + })) + _, err := jsplugin.DefaultRegistry.Register(v1Source, jsplugin.Options{Key: key, Version: "1.0.0"}) + require.NoError(t, err) + taskPluginSyncState.Lock() + taskPluginSyncState.hashes[key] = "hash-v1" + taskPluginSyncState.errors[key] = "stale compile error" + taskPluginSyncState.Unlock() + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: key}, {Key: "version", Value: "1.0.0"}} + context.Request = httptest.NewRequest(http.MethodDelete, "/api/plugin/task/"+key+"/versions/1.0.0", nil) + DeleteTaskPluginVersion(context) + + assert.Contains(t, recorder.Body.String(), `"success":true`) + active, err := model.GetTaskPluginVersion(key, "") + require.NoError(t, err) + assert.Equal(t, "2.0.0", active.Version) + runtimePlugin, ok := jsplugin.DefaultRegistry.Get(key) + require.True(t, ok) + assert.Equal(t, "2.0.0", runtimePlugin.Meta.Version) + taskPluginSyncState.Lock() + syncedHash := taskPluginSyncState.hashes[key] + _, hasSyncError := taskPluginSyncState.errors[key] + taskPluginSyncState.Unlock() + assert.Equal(t, "hash-v2", syncedHash) + assert.False(t, hasSyncError) +} + +func TestUploadTaskPluginRefreshesRuntimeSyncState(t *testing.T) { + setupTaskPluginControllerTest(t) + key := "upload-sync-probe" + cleanupTaskPluginControllerRuntime(t, key) + source := taskPluginControllerTestSource(key, "1.0.0") + taskPluginSyncState.Lock() + taskPluginSyncState.hashes[key] = "stale-hash" + taskPluginSyncState.errors[key] = "stale compile error" + taskPluginSyncState.Unlock() + body, err := common.Marshal(map[string]any{"source": source}) + require.NoError(t, err) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task", strings.NewReader(string(body))) + context.Request.Header.Set("Content-Type", "application/json") + UploadTaskPlugin(context) + + assert.Contains(t, recorder.Body.String(), `"success":true`) + stored, err := model.GetTaskPluginVersion(key, "") + require.NoError(t, err) + taskPluginSyncState.Lock() + syncedHash := taskPluginSyncState.hashes[key] + _, hasSyncError := taskPluginSyncState.errors[key] + taskPluginSyncState.Unlock() + assert.Equal(t, stored.SourceHash, syncedHash) + assert.False(t, hasSyncError) +} + +func TestActivateTaskPluginRefreshesRuntimeSyncState(t *testing.T) { + setupTaskPluginControllerTest(t) + key := "activate-sync-probe" + cleanupTaskPluginControllerRuntime(t, key) + v1Source := taskPluginControllerTestSource(key, "1.0.0") + v2Source := taskPluginControllerTestSource(key, "2.0.0") + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "1.0.0", Source: v1Source, SourceHash: "hash-v1", Enabled: true, + })) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "2.0.0", Source: v2Source, SourceHash: "hash-v2", Enabled: true, + })) + _, err := jsplugin.DefaultRegistry.Register(v1Source, jsplugin.Options{Key: key, Version: "1.0.0"}) + require.NoError(t, err) + taskPluginSyncState.Lock() + taskPluginSyncState.hashes[key] = "hash-v1" + taskPluginSyncState.errors[key] = "stale compile error" + taskPluginSyncState.Unlock() + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: key}} + context.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task/"+key+"/activate", strings.NewReader(`{"version":"2.0.0"}`)) + context.Request.Header.Set("Content-Type", "application/json") + ActivateTaskPlugin(context) + + assert.Contains(t, recorder.Body.String(), `"success":true`) + runtimePlugin, ok := jsplugin.DefaultRegistry.Get(key) + require.True(t, ok) + assert.Equal(t, "2.0.0", runtimePlugin.Meta.Version) + taskPluginSyncState.Lock() + syncedHash := taskPluginSyncState.hashes[key] + _, hasSyncError := taskPluginSyncState.errors[key] + taskPluginSyncState.Unlock() + assert.Equal(t, "hash-v2", syncedHash) + assert.False(t, hasSyncError) +} + +func TestSyncTaskPluginsPublishesOneGenerationForWholeBatch(t *testing.T) { + setupTaskPluginControllerTest(t) + firstKey := "batch-sync-first" + secondKey := "batch-sync-second" + cleanupTaskPluginControllerRuntime(t, firstKey) + cleanupTaskPluginControllerRuntime(t, secondKey) + firstSource := taskPluginControllerTestSource(firstKey, "1.0.0") + secondSource := taskPluginControllerTestSource(secondKey, "1.0.0") + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: firstKey, APIVersion: 1, Version: "1.0.0", Source: firstSource, SourceHash: "first-hash", Enabled: true, + })) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: secondKey, APIVersion: 1, Version: "1.0.0", Source: secondSource, SourceHash: "second-hash", Enabled: true, + })) + before := jsplugin.DefaultRegistry.Generation().Number + + SyncTaskPluginsOnce() + + assert.Equal(t, before+1, jsplugin.DefaultRegistry.Generation().Number) + _, firstRegistered := jsplugin.DefaultRegistry.Get(firstKey) + _, secondRegistered := jsplugin.DefaultRegistry.Get(secondKey) + assert.True(t, firstRegistered) + assert.True(t, secondRegistered) + + published := jsplugin.DefaultRegistry.Generation() + SyncTaskPluginsOnce() + assert.Same(t, published, jsplugin.DefaultRegistry.Generation()) +} + +func TestTaskPluginRuntimeExposesDatabaseRevisionAheadOfLocalGeneration(t *testing.T) { + setupTaskPluginControllerTest(t) + key := "runtime-revision-probe" + cleanupTaskPluginControllerRuntime(t, key) + taskPluginSyncState.Lock() + previousRebuild := taskPluginSyncState.lastRebuild + taskPluginSyncState.Unlock() + t.Cleanup(func() { + taskPluginSyncState.Lock() + taskPluginSyncState.lastRebuild = previousRebuild + taskPluginSyncState.Unlock() + }) + + v1Source := taskPluginControllerTestSource(key, "1.0.0") + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "1.0.0", + Source: v1Source, SourceHash: "runtime-v1", Enabled: true, + })) + require.NoError(t, syncTaskPluginsOnce()) + localGeneration := jsplugin.DefaultRegistry.Generation().Number + taskPluginSyncState.Lock() + syncedRevision := taskPluginSyncState.lastRebuild.DatabaseRevision + taskPluginSyncState.Unlock() + + v2Source := taskPluginControllerTestSource(key, "2.0.0") + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "2.0.0", + Source: v2Source, SourceHash: "runtime-v2", Enabled: true, + })) + require.NoError(t, model.ActivateTaskPlugin(key, "2.0.0")) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task/runtime/status", nil) + GetTaskPluginRuntime(context) + + var response struct { + Success bool `json:"success"` + Data taskPluginRuntimeStatus `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + assert.Equal(t, localGeneration, response.Data.CurrentGeneration) + assert.NotZero(t, response.Data.GenerationPublishedAt) + assert.NotEqual(t, syncedRevision, response.Data.DatabaseRevision) + assert.Equal(t, "success", response.Data.LastRebuild.Status) + assert.Equal(t, syncedRevision, response.Data.LastRebuild.DatabaseRevision) + assert.Empty(t, response.Data.PluginErrors) + + active, ok := jsplugin.DefaultRegistry.Get(key) + require.True(t, ok) + assert.Equal(t, "1.0.0", active.Meta.Version) +} + +func TestTaskPluginRuntimeReportsPluginLevelCompileErrors(t *testing.T) { + setupTaskPluginControllerTest(t) + key := "runtime-error-probe" + cleanupTaskPluginControllerRuntime(t, key) + taskPluginSyncState.Lock() + previousRebuild := taskPluginSyncState.lastRebuild + taskPluginSyncState.Unlock() + t.Cleanup(func() { + taskPluginSyncState.Lock() + taskPluginSyncState.lastRebuild = previousRebuild + taskPluginSyncState.Unlock() + }) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "1.0.0", + Source: "export const meta = {", SourceHash: "broken-source", Enabled: true, + })) + require.NoError(t, syncTaskPluginsOnce()) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task/runtime/status", nil) + GetTaskPluginRuntime(context) + + var response struct { + Success bool `json:"success"` + Data taskPluginRuntimeStatus `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + assert.Equal(t, "partial", response.Data.LastRebuild.Status) + assert.Equal(t, response.Data.DatabaseRevision, response.Data.LastRebuild.DatabaseRevision) + assert.GreaterOrEqual(t, response.Data.LastRebuild.PluginErrorCount, 1) + assert.NotEmpty(t, response.Data.PluginErrors[key]) + _, registered := jsplugin.DefaultRegistry.Get(key) + assert.False(t, registered) +} + +func TestTaskPluginRuntimeSurvivesDatabaseSyncFailure(t *testing.T) { + setupTaskPluginControllerTest(t) + key := "runtime-database-failure" + cleanupTaskPluginControllerRuntime(t, key) + taskPluginSyncState.Lock() + previousRebuild := taskPluginSyncState.lastRebuild + taskPluginSyncState.Unlock() + t.Cleanup(func() { + taskPluginSyncState.Lock() + taskPluginSyncState.lastRebuild = previousRebuild + taskPluginSyncState.Unlock() + }) + + source := taskPluginControllerTestSource(key, "1.0.0") + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "1.0.0", + Source: source, SourceHash: "runtime-database-v1", Enabled: true, + })) + require.NoError(t, syncTaskPluginsOnce()) + generation := jsplugin.DefaultRegistry.Generation().Number + taskPluginSyncState.Lock() + syncedRevision := taskPluginSyncState.lastRebuild.DatabaseRevision + taskPluginSyncState.Unlock() + + sqlDatabase, err := model.DB.DB() + require.NoError(t, err) + require.NoError(t, sqlDatabase.Close()) + require.Error(t, syncTaskPluginsOnce()) + + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task/runtime/status", nil) + GetTaskPluginRuntime(context) + + var response struct { + Success bool `json:"success"` + Data taskPluginRuntimeStatus `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + assert.Equal(t, generation, response.Data.CurrentGeneration) + assert.Equal(t, syncedRevision, response.Data.DatabaseRevision) + assert.Equal(t, "database snapshot unavailable", response.Data.DatabaseError) + assert.Equal(t, "failed", response.Data.LastRebuild.Status) + assert.Equal(t, syncedRevision, response.Data.LastRebuild.DatabaseRevision) + assert.Contains(t, response.Data.LastRebuild.Error, "sync task plugins") +} + +func TestSyncTaskPluginsCachesRejectedDesiredSourceWithoutLosingIncumbent(t *testing.T) { + setupTaskPluginControllerTest(t) + pluginKey := "sync-retained-plugin" + ownerKey := "sync-retained-owner" + cleanupTaskPluginControllerRuntime(t, pluginKey) + cleanupTaskPluginControllerRuntime(t, ownerKey) + v1Source := taskPluginControllerChannelSource(pluginKey, "1.0.0", 9001) + v2Source := taskPluginControllerChannelSource(pluginKey, "2.0.0", 9002) + ownerSource := taskPluginControllerChannelSource(ownerKey, "1.0.0", 9002) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: pluginKey, APIVersion: 1, Version: "1.0.0", Source: v1Source, SourceHash: "retained-v1", Enabled: true, + })) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: ownerKey, APIVersion: 1, Version: "1.0.0", Source: ownerSource, SourceHash: "owner-v1", Enabled: true, + })) + require.NoError(t, syncTaskPluginsOnce()) + + incumbent, ok := jsplugin.DefaultRegistry.Get(pluginKey) + require.True(t, ok) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: pluginKey, APIVersion: 1, Version: "2.0.0", Source: v2Source, SourceHash: "retained-v2", Enabled: true, + })) + require.NoError(t, model.ActivateTaskPlugin(pluginKey, "2.0.0")) + + require.NoError(t, syncTaskPluginsOnce()) + rejectedGeneration := jsplugin.DefaultRegistry.Generation() + for range 2 { + require.NoError(t, syncTaskPluginsOnce()) + active, found := jsplugin.DefaultRegistry.Get(pluginKey) + require.True(t, found) + assert.Same(t, incumbent, active) + assert.Equal(t, "2.0.0", jsplugin.DefaultRegistry.OverridePlugins()[pluginKey].Meta.Version) + assert.Same(t, incumbent, jsplugin.DefaultRegistry.ActiveOverridePlugins()[pluginKey]) + assert.Contains(t, jsplugin.DefaultRegistry.RoutingErrors()[pluginKey], "channelType 9002 conflicts") + taskPluginSyncState.Lock() + assert.Equal(t, "retained-v2", taskPluginSyncState.hashes[pluginKey]) + taskPluginSyncState.Unlock() + assert.Same(t, rejectedGeneration, jsplugin.DefaultRegistry.Generation()) + } + + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: ownerKey, APIVersion: 1, Version: "2.0.0", + Source: taskPluginControllerChannelSource(ownerKey, "2.0.0", 9003), SourceHash: "owner-v2", Enabled: true, + })) + require.NoError(t, model.ActivateTaskPlugin(ownerKey, "2.0.0")) + require.NoError(t, syncTaskPluginsOnce()) + active, ok := jsplugin.DefaultRegistry.Get(pluginKey) + require.True(t, ok) + assert.Equal(t, "2.0.0", active.Meta.Version) + assert.NotContains(t, jsplugin.DefaultRegistry.RoutingErrors(), pluginKey) +} + +func TestSyncTaskPluginsPreservesLastCompiledOverrideWhileOverridesAreDisabled(t *testing.T) { + setupTaskPluginControllerTest(t) + key := "sync-disabled-override" + cleanupTaskPluginControllerRuntime(t, key) + v1Source := taskPluginControllerTestSource(key, "1.0.0") + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "1.0.0", Source: v1Source, SourceHash: "disabled-v1", Enabled: true, + })) + require.NoError(t, syncTaskPluginsOnce()) + jsplugin.DefaultRegistry.SetOverrideEnabled(false) + t.Cleanup(func() { jsplugin.DefaultRegistry.SetOverrideEnabled(true) }) + + disabledGeneration := jsplugin.DefaultRegistry.Generation() + require.NoError(t, syncTaskPluginsOnce()) + assert.Same(t, disabledGeneration, jsplugin.DefaultRegistry.Generation()) + + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{ + Key: key, APIVersion: 1, Version: "2.0.0", + Source: "export const meta = {", SourceHash: "disabled-v2", Enabled: true, + })) + require.NoError(t, model.ActivateTaskPlugin(key, "2.0.0")) + require.NoError(t, syncTaskPluginsOnce()) + assert.Equal(t, "1.0.0", jsplugin.DefaultRegistry.OverridePlugins()[key].Meta.Version) + taskPluginSyncState.Lock() + assert.Equal(t, "disabled-v1", taskPluginSyncState.hashes[key]) + assert.NotEmpty(t, taskPluginSyncState.errors[key]) + taskPluginSyncState.Unlock() + + jsplugin.DefaultRegistry.SetOverrideEnabled(true) + active, ok := jsplugin.DefaultRegistry.Get(key) + require.True(t, ok) + assert.Equal(t, "1.0.0", active.Meta.Version) +} + +const dryRunPluginSource = ` +export const meta = {apiVersion: 1, key: "dryrun-probe", name: "DryRun", version: "1.0.0", author: {name: "Test"}, models: ["doc-1"], fetchMode: "per_task"}; +export function buildSubmitRequest(payload) { + if (!payload || !payload.model) { throw new Error("model required"); } + return {model: payload.model}; +} +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +export const native = { info: function(ctx, task) { return "task:" + task.id; } }; +` + +func runTaskPluginDryRun(t *testing.T, body string) *httptest.ResponseRecorder { + t.Helper() + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: "dryrun-probe"}} + context.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task/dryrun-probe/dryrun", strings.NewReader(body)) + context.Request.Header.Set("Content-Type", "application/json") + DryRunTaskPlugin(context) + return recorder +} + +func TestDryRunTaskPluginExecutesHookAndRendererMember(t *testing.T) { + setupTaskPluginControllerTest(t) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{Key: "dryrun-probe", APIVersion: 1, Version: "1.0.0", Source: dryRunPluginSource, SourceHash: "hash", Enabled: true})) + + hookRecorder := runTaskPluginDryRun(t, `{"hook":"buildSubmitRequest","args":[{"model":"doc-1"}]}`) + assert.Contains(t, hookRecorder.Body.String(), `"success":true`) + assert.Contains(t, hookRecorder.Body.String(), `"model":"doc-1"`) + + memberRecorder := runTaskPluginDryRun(t, `{"hook":"native","member":"info","args":[{}, {"id":"t-1"}]}`) + assert.Contains(t, memberRecorder.Body.String(), `"success":true`) + assert.Contains(t, memberRecorder.Body.String(), "task:t-1") +} + +func TestDryRunTaskPluginReportsUnknownHook(t *testing.T) { + setupTaskPluginControllerTest(t) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{Key: "dryrun-probe", APIVersion: 1, Version: "1.0.0", Source: dryRunPluginSource, SourceHash: "hash", Enabled: true})) + + recorder := runTaskPluginDryRun(t, `{"hook":"missingHook"}`) + + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), `plugin export \"missingHook\" not found`) +} + +func TestDryRunTaskPluginSurfacesBadArgumentErrors(t *testing.T) { + setupTaskPluginControllerTest(t) + require.NoError(t, model.SaveTaskPlugin(&model.TaskPlugin{Key: "dryrun-probe", APIVersion: 1, Version: "1.0.0", Source: dryRunPluginSource, SourceHash: "hash", Enabled: true})) + + malformedRecorder := runTaskPluginDryRun(t, `{"hook":"buildSubmitRequest","args":[{`) + assert.Contains(t, malformedRecorder.Body.String(), `"success":false`) + + rejectedRecorder := runTaskPluginDryRun(t, `{"hook":"buildSubmitRequest","args":[{}]}`) + assert.Contains(t, rejectedRecorder.Body.String(), `"success":false`) + assert.Contains(t, rejectedRecorder.Body.String(), "model required") +} + +func TestUploadTaskPluginPreflightConflict(t *testing.T) { + setupTaskPluginControllerTest(t) + enabledFalse := false + tests := []struct { + name string + key string + enabled *bool + force bool + wantSuccess bool + wantError string + }{ + { + name: "enabled conflict rejected", + key: "preflight-reject", + wantError: "channelType 50 conflicts", + }, + { + name: "force saves despite conflict", + key: "preflight-force", + force: true, + wantSuccess: true, + }, + { + name: "disabled skips preflight", + key: "preflight-disabled", + enabled: &enabledFalse, + wantSuccess: true, + }, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + cleanupTaskPluginControllerRuntime(t, testCase.key) + source := taskPluginControllerChannelSource(testCase.key, "1.0.0", 50) + payload := map[string]any{"source": source} + if testCase.enabled != nil { + payload["enabled"] = *testCase.enabled + } + if testCase.force { + payload["force"] = true + } + body, err := common.Marshal(payload) + require.NoError(t, err) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task", strings.NewReader(string(body))) + context.Request.Header.Set("Content-Type", "application/json") + + UploadTaskPlugin(context) + + if testCase.wantSuccess { + assert.Contains(t, recorder.Body.String(), `"success":true`) + _, err = model.GetTaskPluginVersion(testCase.key, "") + require.NoError(t, err) + return + } + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), testCase.wantError) + assert.Contains(t, recorder.Body.String(), "kling") + var count int64 + require.NoError(t, model.DB.Model(&model.TaskPlugin{}).Where("key = ?", testCase.key).Count(&count).Error) + assert.Zero(t, count) + }) + } +} + +func TestUploadTaskPluginRejectsMetaViolatingV1Schema(t *testing.T) { + setupTaskPluginControllerTest(t) + cases := []struct { + name string + meta string + expectedError string + }{ + { + name: "key with uppercase characters", + meta: `{apiVersion: 1, key: "Bad-Key", name: "Bad", version: "1.0.0", author: {name: "Test"}, models: ["doc-1"], fetchMode: "per_task"}`, + expectedError: "plugin meta key must match", + }, + { + name: "version that is not semver", + meta: `{apiVersion: 1, key: "bad-plugin", name: "Bad", version: "one", author: {name: "Test"}, models: ["doc-1"], fetchMode: "per_task"}`, + expectedError: "plugin meta version must be semver", + }, + { + name: "unsupported fetch mode", + meta: `{apiVersion: 1, key: "bad-plugin", name: "Bad", version: "1.0.0", author: {name: "Test"}, models: ["doc-1"], fetchMode: "sometimes"}`, + expectedError: "plugin meta fetchMode must be per_task or batch", + }, + { + name: "empty model list", + meta: `{apiVersion: 1, key: "bad-plugin", name: "Bad", version: "1.0.0", author: {name: "Test"}, models: [], fetchMode: "per_task"}`, + expectedError: "plugin meta models must contain at least one model", + }, + } + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + source := "export const meta = " + testCase.meta + `; +export function buildSubmitRequest() { return {}; } +export function parseSubmitResponse() { return {}; } +export function buildQueryRequest() { return {}; } +export function parseTaskResult() { return {}; } +` + body, err := common.Marshal(map[string]any{"source": source}) + require.NoError(t, err) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task", strings.NewReader(string(body))) + context.Request.Header.Set("Content-Type", "application/json") + + UploadTaskPlugin(context) + + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), testCase.expectedError) + var count int64 + require.NoError(t, model.DB.Model(&model.TaskPlugin{}).Count(&count).Error) + assert.Zero(t, count) + }) + } +} + +func TestDeletePureFactoryPluginIsRejected(t *testing.T) { + setupTaskPluginControllerTest(t) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Params = gin.Params{{Key: "key", Value: "kling"}, {Key: "version", Value: "1.0.0"}} + context.Request = httptest.NewRequest(http.MethodDelete, "/api/plugin/task/kling/versions/1.0.0", nil) + + DeleteTaskPluginVersion(context) + + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), "factory plugins cannot be deleted") + _, ok := jsplugin.DefaultRegistry.Get("kling") + assert.True(t, ok) +} + +func TestUploadTaskPluginSourceSha256(t *testing.T) { + setupTaskPluginControllerTest(t) + tests := []struct { + name string + key string + withHash bool + hash string + wantSuccess bool + wantError string + }{ + { + name: "matching hash succeeds", + key: "sha256-match", + withHash: true, + wantSuccess: true, + }, + { + name: "mismatching hash rejected", + key: "sha256-mismatch", + withHash: true, + hash: "deadbeef", + wantError: "plugin source sha256 mismatch", + }, + { + name: "absent field unchanged", + key: "sha256-absent", + wantSuccess: true, + }, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + cleanupTaskPluginControllerRuntime(t, testCase.key) + source := taskPluginControllerTestSource(testCase.key, "1.0.0") + payload := map[string]any{"source": source} + if testCase.withHash { + hash := testCase.hash + if hash == "" { + hash = " " + strings.ToUpper(fmt.Sprintf("%x", sha256.Sum256([]byte(source)))) + " " + } + payload["sourceSha256"] = hash + } + body, err := common.Marshal(payload) + require.NoError(t, err) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodPost, "/api/plugin/task", strings.NewReader(string(body))) + context.Request.Header.Set("Content-Type", "application/json") + + UploadTaskPlugin(context) + + if testCase.wantSuccess { + assert.Contains(t, recorder.Body.String(), `"success":true`) + _, err = model.GetTaskPluginVersion(testCase.key, "") + require.NoError(t, err) + return + } + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), testCase.wantError) + var count int64 + require.NoError(t, model.DB.Model(&model.TaskPlugin{}).Where("key = ?", testCase.key).Count(&count).Error) + assert.Zero(t, count) + }) + } +} + +func setupTaskPluginMarketplaceSourcesTest(t *testing.T) { + t.Helper() + setupTaskPluginControllerTest(t) + originalMap := common.OptionMap + common.OptionMapRWMutex.Lock() + common.OptionMap = map[string]string{} + common.OptionMapRWMutex.Unlock() + t.Cleanup(func() { + common.OptionMapRWMutex.Lock() + common.OptionMap = originalMap + common.OptionMapRWMutex.Unlock() + }) +} + +func TestGetTaskPluginMarketplaceSourcesDefaultWhenUnset(t *testing.T) { + setupTaskPluginMarketplaceSourcesTest(t) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task/marketplace/sources", nil) + + GetTaskPluginMarketplaceSources(context) + + var response struct { + Success bool `json:"success"` + Data []setting.TaskPluginMarketplaceSource `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success) + require.Equal(t, []setting.TaskPluginMarketplaceSource{ + {Name: "Official", IndexURL: "https://www.newapi.ai/api/v1/plugins/index.json"}, + {Name: "GitHub", IndexURL: "https://raw.githubusercontent.com/QuantumNous/new-api-plugins/main/index.json"}, + }, response.Data) + var count int64 + require.NoError(t, model.DB.Model(&model.Option{}).Where("key = ?", setting.TaskPluginMarketplaceSourcesKey).Count(&count).Error) + assert.Zero(t, count) +} + +func TestUpdateTaskPluginMarketplaceSourcesRoundTrip(t *testing.T) { + setupTaskPluginMarketplaceSourcesTest(t) + payload := []setting.TaskPluginMarketplaceSource{ + {Name: "Mirror", IndexURL: "https://example.com/plugins/index.json"}, + {Name: "Official", IndexURL: "https://www.newapi.ai/api/v1/plugins/index.json"}, + } + body, err := common.Marshal(payload) + require.NoError(t, err) + putRecorder := httptest.NewRecorder() + putContext, _ := gin.CreateTestContext(putRecorder) + putContext.Request = httptest.NewRequest(http.MethodPut, "/api/plugin/task/marketplace/sources", strings.NewReader(string(body))) + putContext.Request.Header.Set("Content-Type", "application/json") + + UpdateTaskPluginMarketplaceSources(putContext) + + assert.Contains(t, putRecorder.Body.String(), `"success":true`) + + getRecorder := httptest.NewRecorder() + getContext, _ := gin.CreateTestContext(getRecorder) + getContext.Request = httptest.NewRequest(http.MethodGet, "/api/plugin/task/marketplace/sources", nil) + GetTaskPluginMarketplaceSources(getContext) + + var response struct { + Success bool `json:"success"` + Data []setting.TaskPluginMarketplaceSource `json:"data"` + } + require.NoError(t, common.Unmarshal(getRecorder.Body.Bytes(), &response)) + require.True(t, response.Success) + assert.Equal(t, payload, response.Data) +} + +func TestUpdateTaskPluginMarketplaceSourcesValidation(t *testing.T) { + setupTaskPluginMarketplaceSourcesTest(t) + tests := []struct { + name string + body string + wantErr string + }{ + { + name: "empty name", + body: `[{"name":"","index_url":"https://example.com/index.json"}]`, + wantErr: "marketplace source name is required", + }, + { + name: "invalid URL", + body: `[{"name":"Local","index_url":"not-a-url"}]`, + wantErr: "marketplace source index_url must be an absolute http(s) URL", + }, + } + for _, testCase := range tests { + t.Run(testCase.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Request = httptest.NewRequest(http.MethodPut, "/api/plugin/task/marketplace/sources", strings.NewReader(testCase.body)) + context.Request.Header.Set("Content-Type", "application/json") + + UpdateTaskPluginMarketplaceSources(context) + + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), testCase.wantErr) + assert.Empty(t, common.OptionMap[setting.TaskPluginMarketplaceSourcesKey]) + var count int64 + require.NoError(t, model.DB.Model(&model.Option{}).Where("key = ?", setting.TaskPluginMarketplaceSourcesKey).Count(&count).Error) + assert.Zero(t, count) + }) + } +} diff --git a/controller/token.go b/controller/token.go index 836e9b2952ac..09d14fcab373 100644 --- a/controller/token.go +++ b/controller/token.go @@ -7,30 +7,126 @@ import ( "strings" "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/service" + "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/operation_setting" "github.com/gin-gonic/gin" + "github.com/shopspring/decimal" ) -func buildMaskedTokenResponse(token *model.Token) *model.Token { +type tokenAutoGroupsInput struct { + Set bool + Groups []string +} + +func (input *tokenAutoGroupsInput) UnmarshalJSON(data []byte) error { + input.Set = true + if strings.TrimSpace(string(data)) == "null" { + input.Groups = nil + return nil + } + return common.Unmarshal(data, &input.Groups) +} + +type tokenRequest struct { + model.Token + AutoGroups tokenAutoGroupsInput `json:"auto_groups"` +} + +type tokenResponse struct { + *model.Token + AutoGroups []string `json:"auto_groups"` +} + +func maxTokenQuota() int { + quota, err := common.WalletQuotaFromDecimalStrict( + decimal.NewFromInt(1_000_000_000).Mul(decimal.NewFromFloat(common.QuotaPerUnit)), + ) + if err != nil { + return common.MaxWalletQuota + } + return quota +} + +func buildMaskedTokenResponse(token *model.Token) *tokenResponse { if token == nil { return nil } maskedToken := *token maskedToken.Key = token.GetMaskedKey() - return &maskedToken + autoGroups, err := token.GetAutoGroups() + if err != nil { + common.SysError(fmt.Sprintf("failed to parse auto groups for token %d: %v", token.Id, err)) + autoGroups = nil + } + if len(autoGroups) == 0 { + autoGroups = nil + } + return &tokenResponse{Token: &maskedToken, AutoGroups: autoGroups} } -func buildMaskedTokenResponses(tokens []*model.Token) []*model.Token { - maskedTokens := make([]*model.Token, 0, len(tokens)) +func buildMaskedTokenResponses(tokens []*model.Token) []*tokenResponse { + maskedTokens := make([]*tokenResponse, 0, len(tokens)) for _, token := range tokens { maskedTokens = append(maskedTokens, buildMaskedTokenResponse(token)) } return maskedTokens } +func getTokenRequestUserGroup(c *gin.Context) (string, error) { + if userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup); userGroup != "" { + return userGroup, nil + } + if userGroup := c.GetString("group"); userGroup != "" { + return userGroup, nil + } + return model.GetUserGroup(c.GetInt("id"), false) +} + +func setTokenAutoGroups(c *gin.Context, token *model.Token, groups []string) bool { + if len(groups) == 0 { + if err := token.SetAutoGroups(nil); err != nil { + common.ApiError(c, err) + return false + } + return true + } + + maxCount := setting.GetMaxTokenAutoGroups() + if len(groups) > 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(groups)) + for _, group := range groups { + if _, ok := seen[group]; ok { + common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsDuplicate, map[string]any{"Group": group}) + return false + } + seen[group] = struct{}{} + if !service.IsUserSelectableGroup(userGroup, group) { + common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsInvalid, map[string]any{"Group": group}) + return false + } + } + + if err := token.SetAutoGroups(groups); err != nil { + common.ApiError(c, err) + return false + } + return true +} + func GetAllTokens(c *gin.Context) { userId := c.GetInt("id") pageInfo := common.GetPageQuery(c) @@ -77,6 +173,18 @@ func GetToken(c *gin.Context) { common.ApiSuccess(c, buildMaskedTokenResponse(token)) } +func GetTokenAutoGroups(c *gin.Context) { + userGroup, err := getTokenRequestUserGroup(c) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{ + "groups": service.GetUserAutoGroup(userGroup), + "max_count": setting.GetMaxTokenAutoGroups(), + }) +} + func GetTokenKey(c *gin.Context) { id, err := strconv.Atoi(c.Param("id")) userId := c.GetInt("id") @@ -165,12 +273,13 @@ func GetTokenUsage(c *gin.Context) { } func AddToken(c *gin.Context) { - token := model.Token{} - err := c.ShouldBindJSON(&token) + request := tokenRequest{} + err := c.ShouldBindJSON(&request) if err != nil { common.ApiError(c, err) return } + token := request.Token if len(token.Name) > 50 { common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong) return @@ -181,7 +290,7 @@ func AddToken(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgTokenQuotaNegative) return } - maxQuotaValue := int((1000000000 * common.QuotaPerUnit)) + maxQuotaValue := maxTokenQuota() if token.RemainQuota > maxQuotaValue { common.ApiErrorI18n(c, i18n.MsgTokenQuotaExceedMax, map[string]any{"Max": maxQuotaValue}) return @@ -201,6 +310,14 @@ func AddToken(c *gin.Context) { }) return } + if token.Group == "auto" { + if !setTokenAutoGroups(c, &token, request.AutoGroups.Groups) { + return + } + } else { + token.CrossGroupRetry = false + _ = token.SetAutoGroups(nil) + } key, err := common.GenerateKey() if err != nil { common.ApiErrorI18n(c, i18n.MsgTokenGenerateFailed) @@ -221,6 +338,7 @@ func AddToken(c *gin.Context) { AllowIps: token.AllowIps, Group: token.Group, CrossGroupRetry: token.CrossGroupRetry, + AutoGroups: token.AutoGroups, } err = cleanToken.Insert() if err != nil { @@ -250,12 +368,13 @@ func DeleteToken(c *gin.Context) { func UpdateToken(c *gin.Context) { userId := c.GetInt("id") statusOnly := c.Query("status_only") - token := model.Token{} - err := c.ShouldBindJSON(&token) + request := tokenRequest{} + err := c.ShouldBindJSON(&request) if err != nil { common.ApiError(c, err) return } + token := request.Token if len(token.Name) > 50 { common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong) return @@ -265,7 +384,7 @@ func UpdateToken(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgTokenQuotaNegative) return } - maxQuotaValue := int((1000000000 * common.QuotaPerUnit)) + maxQuotaValue := maxTokenQuota() if token.RemainQuota > maxQuotaValue { common.ApiErrorI18n(c, i18n.MsgTokenQuotaExceedMax, map[string]any{"Max": maxQuotaValue}) return @@ -299,6 +418,14 @@ func UpdateToken(c *gin.Context) { cleanToken.AllowIps = token.AllowIps cleanToken.Group = token.Group cleanToken.CrossGroupRetry = token.CrossGroupRetry + if token.Group != "auto" { + cleanToken.CrossGroupRetry = false + _ = cleanToken.SetAutoGroups(nil) + } else if request.AutoGroups.Set { + if !setTokenAutoGroups(c, cleanToken, request.AutoGroups.Groups) { + return + } + } } err = cleanToken.Update() if err != nil { diff --git a/controller/token_auto_groups_test.go b/controller/token_auto_groups_test.go new file mode 100644 index 000000000000..3da2969575c1 --- /dev/null +++ b/controller/token_auto_groups_test.go @@ -0,0 +1,234 @@ +package controller + +import ( + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func configureTokenAutoGroupsTest(t *testing.T, maxCount string, autoGroups string) { + t.Helper() + originalMax := setting.GetMaxTokenAutoGroups() + originalAutoGroups := setting.AutoGroups2JsonString() + originalUsableGroups := setting.UserUsableGroups2JSONString() + originalRatios := ratio_setting.GroupRatio2JSONString() + require.NoError(t, setting.UpdateMaxTokenAutoGroups(maxCount)) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(autoGroups)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`)) + t.Cleanup(func() { + require.NoError(t, setting.UpdateMaxTokenAutoGroups(stringInt(originalMax))) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(originalAutoGroups)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios)) + }) +} + +func stringInt(value int) string { + return fmt.Sprintf("%d", value) +} + +func setupTokenAutoGroupsControllerTest(t *testing.T) *model.User { + t.Helper() + db := setupTokenControllerTestDB(t) + require.NoError(t, db.AutoMigrate(&model.User{})) + user := &model.User{ + Id: 101, + Username: "token-auto-user", + Password: "password", + Group: "default", + Status: common.UserStatusEnabled, + } + require.NoError(t, db.Create(user).Error) + return user +} + +func baseAutoTokenRequest(name string) map[string]any { + return map[string]any{ + "name": name, + "expired_time": -1, + "remain_quota": 0, + "unlimited_quota": true, + "group": "auto", + "cross_group_retry": true, + } +} + +func newTokenAutoGroupsAuthenticatedContext(t *testing.T, method string, target string, body any, userID int) (*gin.Context, *httptest.ResponseRecorder) { + t.Helper() + ctx, recorder := newAuthenticatedContext(t, method, target, body, userID) + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") + return ctx, recorder +} + +func TestAddTokenEmptyAutoGroupsInheritGlobalAuto(t *testing.T) { + tests := []struct { + name string + includeField bool + value any + }{ + {name: "omitted"}, + {name: "null", includeField: true, value: nil}, + {name: "empty array", includeField: true, value: []string{}}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + configureTokenAutoGroupsTest(t, "5", `["default","vip"]`) + user := setupTokenAutoGroupsControllerTest(t) + request := baseAutoTokenRequest("create-" + test.name) + if test.includeField { + request["auto_groups"] = test.value + } + + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPost, "/api/token/", request, user.Id) + AddToken(ctx) + + response := decodeAPIResponse(t, recorder) + require.True(t, response.Success, response.Message) + var token model.Token + require.NoError(t, model.DB.Where("name = ?", request["name"]).First(&token).Error) + assert.Empty(t, token.AutoGroups) + assert.True(t, token.CrossGroupRetry) + payload, err := common.Marshal(buildMaskedTokenResponse(&token)) + require.NoError(t, err) + var responseData map[string]any + require.NoError(t, common.Unmarshal(payload, &responseData)) + assert.Nil(t, responseData["auto_groups"]) + }) + } +} + +func TestAddTokenPersistsOrderedAutoGroupsSnapshot(t *testing.T) { + configureTokenAutoGroupsTest(t, "5", `["default","vip"]`) + user := setupTokenAutoGroupsControllerTest(t) + request := baseAutoTokenRequest("ordered-snapshot") + request["auto_groups"] = []string{"vip", "default"} + + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPost, "/api/token/", request, user.Id) + AddToken(ctx) + require.True(t, decodeAPIResponse(t, recorder).Success) + + var token model.Token + require.NoError(t, model.DB.Where("name = ?", "ordered-snapshot").First(&token).Error) + assert.JSONEq(t, `["vip","default"]`, token.AutoGroups) + + getCtx, getRecorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodGet, "/api/token/"+stringInt(token.Id), nil, user.Id) + getCtx.Params = append(getCtx.Params, gin.Param{Key: "id", Value: stringInt(token.Id)}) + GetToken(getCtx) + getResponse := decodeAPIResponse(t, getRecorder) + require.True(t, getResponse.Success) + var data struct { + AutoGroups []string `json:"auto_groups"` + } + require.NoError(t, common.Unmarshal(getResponse.Data, &data)) + assert.Equal(t, []string{"vip", "default"}, data.AutoGroups) +} + +func TestUpdateTokenAutoGroupsTriStateAndNonAutoCleanup(t *testing.T) { + tests := []struct { + name string + includeField bool + value any + group string + expectedAutoGroups string + expectedRetry bool + }{ + {name: "omitted preserves", group: "auto", expectedAutoGroups: `["vip","default"]`, expectedRetry: true}, + {name: "null inherits", includeField: true, value: nil, group: "auto", expectedRetry: true}, + {name: "empty inherits", includeField: true, value: []string{}, group: "auto", expectedRetry: true}, + {name: "non auto clears and disables retry", includeField: true, value: []string{"vip"}, group: "default"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + configureTokenAutoGroupsTest(t, "5", `["default","vip"]`) + user := setupTokenAutoGroupsControllerTest(t) + token := seedToken(t, model.DB, user.Id, "update-auto", "update-auto-key") + token.Group = "auto" + token.CrossGroupRetry = true + require.NoError(t, token.SetAutoGroups([]string{"vip", "default"})) + require.NoError(t, model.DB.Save(token).Error) + + request := baseAutoTokenRequest("updated-auto") + request["id"] = token.Id + request["status"] = common.TokenStatusEnabled + request["group"] = test.group + if test.includeField { + request["auto_groups"] = test.value + } + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPut, "/api/token/", request, user.Id) + UpdateToken(ctx) + response := decodeAPIResponse(t, recorder) + require.True(t, response.Success, response.Message) + + var updated model.Token + require.NoError(t, model.DB.First(&updated, token.Id).Error) + if test.expectedAutoGroups == "" { + assert.Empty(t, updated.AutoGroups) + } else { + assert.JSONEq(t, test.expectedAutoGroups, updated.AutoGroups) + } + assert.Equal(t, test.expectedRetry, updated.CrossGroupRetry) + }) + } +} + +func TestAddTokenRejectsInvalidAutoGroups(t *testing.T) { + tests := []struct { + name string + maxCount string + groups []string + }{ + {name: "over limit", maxCount: "1", groups: []string{"default", "vip"}}, + {name: "duplicate", maxCount: "5", groups: []string{"default", "default"}}, + {name: "auto pseudo group", maxCount: "5", groups: []string{"auto"}}, + {name: "unavailable", maxCount: "5", groups: []string{"missing"}}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + configureTokenAutoGroupsTest(t, test.maxCount, `["default","vip"]`) + user := setupTokenAutoGroupsControllerTest(t) + request := baseAutoTokenRequest("invalid-" + test.name) + request["auto_groups"] = test.groups + + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPost, "/api/token/", request, user.Id) + AddToken(ctx) + + response := decodeAPIResponse(t, recorder) + assert.False(t, response.Success) + var count int64 + require.NoError(t, model.DB.Model(&model.Token{}).Count(&count).Error) + assert.Zero(t, count) + }) + } +} + +func TestGetTokenAutoGroupsReturnsFullFilteredGlobalOrderAndLimit(t *testing.T) { + configureTokenAutoGroupsTest(t, "1", `["vip","missing","default"]`) + user := setupTokenAutoGroupsControllerTest(t) + + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodGet, "/api/token/auto-groups", nil, user.Id) + GetTokenAutoGroups(ctx) + + response := decodeAPIResponse(t, recorder) + require.True(t, response.Success, response.Message) + var data struct { + Groups []string `json:"groups"` + MaxCount int `json:"max_count"` + } + require.NoError(t, common.Unmarshal(response.Data, &data)) + assert.Equal(t, []string{"vip", "default"}, data.Groups) + assert.Equal(t, 1, data.MaxCount) +} diff --git a/controller/token_test.go b/controller/token_test.go index 12b1cbdd84fb..9cca168ab100 100644 --- a/controller/token_test.go +++ b/controller/token_test.go @@ -273,6 +273,34 @@ func getTokenKeyColumnType(t *testing.T, db *gorm.DB, dialect string) string { } } +func getTokenAutoGroupsColumnType(t *testing.T, db *gorm.DB, dialect string) string { + t.Helper() + + switch dialect { + case "sqlite": + return getSQLiteColumnType(t, db, "tokens", "auto_groups") + case "mysql": + var columnType string + if err := db.Raw(`SELECT DATA_TYPE FROM information_schema.columns + WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?`, + "tokens", "auto_groups").Scan(&columnType).Error; err != nil { + t.Fatalf("failed to inspect mysql token auto_groups column: %v", err) + } + return strings.ToLower(columnType) + case "postgres": + var dataType string + if err := db.Raw(`SELECT data_type FROM information_schema.columns + WHERE table_schema = current_schema() AND table_name = ? AND column_name = ?`, + "tokens", "auto_groups").Scan(&dataType).Error; err != nil { + t.Fatalf("failed to inspect postgres token auto_groups column: %v", err) + } + return strings.ToLower(dataType) + default: + t.Fatalf("unsupported dialect %q", dialect) + return "" + } +} + func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect string, managedTokensTable *bool) { t.Helper() @@ -314,6 +342,12 @@ func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect strin if got := getTokenKeyColumnType(t, db, dialect); got != "varchar(128)" { t.Fatalf("expected migrated key column type varchar(128), got %q", got) } + if !db.Migrator().HasColumn(&model.Token{}, "auto_groups") { + t.Fatal("expected migration to add auto_groups column") + } + if got := getTokenAutoGroupsColumnType(t, db, dialect); got != "text" { + t.Fatalf("expected migrated auto_groups column type text, got %q", got) + } var migratedToken model.Token if err := db.First(&migratedToken, "name = ?", "legacy-token").Error; err != nil { @@ -325,6 +359,9 @@ func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect strin if migratedToken.Name != "legacy-token" { t.Fatalf("expected migrated token name to be preserved, got %q", migratedToken.Name) } + if migratedToken.AutoGroups != "" { + t.Fatalf("expected legacy token to inherit global Auto groups, got %q", migratedToken.AutoGroups) + } inserted := model.Token{ UserId: 8, @@ -362,6 +399,9 @@ func TestTokenAutoMigrateUsesVarchar128KeyColumn(t *testing.T) { if got := getTokenKeyColumnType(t, db, "sqlite"); got != "varchar(128)" { t.Fatalf("expected key column type varchar(128), got %q", got) } + if got := getSQLiteColumnType(t, db, "tokens", "auto_groups"); got != "text" { + t.Fatalf("expected auto_groups column type text, got %q", got) + } } func TestTokenMigrationFromChar48ToVarchar128(t *testing.T) { diff --git a/controller/topup.go b/controller/topup.go index 390f53f7dce8..30f8d2215784 100644 --- a/controller/topup.go +++ b/controller/topup.go @@ -1,6 +1,7 @@ package controller import ( + "errors" "fmt" "net/http" "net/url" @@ -181,11 +182,91 @@ func getMinTopup() int64 { if operation_setting.GetQuotaDisplayType() == operation_setting.QuotaDisplayTypeTokens { dMinTopup := decimal.NewFromInt(int64(minTopup)) dQuotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit) - minTopup = int(dMinTopup.Mul(dQuotaPerUnit).IntPart()) + quota, err := common.WalletQuotaFromDecimalStrict(dMinTopup.Mul(dQuotaPerUnit)) + if err != nil { + return common.MaxWalletQuota + } + minTopup = quota } return int64(minTopup) } +func getTopUpQuota(amount int64) (int, error) { + quota := decimal.NewFromInt(amount) + if operation_setting.GetQuotaDisplayType() == operation_setting.QuotaDisplayTypeTokens { + quotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit) + quota = decimal.NewFromInt(quota.Div(quotaPerUnit).IntPart()).Mul(quotaPerUnit) + } else { + quota = quota.Mul(decimal.NewFromFloat(common.QuotaPerUnit)) + } + return common.WalletQuotaFromDecimalStrict(quota) +} + +func getMaxTopUpAmount() int64 { + if common.QuotaPerUnit <= 0 { + return 0 + } + quotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit) + maxStoredAmount := decimal.NewFromInt(common.MaxWalletQuota). + Div(quotaPerUnit). + Floor() + if operation_setting.GetQuotaDisplayType() == operation_setting.QuotaDisplayTypeTokens { + return maxStoredAmount.Add(decimal.NewFromInt(1)). + Mul(quotaPerUnit). + Ceil(). + Sub(decimal.NewFromInt(1)). + IntPart() + } + return maxStoredAmount.IntPart() +} + +func validateCreditedQuota(quota decimal.Decimal) (int, error) { + value, err := common.WalletQuotaFromDecimalStrict(quota) + if err != nil { + return 0, errors.New("充值额度超出系统可表示范围") + } + if value <= 0 { + return 0, errors.New("充值额度必须大于 0") + } + return value, nil +} + +func validateTopUpQuota(amount int64) (int, error) { + quota, err := getTopUpQuota(amount) + if err == nil && quota > 0 { + return quota, nil + } + maxAmount := getMaxTopUpAmount() + if maxAmount > 0 && amount > maxAmount { + return 0, fmt.Errorf("单笔充值数量不能大于 %d", maxAmount) + } + return 0, errors.New("充值数量无效") +} + +func rejectInvalidCreditedQuota(c *gin.Context, userId int, quota decimal.Decimal) bool { + creditedQuota, err := validateCreditedQuota(quota) + if err == nil { + err = model.ValidateTopUpQuotaCapacity(userId, creditedQuota) + } + if err != nil { + c.JSON(http.StatusOK, gin.H{"message": "error", "data": err.Error()}) + return true + } + return false +} + +func rejectInvalidTopUpQuota(c *gin.Context, userId int, amount int64) bool { + creditedQuota, err := validateTopUpQuota(amount) + if err == nil { + err = model.ValidateTopUpQuotaCapacity(userId, creditedQuota) + } + if err != nil { + c.JSON(http.StatusOK, gin.H{"message": "error", "data": err.Error()}) + return true + } + return false +} + func RequestEpay(c *gin.Context) { var req EpayRequest err := c.ShouldBindJSON(&req) @@ -197,8 +278,11 @@ func RequestEpay(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"message": "error", "data": fmt.Sprintf("充值数量不能小于 %d", getMinTopup())}) return } - id := c.GetInt("id") + if rejectInvalidTopUpQuota(c, id, req.Amount) { + return + } + group, err := model.GetUserGroup(id, true) if err != nil { c.JSON(http.StatusOK, gin.H{"message": "error", "data": "获取用户分组失败"}) @@ -351,16 +435,9 @@ func EpayNotify(c *gin.Context) { return } verifyInfo, err := client.Verify(params) - if err == nil && verifyInfo.VerifyStatus { - logger.LogInfo(c.Request.Context(), fmt.Sprintf("易支付 webhook 验签成功 trade_no=%s callback_type=%s trade_status=%s client_ip=%s verify_info=%q", verifyInfo.ServiceTradeNo, verifyInfo.Type, verifyInfo.TradeStatus, c.ClientIP(), common.GetJsonString(verifyInfo))) - _, err := c.Writer.Write([]byte("success")) - if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("易支付 webhook 响应写入失败 trade_no=%s client_ip=%s error=%q", verifyInfo.ServiceTradeNo, c.ClientIP(), err.Error())) - } - } else { - _, err := c.Writer.Write([]byte("fail")) - if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("易支付 webhook 响应写入失败 path=%q client_ip=%s error=%q", c.Request.RequestURI, c.ClientIP(), err.Error())) + if err != nil || !verifyInfo.VerifyStatus { + if _, writeErr := c.Writer.Write([]byte("fail")); writeErr != nil { + logger.LogError(c.Request.Context(), fmt.Sprintf("易支付 webhook 响应写入失败 path=%q client_ip=%s error=%q", c.Request.RequestURI, c.ClientIP(), writeErr.Error())) } if err != nil { logger.LogWarn(c.Request.Context(), fmt.Sprintf("易支付 webhook 验签失败 path=%q client_ip=%s verify_error=%q", c.Request.RequestURI, c.ClientIP(), err.Error())) @@ -369,46 +446,41 @@ func EpayNotify(c *gin.Context) { } return } + logger.LogInfo(c.Request.Context(), fmt.Sprintf("易支付 webhook 验签成功 trade_no=%s callback_type=%s trade_status=%s client_ip=%s verify_info=%q", verifyInfo.ServiceTradeNo, verifyInfo.Type, verifyInfo.TradeStatus, c.ClientIP(), common.GetJsonString(verifyInfo))) if verifyInfo.TradeStatus == epay.StatusTradeSuccess { + // 进程内锁只是优化;重复/并发回调的正确性由 RechargeEpay 的 + // 数据库行锁 + 事务内状态校验保证(多实例部署下同样安全)。 LockOrder(verifyInfo.ServiceTradeNo) defer UnlockOrder(verifyInfo.ServiceTradeNo) - topUp := model.GetTopUpByTradeNo(verifyInfo.ServiceTradeNo) - if topUp == nil { - logger.LogWarn(c.Request.Context(), fmt.Sprintf("易支付 回调订单不存在 trade_no=%s callback_type=%s client_ip=%s verify_info=%q", verifyInfo.ServiceTradeNo, verifyInfo.Type, c.ClientIP(), common.GetJsonString(verifyInfo))) - return - } - if topUp.PaymentProvider != model.PaymentProviderEpay { - logger.LogWarn(c.Request.Context(), fmt.Sprintf("易支付 订单支付网关不匹配 trade_no=%s order_provider=%s callback_type=%s client_ip=%s", verifyInfo.ServiceTradeNo, topUp.PaymentProvider, verifyInfo.Type, c.ClientIP())) - return - } - if topUp.Status == common.TopUpStatusPending { - if topUp.PaymentMethod != verifyInfo.Type { - logger.LogInfo(c.Request.Context(), fmt.Sprintf("易支付 实际支付方式与订单不同 trade_no=%s order_payment_method=%s actual_type=%s client_ip=%s", verifyInfo.ServiceTradeNo, topUp.PaymentMethod, verifyInfo.Type, c.ClientIP())) - topUp.PaymentMethod = verifyInfo.Type - } - topUp.Status = common.TopUpStatusSuccess - err := topUp.Update() - if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("易支付 更新充值订单失败 trade_no=%s user_id=%d client_ip=%s error=%q topup=%q", topUp.TradeNo, topUp.UserId, c.ClientIP(), err.Error(), common.GetJsonString(topUp))) - return + alreadyDone, err := model.RechargeEpay(verifyInfo.ServiceTradeNo, verifyInfo.Type, c.ClientIP()) + if err != nil { + switch { + case errors.Is(err, model.ErrTopUpNotFound): + logger.LogWarn(c.Request.Context(), fmt.Sprintf("易支付 回调订单不存在 trade_no=%s callback_type=%s client_ip=%s verify_info=%q", verifyInfo.ServiceTradeNo, verifyInfo.Type, c.ClientIP(), common.GetJsonString(verifyInfo))) + case errors.Is(err, model.ErrPaymentMethodMismatch): + logger.LogWarn(c.Request.Context(), fmt.Sprintf("易支付 订单支付网关不匹配 trade_no=%s callback_type=%s client_ip=%s", verifyInfo.ServiceTradeNo, verifyInfo.Type, c.ClientIP())) + case errors.Is(err, model.ErrTopUpStatusInvalid): + logger.LogWarn(c.Request.Context(), fmt.Sprintf("易支付 订单状态非法 trade_no=%s callback_type=%s client_ip=%s", verifyInfo.ServiceTradeNo, verifyInfo.Type, c.ClientIP())) + default: + logger.LogError(c.Request.Context(), fmt.Sprintf("易支付 充值处理失败 trade_no=%s client_ip=%s error=%q", verifyInfo.ServiceTradeNo, c.ClientIP(), err.Error())) } - //user, _ := model.GetUserById(topUp.UserId, false) - //user.Quota += topUp.Amount * 500000 - dAmount := decimal.NewFromInt(int64(topUp.Amount)) - dQuotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit) - quotaToAdd := int(dAmount.Mul(dQuotaPerUnit).IntPart()) - err = model.IncreaseUserQuota(topUp.UserId, quotaToAdd, true) - if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("易支付 更新用户额度失败 trade_no=%s user_id=%d client_ip=%s quota_to_add=%d error=%q topup=%q", topUp.TradeNo, topUp.UserId, c.ClientIP(), quotaToAdd, err.Error(), common.GetJsonString(topUp))) - return + if _, writeErr := c.Writer.Write([]byte("fail")); writeErr != nil { + logger.LogError(c.Request.Context(), fmt.Sprintf("易支付 webhook 响应写入失败 trade_no=%s client_ip=%s error=%q", verifyInfo.ServiceTradeNo, c.ClientIP(), writeErr.Error())) } - logger.LogInfo(c.Request.Context(), fmt.Sprintf("易支付 充值成功 trade_no=%s user_id=%d client_ip=%s quota_to_add=%d money=%.2f topup=%q", topUp.TradeNo, topUp.UserId, c.ClientIP(), quotaToAdd, topUp.Money, common.GetJsonString(topUp))) - model.RecordTopupLog(topUp.UserId, fmt.Sprintf("使用在线充值成功,充值金额: %v,支付金额:%f", logger.LogQuota(quotaToAdd), topUp.Money), c.ClientIP(), topUp.PaymentMethod, "epay") + return + } + if alreadyDone { + logger.LogInfo(c.Request.Context(), fmt.Sprintf("易支付 重复回调幂等忽略 trade_no=%s callback_type=%s client_ip=%s", verifyInfo.ServiceTradeNo, verifyInfo.Type, c.ClientIP())) + } else { + logger.LogInfo(c.Request.Context(), fmt.Sprintf("易支付 充值成功 trade_no=%s callback_type=%s client_ip=%s", verifyInfo.ServiceTradeNo, verifyInfo.Type, c.ClientIP())) } } else { logger.LogInfo(c.Request.Context(), fmt.Sprintf("易支付 webhook 忽略事件 trade_no=%s callback_type=%s trade_status=%s client_ip=%s verify_info=%q", verifyInfo.ServiceTradeNo, verifyInfo.Type, verifyInfo.TradeStatus, c.ClientIP(), common.GetJsonString(verifyInfo))) } + if _, writeErr := c.Writer.Write([]byte("success")); writeErr != nil { + logger.LogError(c.Request.Context(), fmt.Sprintf("易支付 webhook 响应写入失败 trade_no=%s client_ip=%s error=%q", verifyInfo.ServiceTradeNo, c.ClientIP(), writeErr.Error())) + } } func RequestAmount(c *gin.Context) { @@ -424,6 +496,9 @@ func RequestAmount(c *gin.Context) { return } id := c.GetInt("id") + if rejectInvalidTopUpQuota(c, id, req.Amount) { + return + } group, err := model.GetUserGroup(id, true) if err != nil { c.JSON(http.StatusOK, gin.H{"message": "error", "data": "获取用户分组失败"}) diff --git a/controller/topup_creem.go b/controller/topup_creem.go index 7472690e22fb..52696acc6723 100644 --- a/controller/topup_creem.go +++ b/controller/topup_creem.go @@ -18,6 +18,7 @@ import ( "time" "github.com/gin-gonic/gin" + "github.com/shopspring/decimal" "github.com/thanhpk/randstr" ) @@ -98,7 +99,15 @@ func (*CreemAdaptor) RequestPay(c *gin.Context, req *CreemPayRequest) { } id := c.GetInt("id") - user, _ := model.GetUserById(id, false) + if rejectInvalidCreditedQuota(c, id, decimal.NewFromInt(selectedProduct.Quota)) { + return + } + + user, err := model.GetUserById(id, false) + if err != nil || user == nil { + c.JSON(http.StatusOK, gin.H{"message": "error", "data": "用户不存在"}) + return + } // 生成唯一的订单引用ID reference := fmt.Sprintf("creem-api-ref-%d-%d-%s", user.Id, time.Now().UnixMilli(), randstr.String(4)) diff --git a/controller/topup_quota_limit_test.go b/controller/topup_quota_limit_test.go new file mode 100644 index 000000000000..710e87a31556 --- /dev/null +++ b/controller/topup_quota_limit_test.go @@ -0,0 +1,203 @@ +package controller + +import ( + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/shopspring/decimal" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestTopUpQuotaValidation(t *testing.T) { + oldQuotaPerUnit := common.QuotaPerUnit + oldDisplayType := operation_setting.GetGeneralSetting().QuotaDisplayType + common.QuotaPerUnit = 500000 + t.Cleanup(func() { + common.QuotaPerUnit = oldQuotaPerUnit + operation_setting.GetGeneralSetting().QuotaDisplayType = oldDisplayType + }) + + testCases := []struct { + name string + displayType string + amount int64 + wantQuota int + wantErr bool + }{ + { + name: "currency amount below limit", + displayType: operation_setting.QuotaDisplayTypeUSD, + amount: 4294, + wantQuota: 2_147_000_000, + }, + { + name: "currency amount above limit", + displayType: operation_setting.QuotaDisplayTypeUSD, + amount: 4295, + wantQuota: 2_147_500_000, + }, + { + name: "token amount preserves settlement truncation", + displayType: operation_setting.QuotaDisplayTypeTokens, + amount: 2_147_500_000, + wantQuota: 2_147_500_000, + }, + { + name: "token amount above legacy int32 range", + displayType: operation_setting.QuotaDisplayTypeTokens, + amount: 4_294_500_000, + wantQuota: 4_294_500_000, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + operation_setting.GetGeneralSetting().QuotaDisplayType = tc.displayType + quota, err := getTopUpQuota(tc.amount) + if tc.wantErr { + require.Error(t, err) + return + } + require.NoError(t, err) + assert.Equal(t, tc.wantQuota, quota) + }) + } +} + +func TestValidateTopUpQuotaReturnsMaximumAmount(t *testing.T) { + oldQuotaPerUnit := common.QuotaPerUnit + oldDisplayType := operation_setting.GetGeneralSetting().QuotaDisplayType + common.QuotaPerUnit = 500000 + operation_setting.GetGeneralSetting().QuotaDisplayType = operation_setting.QuotaDisplayTypeUSD + t.Cleanup(func() { + common.QuotaPerUnit = oldQuotaPerUnit + operation_setting.GetGeneralSetting().QuotaDisplayType = oldDisplayType + }) + + maxAmount := decimal.NewFromInt(common.MaxWalletQuota). + Div(decimal.NewFromFloat(common.QuotaPerUnit)). + Floor().IntPart() + + _, err := validateTopUpQuota(maxAmount) + require.NoError(t, err) + _, err = validateTopUpQuota(maxAmount + 1) + require.EqualError(t, err, fmt.Sprintf("单笔充值数量不能大于 %d", maxAmount)) +} + +func TestRequestAmountRejectsTopUpThatCannotBeSettled(t *testing.T) { + oldQuotaPerUnit := common.QuotaPerUnit + oldDisplayType := operation_setting.GetGeneralSetting().QuotaDisplayType + common.QuotaPerUnit = 500000 + operation_setting.GetGeneralSetting().QuotaDisplayType = operation_setting.QuotaDisplayTypeUSD + t.Cleanup(func() { + common.QuotaPerUnit = oldQuotaPerUnit + operation_setting.GetGeneralSetting().QuotaDisplayType = oldDisplayType + }) + + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + maxAmount := decimal.NewFromInt(common.MaxWalletQuota). + Div(decimal.NewFromFloat(common.QuotaPerUnit)). + Floor().IntPart() + ctx.Request = httptest.NewRequest( + http.MethodPost, + "/api/user/amount", + strings.NewReader(fmt.Sprintf(`{"amount":%d}`, maxAmount+1)), + ) + ctx.Request.Header.Set("Content-Type", "application/json") + + RequestAmount(ctx) + + assert.Equal(t, http.StatusOK, recorder.Code) + assert.JSONEq(t, fmt.Sprintf(`{"message":"error","data":"单笔充值数量不能大于 %d"}`, maxAmount), recorder.Body.String()) +} + +func TestRequestAmountRejectsTopUpThatWouldOverflowWallet(t *testing.T) { + oldQuotaPerUnit := common.QuotaPerUnit + oldDisplayType := operation_setting.GetGeneralSetting().QuotaDisplayType + oldDB := model.DB + common.QuotaPerUnit = 500000 + operation_setting.GetGeneralSetting().QuotaDisplayType = operation_setting.QuotaDisplayTypeUSD + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.User{})) + model.DB = db + t.Cleanup(func() { + common.QuotaPerUnit = oldQuotaPerUnit + operation_setting.GetGeneralSetting().QuotaDisplayType = oldDisplayType + model.DB = oldDB + sqlDB, dbErr := db.DB() + if dbErr == nil { + require.NoError(t, sqlDB.Close()) + } + }) + + require.NoError(t, model.DB.Create(&model.User{ + Id: 42, + Username: "topup_capacity_user", + Quota: common.MaxWalletQuota - 100_000, + Status: common.UserStatusEnabled, + }).Error) + + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Set("id", 42) + ctx.Request = httptest.NewRequest( + http.MethodPost, + "/api/user/amount", + strings.NewReader(`{"amount":1}`), + ) + ctx.Request.Header.Set("Content-Type", "application/json") + + RequestAmount(ctx) + + assert.Equal(t, http.StatusOK, recorder.Code) + assert.JSONEq(t, `{"message":"error","data":"top-up quota limit exceeded"}`, recorder.Body.String()) +} + +func TestValidateCreditedQuotaRejectsOverflow(t *testing.T) { + _, err := validateCreditedQuota(decimal.NewFromInt(int64(common.MaxWalletQuota / 2))) + require.NoError(t, err) + _, err = validateCreditedQuota(decimal.Zero) + require.EqualError(t, err, "充值额度必须大于 0") + _, err = validateCreditedQuota(decimal.NewFromInt(common.MaxWalletQuota + 1)) + require.EqualError( + t, + err, + "充值额度超出系统可表示范围", + ) +} + +func TestStripeCreditedQuotaIncludesGroupRatio(t *testing.T) { + oldQuotaPerUnit := common.QuotaPerUnit + oldTopupGroupRatio := common.TopupGroupRatio2JSONString() + common.QuotaPerUnit = 500000 + require.NoError(t, common.UpdateTopupGroupRatioByJSONString(`{"vip":2}`)) + t.Cleanup(func() { + common.QuotaPerUnit = oldQuotaPerUnit + require.NoError(t, common.UpdateTopupGroupRatioByJSONString(oldTopupGroupRatio)) + }) + + _, err := validateCreditedQuota(getStripeCreditedQuota(2147, "vip")) + require.NoError(t, err) + _, err = validateCreditedQuota(getStripeCreditedQuota(2148, "vip")) + require.NoError(t, err) + _, err = validateCreditedQuota(getStripeCreditedQuota(int64(common.MaxWalletQuota), "vip")) + require.Error(t, err) + + require.NoError(t, common.UpdateTopupGroupRatioByJSONString(`{"free":0}`)) + assert.True(t, decimal.NewFromInt(500000).Equal(getStripeCreditedQuota(1, "free"))) +} diff --git a/controller/topup_stripe.go b/controller/topup_stripe.go index 8a39576659e6..268b4046eac1 100644 --- a/controller/topup_stripe.go +++ b/controller/topup_stripe.go @@ -17,6 +17,7 @@ import ( "github.com/QuantumNous/new-api/setting/operation_setting" "github.com/gin-gonic/gin" + "github.com/shopspring/decimal" "github.com/stripe/stripe-go/v81" "github.com/stripe/stripe-go/v81/checkout/session" "github.com/stripe/stripe-go/v81/webhook" @@ -47,12 +48,19 @@ func (*StripeAdaptor) RequestAmount(c *gin.Context, req *StripePayRequest) { c.JSON(http.StatusOK, gin.H{"message": "error", "data": fmt.Sprintf("充值数量不能小于 %d", getStripeMinTopup())}) return } + if req.Amount > 10000 { + c.JSON(http.StatusOK, gin.H{"message": "error", "data": "充值数量不能大于 10000"}) + return + } id := c.GetInt("id") group, err := model.GetUserGroup(id, true) if err != nil { c.JSON(http.StatusOK, gin.H{"message": "error", "data": "获取用户分组失败"}) return } + if rejectInvalidCreditedQuota(c, id, getStripeCreditedQuota(req.Amount, group)) { + return + } payMoney := getStripePayMoney(float64(req.Amount), group) if payMoney <= 0.01 { c.JSON(http.StatusOK, gin.H{"message": "error", "data": "充值金额过低"}) @@ -86,8 +94,17 @@ func (*StripeAdaptor) RequestPay(c *gin.Context, req *StripePayRequest) { } id := c.GetInt("id") - user, _ := model.GetUserById(id, false) + user, err := model.GetUserById(id, false) + if err != nil || user == nil { + c.JSON(http.StatusOK, gin.H{"message": "error", "data": "用户不存在"}) + return + } chargedMoney := GetChargedAmount(float64(req.Amount), *user) + if rejectInvalidCreditedQuota(c, id, + decimal.NewFromFloat(chargedMoney).Mul(decimal.NewFromFloat(common.QuotaPerUnit)), + ) { + return + } reference := fmt.Sprintf("new-api-ref-%d-%d-%s", user.Id, time.Now().UnixMilli(), randstr.String(4)) referenceId := "ref_" + common.Sha1([]byte(reference)) @@ -394,6 +411,16 @@ func GetChargedAmount(count float64, user model.User) float64 { return count * topUpGroupRatio } +func getStripeCreditedQuota(amount int64, group string) decimal.Decimal { + topUpGroupRatio := common.GetTopupGroupRatio(group) + if topUpGroupRatio == 0 { + topUpGroupRatio = 1 + } + return decimal.NewFromInt(amount). + Mul(decimal.NewFromFloat(topUpGroupRatio)). + Mul(decimal.NewFromFloat(common.QuotaPerUnit)) +} + func getStripePayMoney(amount float64, group string) float64 { originalAmount := amount if operation_setting.GetQuotaDisplayType() == operation_setting.QuotaDisplayTypeTokens { diff --git a/controller/topup_waffo.go b/controller/topup_waffo.go index 4ac3b2b5ddd7..d803c881ffca 100644 --- a/controller/topup_waffo.go +++ b/controller/topup_waffo.go @@ -123,8 +123,11 @@ func RequestWaffoAmount(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"message": "error", "data": fmt.Sprintf("充值数量不能小于 %d", waffoMinTopup)}) return } - id := c.GetInt("id") + if rejectInvalidTopUpQuota(c, id, req.Amount) { + return + } + group, err := model.GetUserGroup(id, true) if err != nil { c.JSON(http.StatusOK, gin.H{"message": "error", "data": "获取用户分组失败"}) @@ -157,8 +160,11 @@ func RequestWaffoPay(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"message": "error", "data": fmt.Sprintf("充值数量不能小于 %d", waffoMinTopup)}) return } - id := c.GetInt("id") + if rejectInvalidTopUpQuota(c, id, req.Amount) { + return + } + user, err := model.GetUserById(id, false) if err != nil || user == nil { c.JSON(http.StatusOK, gin.H{"message": "error", "data": "用户不存在"}) diff --git a/controller/topup_waffo_pancake.go b/controller/topup_waffo_pancake.go index beb73ebee681..b7366ec8f5db 100644 --- a/controller/topup_waffo_pancake.go +++ b/controller/topup_waffo_pancake.go @@ -33,8 +33,11 @@ func RequestWaffoPancakeAmount(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"message": "error", "data": fmt.Sprintf("充值数量不能小于 %d", setting.WaffoPancakeMinTopUp)}) return } - id := c.GetInt("id") + if rejectInvalidTopUpQuota(c, id, req.Amount) { + return + } + group, err := model.GetUserGroup(id, true) if err != nil { c.JSON(http.StatusOK, gin.H{"message": "error", "data": "获取用户分组失败"}) @@ -351,8 +354,11 @@ func RequestWaffoPancakePay(c *gin.Context) { c.JSON(http.StatusOK, gin.H{"message": "error", "data": fmt.Sprintf("充值数量不能小于 %d", setting.WaffoPancakeMinTopUp)}) return } - id := c.GetInt("id") + if rejectInvalidTopUpQuota(c, id, req.Amount) { + return + } + user, err := model.GetUserById(id, false) if err != nil || user == nil { c.JSON(http.StatusOK, gin.H{"message": "error", "data": "用户不存在"}) diff --git a/controller/user.go b/controller/user.go index 59aaf5847238..f2eb20a56356 100644 --- a/controller/user.go +++ b/controller/user.go @@ -28,8 +28,10 @@ import ( ) type LoginRequest struct { - Username string `json:"username"` - Password string `json:"password"` + Username string `json:"username"` + Password string `json:"password"` + PasswordEncrypted string `json:"password_encrypted"` + EncryptionKeyID string `json:"encryption_key_id"` } var ( @@ -37,6 +39,23 @@ var ( errOriginalPasswordFail = errors.New("original password is incorrect") ) +func GetPasswordEncryptionKey(c *gin.Context) { + if !common.PasswordLoginEncryptionEnabled { + common.ApiSuccess(c, gin.H{"enabled": false}) + return + } + keyID, publicKey := common.PasswordEncryptionPublicKey() + if keyID == "" || publicKey == "" { + common.ApiErrorI18n(c, i18n.MsgDatabaseError) + return + } + common.ApiSuccess(c, gin.H{ + "enabled": true, + "kid": keyID, + "public_key": publicKey, + }) +} + func Login(c *gin.Context) { if !common.PasswordLoginEnabled { common.ApiErrorI18n(c, i18n.MsgUserPasswordLoginDisabled) @@ -50,6 +69,17 @@ func Login(c *gin.Context) { } username := loginRequest.Username password := loginRequest.Password + if common.PasswordLoginEncryptionEnabled { + if loginRequest.PasswordEncrypted == "" || loginRequest.EncryptionKeyID == "" { + common.ApiErrorI18n(c, i18n.MsgInvalidParams) + return + } + password, err = common.DecryptPassword(loginRequest.PasswordEncrypted, loginRequest.EncryptionKeyID) + if err != nil { + common.ApiErrorI18n(c, i18n.MsgUserUsernameOrPasswordError) + return + } + } if username == "" || password == "" { common.ApiErrorI18n(c, i18n.MsgInvalidParams) return @@ -399,11 +429,6 @@ func GetUser(c *gin.Context) { func GenerateAccessToken(c *gin.Context) { id := c.GetInt("id") - user, err := model.GetUserById(id, true) - if err != nil { - common.ApiError(c, err) - return - } // get rand int 28-32 randI := common.GetRandomInt(4) key, err := common.GenerateRandomKey(29 + randI) @@ -412,14 +437,12 @@ func GenerateAccessToken(c *gin.Context) { common.SysLog("failed to generate key: " + err.Error()) return } - user.SetAccessToken(key) - - if model.DB.Where("access_token = ?", user.AccessToken).First(user).RowsAffected != 0 { + if model.DB.Where("access_token = ?", key).First(&model.User{}).RowsAffected != 0 { common.ApiErrorI18n(c, i18n.MsgUuidDuplicate) return } - if err := user.Update(false); err != nil { + if err := model.UpdateUserAccessToken(id, key); err != nil { common.ApiError(c, err) return } @@ -427,7 +450,7 @@ func GenerateAccessToken(c *gin.Context) { c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", - "data": user.AccessToken, + "data": key, }) return } @@ -1174,6 +1197,10 @@ func ManageUser(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgUserQuotaChangeZero) return } + if err := common.ValidateWalletQuota(req.Value); err != nil { + common.ApiError(c, err) + return + } if err := model.IncreaseUserQuota(user.Id, req.Value, true); err != nil { common.ApiError(c, err) return @@ -1194,6 +1221,10 @@ func ManageUser(c *gin.Context) { "quota": logger.LogQuota(req.Value), }) case "override": + if err := common.ValidateWalletQuota(req.Value); err != nil { + common.ApiError(c, err) + return + } oldQuota := user.Quota if err := model.DB.Model(&model.User{}).Where("id = ?", user.Id).Update("quota", req.Value).Error; err != nil { common.ApiError(c, err) diff --git a/controller/user_manage_test.go b/controller/user_manage_test.go index 1b52ece08835..985640786308 100644 --- a/controller/user_manage_test.go +++ b/controller/user_manage_test.go @@ -159,3 +159,24 @@ func TestManageUserDeleteReturnsImmediatelyAndUnknownActionFails(t *testing.T) { assert.EqualValues(t, 1, unchanged.AuthVersion) assert.Equal(t, common.UserStatusEnabled, unchanged.Status) } + +func TestManageUserQuotaRespectsWalletCeiling(t *testing.T) { + db := setupManageUserTestDB(t) + user := model.User{ + Username: "managed-quota-user", Password: "password", Role: common.RoleCommonUser, + Status: common.UserStatusEnabled, Group: "default", Quota: common.MaxWalletQuota - 1, + } + require.NoError(t, db.Create(&user).Error) + + recorder := performManageUserRequest(t, fmt.Sprintf(`{"id":%d,"action":"add_quota","mode":"add","value":2}`, user.Id)) + assert.Contains(t, recorder.Body.String(), `"success":false`) + + var updated model.User + require.NoError(t, db.First(&updated, user.Id).Error) + assert.Equal(t, common.MaxWalletQuota-1, updated.Quota) + + recorder = performManageUserRequest(t, fmt.Sprintf(`{"id":%d,"action":"add_quota","mode":"override","value":%d}`, user.Id, common.MaxWalletQuota+1)) + assert.Contains(t, recorder.Body.String(), `"success":false`) + require.NoError(t, db.First(&updated, user.Id).Error) + assert.Equal(t, common.MaxWalletQuota-1, updated.Quota) +} diff --git a/controller/video_proxy.go b/controller/video_proxy.go index 996d084d88fa..230d2bddfd38 100644 --- a/controller/video_proxy.go +++ b/controller/video_proxy.go @@ -1,27 +1,56 @@ package controller import ( + "bytes" "context" "encoding/base64" + "errors" "fmt" "io" + "net" "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/logger" "github.com/QuantumNous/new-api/model" + relaychannel "github.com/QuantumNous/new-api/relay/channel" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting/system_setting" - "github.com/gin-gonic/gin" + "golang.org/x/net/http/httpguts" ) +var errTaskMediaRequestRejected = errors.New("task media request rejected") + +var taskMediaResponseHeaderTimeout = 60 * time.Second +var taskMediaDataURLMaxEncodedBytes = 64 << 20 + +type taskMediaProxyError struct { + status int + code string + message string + err error +} + +func (e *taskMediaProxyError) Error() string { + if e.err == nil { + return e.message + } + return e.message + ": " + e.err.Error() +} + +func (e *taskMediaProxyError) Unwrap() error { + return e.err +} + // videoProxyError returns a standardized OpenAI-style error response. func videoProxyError(c *gin.Context, status int, errType, message string) { + c.Header("Cache-Control", "private, no-store") c.JSON(status, gin.H{ "error": gin.H{ "message": message, @@ -37,8 +66,7 @@ func VideoProxy(c *gin.Context) { return } - userID := c.GetInt("id") - task, exists, err := model.GetByTaskId(userID, taskID) + task, exists, err := getTaskForArtifactRequest(c, taskID) if err != nil { logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to query task %s: %s", taskID, err.Error())) videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to query task") @@ -48,141 +76,505 @@ func VideoProxy(c *gin.Context) { videoProxyError(c, http.StatusNotFound, "invalid_request_error", "Task not found") return } - if task.Status != model.TaskStatusSuccess { videoProxyError(c, http.StatusBadRequest, "invalid_request_error", fmt.Sprintf("Task is not completed yet, current status: %s", task.Status)) return } + var descriptor *relaychannel.TaskContentRequest + if taskHasPluginExecution(task) { + artifacts, projectionErr := projectTaskArtifacts(task) + if projectionErr == nil { + for _, artifact := range artifacts { + if artifact.Type != "video" { + continue + } + adaptor, adaptorErr := initTaskArtifactAdaptor(task) + if adaptorErr == nil { + if provider, ok := adaptor.(relaychannel.TaskContentRequestProvider); ok { + descriptor, adaptorErr = provider.BuildContentRequest(task, artifact.Key, relaychannel.TaskArtifactClientRequest{ + Method: c.Request.Method, + Headers: taskArtifactClientHeaders(c.Request.Header), + }) + } + } + if adaptorErr != nil { + logger.LogWarn(c.Request.Context(), fmt.Sprintf("Failed to resolve plugin video content for task %s", taskID)) + descriptor = nil + } + break + } + } else { + logger.LogWarn(c.Request.Context(), fmt.Sprintf("Failed to project plugin video for task %s", taskID)) + } + } + if descriptor == nil { + resultURL := task.GetResultURL() + if isTaskMediaFallbackLoop(resultURL, task.TaskID) { + writeTaskMediaProxyError(c, &taskMediaProxyError{ + status: http.StatusGone, code: "artifact_gone", + message: "Artifact content is no longer available", + }) + return + } + descriptor = &relaychannel.TaskContentRequest{ + URL: resultURL, + Method: c.Request.Method, + Credentialless: true, + } + } + if err := proxyTaskMedia(c, task, descriptor); err != nil { + writeTaskMediaProxyError(c, err) + } +} + +func proxyTaskMedia(c *gin.Context, task *model.Task, descriptor *relaychannel.TaskContentRequest) error { + if descriptor == nil { + return &taskMediaProxyError{ + status: http.StatusInternalServerError, code: "artifact_plugin_error", + message: "Artifact content plugin returned no request", + } + } + rawURL := strings.TrimSpace(descriptor.URL) + if rawURL == "" { + return &taskMediaProxyError{ + status: http.StatusGone, code: "artifact_gone", + message: "Artifact content is no longer available", + } + } + if strings.HasPrefix(rawURL, "data:") { + if len(rawURL) > taskMediaDataURLMaxEncodedBytes { + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_request_rejected", + message: "Artifact request was rejected", err: errTaskMediaRequestRejected, + } + } + if err := writeVideoDataURL(c, rawURL); err != nil { + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_upstream_error", + message: "Failed to decode artifact content", err: err, + } + } + return nil + } + if len(rawURL) > 64<<10 { + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_request_rejected", + message: "Artifact request was rejected", err: errTaskMediaRequestRejected, + } + } + + parsedURL, err := url.Parse(rawURL) + if err != nil || parsedURL == nil || (parsedURL.Scheme != "http" && parsedURL.Scheme != "https") || + parsedURL.Host == "" || parsedURL.User != nil || parsedURL.Fragment != "" { + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_request_rejected", + message: "Artifact request was rejected", err: errTaskMediaRequestRejected, + } + } + if isTaskMediaFallbackLoop(rawURL, task.TaskID) || isSelfTaskMediaURL(c, parsedURL) { + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_request_rejected", + message: "Artifact proxy loop was rejected", err: errTaskMediaRequestRejected, + } + } + + method := strings.ToUpper(strings.TrimSpace(descriptor.Method)) + if method == "" { + method = c.Request.Method + } + switch method { + case http.MethodGet, http.MethodHead, http.MethodPost: + default: + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_request_rejected", + message: "Artifact request method was rejected", err: errTaskMediaRequestRejected, + } + } + if len(descriptor.Body) > 1<<20 { + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_request_rejected", + message: "Artifact request body was rejected", err: errTaskMediaRequestRejected, + } + } + if descriptor.Credentialless && + (method != http.MethodGet && method != http.MethodHead || + descriptor.Body != nil || len(descriptor.Headers) != 0) { + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_request_rejected", + message: "Credentialless artifact request was rejected", err: errTaskMediaRequestRejected, + } + } + channel, err := model.CacheGetChannel(task.ChannelId) if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to get channel for task %s: %s", taskID, err.Error())) - videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to retrieve channel information") - return + return &taskMediaProxyError{ + status: http.StatusServiceUnavailable, code: "artifact_plugin_unavailable", + message: "Artifact channel is unavailable", err: err, + } } - baseURL := channel.GetBaseURL() - if baseURL == "" { - baseURL = "https://api.openai.com" + proxy := strings.TrimSpace(channel.GetSetting().Proxy) + if err := validateTaskMediaURL(rawURL, proxy); err != nil { + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_request_rejected", + message: "Artifact request was rejected", err: err, + } } - var videoURL string - proxy := channel.GetSetting().Proxy client := service.GetSSRFProtectedHTTPClient() if proxy != "" { - // 渠道代理路径的连接由代理侧建立,无法做拨号时逐 IP 校验, - // 因此后面对 videoURL 保留请求前的一次性 SSRF 校验。 client, err = service.GetHttpClientWithProxy(proxy) if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to create proxy client for task %s: %s", taskID, err.Error())) - videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy client") - return + return &taskMediaProxyError{ + status: http.StatusInternalServerError, code: "artifact_internal_error", + message: "Failed to create artifact proxy client", err: err, + } } } + if client == nil { + client = http.DefaultClient + } - ctx, cancel := context.WithTimeout(c.Request.Context(), 60*time.Second) - defer cancel() - req, err := http.NewRequestWithContext(ctx, http.MethodGet, "", nil) + req, err := http.NewRequestWithContext(c.Request.Context(), method, parsedURL.String(), bytes.NewReader(descriptor.Body)) if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to create request: %s", err.Error())) - videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy request") - return + return &taskMediaProxyError{ + status: http.StatusInternalServerError, code: "artifact_internal_error", + message: "Failed to create artifact request", err: err, + } + } + if err := applyTaskMediaRequestHeaders(req.Header, descriptor.Headers); err != nil { + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_request_rejected", + message: "Artifact request headers were rejected", err: err, + } + } + clientHeaders := taskArtifactClientHeaders(c.Request.Header) + for name, value := range clientHeaders { + req.Header.Set(name, value) } - switch channel.Type { - case constant.ChannelTypeGemini: - apiKey := task.PrivateData.Key - if apiKey == "" { - logger.LogError(c.Request.Context(), fmt.Sprintf("Missing stored API key for Gemini task %s", taskID)) - videoProxyError(c, http.StatusInternalServerError, "server_error", "API key not stored for task") - return + client = taskMediaRedirectClient(client, proxy, c, clientHeaders, descriptor.Credentialless) + clientWithoutBodyTimeout := *client + clientWithoutBodyTimeout.Timeout = 0 + resp, err := doTaskMediaRequest(&clientWithoutBodyTimeout, req, taskMediaResponseHeaderTimeout) + if err != nil { + if errors.Is(err, errTaskMediaRequestRejected) { + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_request_rejected", + message: "Artifact redirect was rejected", err: err, + } } - videoURL, err = getGeminiVideoURL(channel, task, apiKey) - if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to resolve Gemini video URL for task %s: %s", taskID, err.Error())) - videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to resolve Gemini video URL") - return + var netErr net.Error + if errors.Is(err, context.DeadlineExceeded) || errors.As(err, &netErr) && netErr.Timeout() { + return &taskMediaProxyError{ + status: http.StatusGatewayTimeout, code: "artifact_upstream_timeout", + message: "Artifact upstream request timed out", err: err, + } } - req.Header.Set("x-goog-api-key", apiKey) - case constant.ChannelTypeVertexAi: - videoURL, err = getVertexVideoURL(channel, task) - if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to resolve Vertex video URL for task %s: %s", taskID, err.Error())) - videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to resolve Vertex video URL") - return + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_upstream_error", + message: "Failed to fetch artifact content", err: err, + } + } + defer resp.Body.Close() + + switch resp.StatusCode { + case http.StatusOK, http.StatusPartialContent, http.StatusNotModified, http.StatusRequestedRangeNotSatisfiable: + copyTaskMediaResponseHeaders(c.Writer.Header(), resp.Header) + setTaskMediaResponseSecurityHeaders(c.Writer.Header()) + c.Status(resp.StatusCode) + c.Writer.WriteHeaderNow() + if c.Request.Method == http.MethodHead || resp.StatusCode == http.StatusNotModified { + return nil + } + if _, err := io.Copy(c.Writer, resp.Body); err != nil { + logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to stream task media: %v", err)) + } + return nil + case http.StatusUnauthorized, http.StatusForbidden: + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_upstream_auth_failed", + message: "Artifact upstream authentication failed", + } + case http.StatusNotFound, http.StatusGone: + return &taskMediaProxyError{ + status: http.StatusGone, code: "artifact_gone", + message: "Artifact content is no longer available", + } + case http.StatusTooManyRequests: + if retryAfter := strings.TrimSpace(resp.Header.Get("Retry-After")); retryAfter != "" && + len(retryAfter) <= 256 && !strings.ContainsAny(retryAfter, "\r\n") { + c.Header("Retry-After", retryAfter) + } + return &taskMediaProxyError{ + status: http.StatusServiceUnavailable, code: "artifact_upstream_busy", + message: "Artifact upstream is busy", } - case constant.ChannelTypeOpenAI, constant.ChannelTypeSora: - videoURL = fmt.Sprintf("%s/v1/videos/%s/content", baseURL, task.GetUpstreamTaskID()) - req.Header.Set("Authorization", "Bearer "+channel.Key) default: - // Video URL is stored in PrivateData.ResultURL (fallback to FailReason for old data) - videoURL = task.GetResultURL() + return &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_upstream_error", + message: fmt.Sprintf("Artifact upstream returned status %d", resp.StatusCode), + } } +} - videoURL = strings.TrimSpace(videoURL) - if videoURL == "" { - logger.LogError(c.Request.Context(), fmt.Sprintf("Video URL is empty for task %s", taskID)) - videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to fetch video content") - return +type taskMediaHTTPResult struct { + response *http.Response + err error +} + +type taskMediaCancelBody struct { + io.ReadCloser + cancel context.CancelFunc + once sync.Once +} + +func (b *taskMediaCancelBody) Close() error { + b.once.Do(b.cancel) + return b.ReadCloser.Close() +} + +func doTaskMediaRequest(client *http.Client, request *http.Request, responseHeaderTimeout time.Duration) (*http.Response, error) { + if client == nil { + client = http.DefaultClient } + requestContext, cancel := context.WithCancel(request.Context()) + request = request.Clone(requestContext) + resultChannel := make(chan taskMediaHTTPResult, 1) + go func() { + response, err := client.Do(request) + resultChannel <- taskMediaHTTPResult{response: response, err: err} + }() - if strings.HasPrefix(videoURL, "data:") { - if err := writeVideoDataURL(c, videoURL); err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to decode video data URL for task %s: %s", taskID, err.Error())) - videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to fetch video content") + timer := time.NewTimer(responseHeaderTimeout) + defer timer.Stop() + cleanupResult := func() { + go func() { + result := <-resultChannel + if result.response != nil && result.response.Body != nil { + _ = result.response.Body.Close() + } + }() + } + + select { + case result := <-resultChannel: + if result.err != nil { + cancel() + if result.response != nil && result.response.Body != nil { + _ = result.response.Body.Close() + } + return nil, result.err } - return + if result.response == nil || result.response.Body == nil { + cancel() + return nil, errors.New("artifact upstream returned no response body") + } + result.response.Body = &taskMediaCancelBody{ + ReadCloser: result.response.Body, + cancel: cancel, + } + return result.response, nil + case <-timer.C: + cancel() + cleanupResult() + return nil, context.DeadlineExceeded + case <-request.Context().Done(): + cancel() + cleanupResult() + return nil, request.Context().Err() + } +} + +func applyTaskMediaRequestHeaders(destination http.Header, headers map[string]string) error { + if len(headers) > 64 { + return errTaskMediaRequestRejected } + for name, value := range headers { + name = strings.TrimSpace(name) + if !httpguts.ValidHeaderFieldName(name) || !httpguts.ValidHeaderFieldValue(value) || len(value) > 8192 { + return errTaskMediaRequestRejected + } + switch strings.ToLower(name) { + case "host", "content-length", "accept-encoding", "connection", "proxy-connection", "keep-alive", + "proxy-authorization", "te", "trailer", "transfer-encoding", "upgrade": + return errTaskMediaRequestRejected + } + destination.Set(name, value) + } + return nil +} - var validateErr error +func taskMediaRedirectClient(base *http.Client, proxy string, c *gin.Context, clientHeaders map[string]string, credentialless bool) *http.Client { + cloned := *base + cloned.CheckRedirect = func(req *http.Request, via []*http.Request) error { + if len(via) >= 10 { + return fmt.Errorf("%w: too many redirects", errTaskMediaRequestRejected) + } + if req.URL == nil || (req.URL.Scheme != "http" && req.URL.Scheme != "https") || + req.URL.Host == "" || req.URL.User != nil || req.URL.Fragment != "" { + return fmt.Errorf("%w: invalid redirect URL", errTaskMediaRequestRejected) + } + if err := validateTaskMediaURL(req.URL.String(), proxy); err != nil { + return fmt.Errorf("%w: %v", errTaskMediaRequestRejected, err) + } + if isSelfTaskMediaURL(c, req.URL) { + return fmt.Errorf("%w: proxy loop", errTaskMediaRequestRejected) + } + if len(via) > 0 && !sameTaskMediaOrigin(via[len(via)-1].URL, req.URL) { + if !credentialless { + return fmt.Errorf("%w: credentialed cross-origin redirect", errTaskMediaRequestRejected) + } + for name := range req.Header { + req.Header.Del(name) + } + req.Body = http.NoBody + req.GetBody = nil + req.ContentLength = 0 + } + for name, value := range clientHeaders { + req.Header.Set(name, value) + } + return nil + } + return &cloned +} + +func validateTaskMediaURL(rawURL, proxy string) error { if proxy == "" { - validateErr = service.ValidateSSRFProtectedFetchURL(videoURL) - } else { - fetchSetting := system_setting.GetFetchSetting() - validateErr = common.ValidateURLWithFetchSetting(videoURL, fetchSetting.EnableSSRFProtection, fetchSetting.AllowPrivateIp, fetchSetting.DomainFilterMode, fetchSetting.IpFilterMode, fetchSetting.DomainList, fetchSetting.IpList, fetchSetting.AllowedPorts, fetchSetting.ApplyIPFilterForDomain) - } - if validateErr != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("Video URL blocked for task %s: %v", taskID, validateErr)) - videoProxyError(c, http.StatusForbidden, "server_error", fmt.Sprintf("request blocked: %v", validateErr)) - return + return service.ValidateSSRFProtectedFetchURL(rawURL) } + fetchSetting := system_setting.GetFetchSetting() + return common.ValidateURLWithFetchSetting( + rawURL, + fetchSetting.EnableSSRFProtection, + fetchSetting.AllowPrivateIp, + fetchSetting.DomainFilterMode, + fetchSetting.IpFilterMode, + fetchSetting.DomainList, + fetchSetting.IpList, + fetchSetting.AllowedPorts, + fetchSetting.ApplyIPFilterForDomain, + ) +} - req.URL, err = url.Parse(videoURL) - if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to parse URL %s: %s", videoURL, err.Error())) - videoProxyError(c, http.StatusInternalServerError, "server_error", "Failed to create proxy request") - return +func sameTaskMediaOrigin(left, right *url.URL) bool { + if left == nil || right == nil { + return false } + return strings.EqualFold(left.Scheme, right.Scheme) && + strings.EqualFold(normalizeTaskMediaHost(left.Scheme, left.Host), normalizeTaskMediaHost(right.Scheme, right.Host)) +} - resp, err := client.Do(req) +func normalizeTaskMediaHost(scheme, host string) string { + host = strings.ToLower(strings.TrimSpace(host)) + hostname, port, err := net.SplitHostPort(host) if err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to fetch video from %s: %s", videoURL, err.Error())) - videoProxyError(c, http.StatusBadGateway, "server_error", "Failed to fetch video content") - return + return strings.TrimSuffix(host, ".") } - defer resp.Body.Close() + hostname = strings.TrimSuffix(strings.ToLower(hostname), ".") + if (strings.EqualFold(scheme, "http") && port == "80") || (strings.EqualFold(scheme, "https") && port == "443") { + return hostname + } + return net.JoinHostPort(hostname, port) +} - if resp.StatusCode != http.StatusOK { - logger.LogError(c.Request.Context(), fmt.Sprintf("Upstream returned status %d for %s", resp.StatusCode, videoURL)) - videoProxyError(c, http.StatusBadGateway, "server_error", - fmt.Sprintf("Upstream service returned status %d", resp.StatusCode)) - return +func isSelfTaskMediaURL(c *gin.Context, target *url.URL) bool { + if c == nil || target == nil || !isTaskMediaProxyPath(target.Path) { + return false + } + targetHost := normalizeTaskMediaHost(target.Scheme, target.Host) + if targetHost == "" { + return true + } + scheme := strings.TrimSpace(strings.Split(c.Request.Header.Get("X-Forwarded-Proto"), ",")[0]) + if scheme == "" { + scheme = "http" + if c.Request.TLS != nil { + scheme = "https" + } + } + hosts := []string{c.Request.Host} + if forwardedHost := strings.TrimSpace(strings.Split(c.Request.Header.Get("X-Forwarded-Host"), ",")[0]); forwardedHost != "" { + hosts = append(hosts, forwardedHost) + } + for _, host := range hosts { + if strings.EqualFold(targetHost, normalizeTaskMediaHost(scheme, host)) { + return true + } + } + return false +} + +func isTaskMediaProxyPath(path string) bool { + if strings.HasPrefix(path, "/v1/videos/") && strings.HasSuffix(path, "/content") { + return true + } + return strings.HasPrefix(path, "/v1/tasks/") && + strings.Contains(path, "/artifacts/") && + strings.HasSuffix(path, "/content") +} + +func isTaskMediaFallbackLoop(rawURL, taskID string) bool { + parsedURL, err := url.Parse(strings.TrimSpace(rawURL)) + if err != nil || parsedURL == nil { + return false + } + path, err := url.PathUnescape(parsedURL.EscapedPath()) + if err != nil { + path = parsedURL.Path } + if path == "/v1/videos/"+taskID+"/content" { + return true + } + artifactPrefix := "/v1/tasks/" + taskID + "/artifacts/" + return strings.HasPrefix(path, artifactPrefix) && strings.HasSuffix(path, "/content") +} - for key, values := range resp.Header { - for _, value := range values { - c.Writer.Header().Add(key, value) +func copyTaskMediaResponseHeaders(destination, source http.Header) { + for _, name := range []string{ + "Content-Type", + "Content-Length", + "Content-Range", + "Accept-Ranges", + "ETag", + "Last-Modified", + "Content-Disposition", + } { + for _, value := range source.Values(name) { + destination.Add(name, value) } } +} + +func setTaskMediaResponseSecurityHeaders(header http.Header) { + header.Set("Cache-Control", "private, no-store") + header.Set("Content-Security-Policy", "sandbox; default-src 'none'") + header.Set("Referrer-Policy", "no-referrer") + header.Set("X-Content-Type-Options", "nosniff") +} - c.Writer.Header().Set("Cache-Control", "public, max-age=86400") - c.Writer.WriteHeader(resp.StatusCode) - if _, err = io.Copy(c.Writer, resp.Body); err != nil { - logger.LogError(c.Request.Context(), fmt.Sprintf("Failed to stream video content: %s", err.Error())) +func writeTaskMediaProxyError(c *gin.Context, err error) { + if c.Writer.Written() { + logger.LogError(c.Request.Context(), err.Error()) + return } + var proxyErr *taskMediaProxyError + if !errors.As(err, &proxyErr) { + proxyErr = &taskMediaProxyError{ + status: http.StatusBadGateway, code: "artifact_upstream_error", + message: "Failed to fetch artifact content", err: err, + } + } + c.Header("Cache-Control", "private, no-store") + writeTaskArtifactError(c, proxyErr.status, proxyErr.code, proxyErr.message) } func writeVideoDataURL(c *gin.Context, dataURL string) error { + if len(dataURL) > taskMediaDataURLMaxEncodedBytes { + return errTaskMediaRequestRejected + } parts := strings.SplitN(dataURL, ",", 2) if len(parts) != 2 { return fmt.Errorf("invalid data url") @@ -199,18 +591,31 @@ func writeVideoDataURL(c *gin.Context, dataURL string) error { if mimeType == "" { mimeType = "video/mp4" } + if len(mimeType) > 255 || !httpguts.ValidHeaderFieldValue(mimeType) { + return fmt.Errorf("invalid data url media type") + } - videoBytes, err := base64.StdEncoding.DecodeString(payload) - if err != nil { - videoBytes, err = base64.RawStdEncoding.DecodeString(payload) - if err != nil { - return err + var encoding *base64.Encoding + var contentLength int64 + for _, candidate := range []*base64.Encoding{base64.StdEncoding, base64.RawStdEncoding} { + decodedLength, err := io.Copy(io.Discard, base64.NewDecoder(candidate, strings.NewReader(payload))) + if err == nil { + encoding = candidate + contentLength = decodedLength + break } } + if encoding == nil { + return fmt.Errorf("invalid base64 data") + } c.Writer.Header().Set("Content-Type", mimeType) - c.Writer.Header().Set("Cache-Control", "public, max-age=86400") + c.Writer.Header().Set("Content-Length", strconv.FormatInt(contentLength, 10)) + setTaskMediaResponseSecurityHeaders(c.Writer.Header()) c.Writer.WriteHeader(http.StatusOK) - _, err = c.Writer.Write(videoBytes) + if c.Request.Method == http.MethodHead { + return nil + } + _, err := io.Copy(c.Writer, base64.NewDecoder(encoding, strings.NewReader(payload))) return err } diff --git a/controller/video_proxy_gemini.go b/controller/video_proxy_gemini.go deleted file mode 100644 index 0c76e33c709a..000000000000 --- a/controller/video_proxy_gemini.go +++ /dev/null @@ -1,294 +0,0 @@ -package controller - -import ( - "fmt" - "io" - "strconv" - "strings" - - "github.com/QuantumNous/new-api/common" - "github.com/QuantumNous/new-api/constant" - "github.com/QuantumNous/new-api/model" - "github.com/QuantumNous/new-api/relay" -) - -func getGeminiVideoURL(channel *model.Channel, task *model.Task, apiKey string) (string, error) { - if channel == nil || task == nil { - return "", fmt.Errorf("invalid channel or task") - } - - if url := extractGeminiVideoURLFromTaskData(task); url != "" { - return ensureAPIKey(url, apiKey), nil - } - - baseURL := constant.ChannelBaseURLs[channel.Type] - if channel.GetBaseURL() != "" { - baseURL = channel.GetBaseURL() - } - - adaptor := relay.GetTaskAdaptor(constant.TaskPlatform(strconv.Itoa(channel.Type))) - if adaptor == nil { - return "", fmt.Errorf("gemini task adaptor not found") - } - - if apiKey == "" { - return "", fmt.Errorf("api key not available for task") - } - - proxy := channel.GetSetting().Proxy - resp, err := adaptor.FetchTask(baseURL, apiKey, map[string]any{ - "task_id": task.GetUpstreamTaskID(), - "action": task.Action, - }, proxy) - if err != nil { - return "", fmt.Errorf("fetch task failed: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return "", fmt.Errorf("read task response failed: %w", err) - } - - taskInfo, parseErr := adaptor.ParseTaskResult(body) - if parseErr == nil && taskInfo != nil && taskInfo.RemoteUrl != "" { - return ensureAPIKey(taskInfo.RemoteUrl, apiKey), nil - } - - if url := extractGeminiVideoURLFromPayload(body); url != "" { - return ensureAPIKey(url, apiKey), nil - } - - if parseErr != nil { - return "", fmt.Errorf("parse task result failed: %w", parseErr) - } - - return "", fmt.Errorf("gemini video url not found") -} - -func extractGeminiVideoURLFromTaskData(task *model.Task) string { - if task == nil || len(task.Data) == 0 { - return "" - } - var payload map[string]any - if err := common.Unmarshal(task.Data, &payload); err != nil { - return "" - } - return extractGeminiVideoURLFromMap(payload) -} - -func extractGeminiVideoURLFromPayload(body []byte) string { - var payload map[string]any - if err := common.Unmarshal(body, &payload); err != nil { - return "" - } - return extractGeminiVideoURLFromMap(payload) -} - -func extractGeminiVideoURLFromMap(payload map[string]any) string { - if payload == nil { - return "" - } - if uri, ok := payload["uri"].(string); ok && uri != "" { - return uri - } - if resp, ok := payload["response"].(map[string]any); ok { - if uri := extractGeminiVideoURLFromResponse(resp); uri != "" { - return uri - } - } - return "" -} - -func extractGeminiVideoURLFromResponse(resp map[string]any) string { - if resp == nil { - return "" - } - if gvr, ok := resp["generateVideoResponse"].(map[string]any); ok { - if uri := extractGeminiVideoURLFromGeneratedSamples(gvr); uri != "" { - return uri - } - } - if videos, ok := resp["videos"].([]any); ok { - for _, video := range videos { - if vm, ok := video.(map[string]any); ok { - if uri, ok := vm["uri"].(string); ok && uri != "" { - return uri - } - } - } - } - if uri, ok := resp["video"].(string); ok && uri != "" { - return uri - } - if uri, ok := resp["uri"].(string); ok && uri != "" { - return uri - } - return "" -} - -func extractGeminiVideoURLFromGeneratedSamples(gvr map[string]any) string { - if gvr == nil { - return "" - } - if samples, ok := gvr["generatedSamples"].([]any); ok { - for _, sample := range samples { - if sm, ok := sample.(map[string]any); ok { - if video, ok := sm["video"].(map[string]any); ok { - if uri, ok := video["uri"].(string); ok && uri != "" { - return uri - } - } - } - } - } - return "" -} - -func getVertexVideoURL(channel *model.Channel, task *model.Task) (string, error) { - if channel == nil || task == nil { - return "", fmt.Errorf("invalid channel or task") - } - if url := strings.TrimSpace(task.GetResultURL()); url != "" && !isTaskProxyContentURL(url, task.TaskID) { - return url, nil - } - if url := extractVertexVideoURLFromTaskData(task); url != "" { - return url, nil - } - - baseURL := constant.ChannelBaseURLs[channel.Type] - if channel.GetBaseURL() != "" { - baseURL = channel.GetBaseURL() - } - - adaptor := relay.GetTaskAdaptor(constant.TaskPlatform(strconv.Itoa(channel.Type))) - if adaptor == nil { - return "", fmt.Errorf("vertex task adaptor not found") - } - - key := getVertexTaskKey(channel, task) - if key == "" { - return "", fmt.Errorf("vertex key not available for task") - } - - resp, err := adaptor.FetchTask(baseURL, key, map[string]any{ - "task_id": task.GetUpstreamTaskID(), - "action": task.Action, - }, channel.GetSetting().Proxy) - if err != nil { - return "", fmt.Errorf("fetch task failed: %w", err) - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - if err != nil { - return "", fmt.Errorf("read task response failed: %w", err) - } - - taskInfo, parseErr := adaptor.ParseTaskResult(body) - if parseErr == nil && taskInfo != nil && strings.TrimSpace(taskInfo.Url) != "" { - return taskInfo.Url, nil - } - if url := extractVertexVideoURLFromPayload(body); url != "" { - return url, nil - } - if parseErr != nil { - return "", fmt.Errorf("parse task result failed: %w", parseErr) - } - return "", fmt.Errorf("vertex video url not found") -} - -func isTaskProxyContentURL(url string, taskID string) bool { - if strings.TrimSpace(url) == "" || strings.TrimSpace(taskID) == "" { - return false - } - return strings.Contains(url, "/v1/videos/"+taskID+"/content") -} - -func getVertexTaskKey(channel *model.Channel, task *model.Task) string { - if task != nil { - if key := strings.TrimSpace(task.PrivateData.Key); key != "" { - return key - } - } - if channel == nil { - return "" - } - keys := channel.GetKeys() - for _, key := range keys { - key = strings.TrimSpace(key) - if key != "" { - return key - } - } - return strings.TrimSpace(channel.Key) -} - -func extractVertexVideoURLFromTaskData(task *model.Task) string { - if task == nil || len(task.Data) == 0 { - return "" - } - return extractVertexVideoURLFromPayload(task.Data) -} - -func extractVertexVideoURLFromPayload(body []byte) string { - var payload map[string]any - if err := common.Unmarshal(body, &payload); err != nil { - return "" - } - resp, ok := payload["response"].(map[string]any) - if !ok || resp == nil { - return "" - } - - if videos, ok := resp["videos"].([]any); ok && len(videos) > 0 { - if video, ok := videos[0].(map[string]any); ok && video != nil { - if b64, _ := video["bytesBase64Encoded"].(string); strings.TrimSpace(b64) != "" { - mime, _ := video["mimeType"].(string) - enc, _ := video["encoding"].(string) - return buildVideoDataURL(mime, enc, b64) - } - } - } - if b64, _ := resp["bytesBase64Encoded"].(string); strings.TrimSpace(b64) != "" { - enc, _ := resp["encoding"].(string) - return buildVideoDataURL("", enc, b64) - } - if video, _ := resp["video"].(string); strings.TrimSpace(video) != "" { - if strings.HasPrefix(video, "data:") || strings.HasPrefix(video, "http://") || strings.HasPrefix(video, "https://") { - return video - } - enc, _ := resp["encoding"].(string) - return buildVideoDataURL("", enc, video) - } - return "" -} - -func buildVideoDataURL(mimeType string, encoding string, base64Data string) string { - mime := strings.TrimSpace(mimeType) - if mime == "" { - enc := strings.TrimSpace(encoding) - if enc == "" { - enc = "mp4" - } - if strings.Contains(enc, "/") { - mime = enc - } else { - mime = "video/" + enc - } - } - return "data:" + mime + ";base64," + base64Data -} - -func ensureAPIKey(uri, key string) string { - if key == "" || uri == "" { - return uri - } - if strings.Contains(uri, "key=") { - return uri - } - if strings.Contains(uri, "?") { - return fmt.Sprintf("%s&key=%s", uri, key) - } - return fmt.Sprintf("%s?key=%s", uri, key) -} diff --git a/controller/wechat.go b/controller/wechat.go index dd185735e2df..f44a2c592751 100644 --- a/controller/wechat.go +++ b/controller/wechat.go @@ -156,21 +156,13 @@ func WeChatBind(c *gin.Context) { }) return } - user := model.User{ - Id: c.GetInt("id"), - } - if user.Id == 0 { + userId := c.GetInt("id") + if userId == 0 { c.JSON(http.StatusUnauthorized, gin.H{"success": false, "message": "未登录"}) return } - err = user.FillUserById() - if err != nil { - common.ApiError(c, err) - return - } - user.WeChatId = wechatId - err = user.Update(false) - if err != nil { + // 只更新绑定列,避免完整用户快照覆盖并发的封禁、降权或分组变更。 + if err := model.UpdateUserBindColumn(userId, "wechat_id", wechatId); err != nil { common.ApiError(c, err) return } diff --git a/docs/authentication.md b/docs/authentication.md index b7b5f870ee00..0ea9d25218ae 100644 --- a/docs/authentication.md +++ b/docs/authentication.md @@ -6,6 +6,7 @@ - Access Token 是有效期 15 分钟的 JWT,只保存在浏览器内存中,通过 `Authorization: Bearer ` 发送。 - Refresh Token 是随机不透明值,有效期最长 30 天。浏览器只通过 `HttpOnly`、`SameSite=Strict` Cookie 持有它;服务端仅保存 HMAC 摘要,并在每次刷新时轮换。 +- `new_api_has_session` 是 Refresh Cookie 的会话提示,值恒为 `1`,`Path=/`、非 `HttpOnly`,与 Refresh Cookie 同时写入、同时清除、同一过期时间。它只声明"曾签发过 Refresh Cookie",不含任何凭据,也不参与任何鉴权判定;伪造它唯一的效果是自费一次注定失败的 refresh。它存在的原因是 Refresh Cookie 被 `HttpOnly` 和 `Path=/api/user/auth` 双重限制,`/` 上的页面无法判断自己是否匿名,否则每次冷启动都要发一次注定 401 的 refresh,而该请求还会占用按 IP 计数的 `CriticalRateLimit` 配额。 - `user_sessions` 是登录会话控制面,记录设备、IP、登录方式、最后活跃时间、到期时间和撤销状态。数据库中的 Session 状态是最终权威;撤销传播速度取决于下文所述的 Redis 拓扑。 - 用户的密码、状态、角色或安全因子发生安全相关变化时,`auth_version` 会递增并使旧登录会话失效。订阅带来的分组升降级只刷新授权缓存,不会退出任何登录设备。 - Redis 缓存保存用户鉴权快照和登录会话快照。版本栅栏和撤销 tombstone 防止旧缓存重新授权;Session 快照使用跟随 `SYNC_FREQUENCY` 的短 TTL,缓存未命中或未启用 Redis 时回退到数据库校验。 @@ -70,6 +71,8 @@ 前端将冷启动状态与登录状态分开管理。网络或服务端临时故障允许后续导航重试 refresh;服务端确认 Refresh Cookie 无效时才进入已完成的匿名状态。内存 SID 与 Cookie SID 不一致时,客户端清除旧内存身份并在不携带旧 SID 的情况下重试一次。 +公开页面的冷启动会先读 `new_api_has_session`:提示不存在且内存中没有任何身份时跳过 refresh,直接按匿名渲染,且**不**把这次跳过记为已完成的匿名判定——跳过只是延后,不是服务端结论。会依据鉴权结果做跳转的位置(受保护路由与登录页)不看提示,内存为空时一律回源。因此提示缺失但 Refresh Cookie 有效的用户(该 Cookie 上线前建立的会话,或只清理了 `/` 站点数据的浏览器)会在公开页显示为匿名,并在进入上述任一位置时自动恢复登录态,不需要重新输入密码。提示因服务端撤销而过期时,那次 refresh 返回 401 并在同一响应里清除提示,浪费的请求只发生一次。 + ## Session 签发限额与保留策略 服务端在所有登录方式的统一 Session 签发出口执行两级账户限制: diff --git a/docs/plugin-api/README.md b/docs/plugin-api/README.md new file mode 100644 index 000000000000..bb188cec040d --- /dev/null +++ b/docs/plugin-api/README.md @@ -0,0 +1,138 @@ +# Task plugin API v1 + +Task plugins are single-file synchronous ECMAScript modules. The plugin contract +is currently unreleased; [`v1.schema.json`](./v1.schema.json) and +[`v1.d.ts`](./v1.d.ts) are the authoritative v1 contract. + +## Contract and lifecycle + +Every plugin exports `meta`, `buildSubmitRequest`, `parseSubmitResponse`, and +`parseTaskResult`. A `per_task` plugin also exports `buildQueryRequest`; a +`batch` plugin exports `buildBatchQueryRequest` and `parseBatchResult`. +`meta.author.name` is required and `meta.author.url`, when present, must be an +absolute HTTP(S) URL. This is self-declared attribution; a future marketplace's +verified publisher identity is a separate host-owned record. +Plugins may declare authenticated vendor-native `meta.routes` and claim +host-owned names through `meta.protocols`. Submit and dynamic routes name a +`native` decoder and presenter; query routes name only a presenter. Protocol +bindings are registered once by the host registry, and protocol decoders receive +the host-parsed `body` union plus the pinned model. Shared protocol hooks are +synchronous transformations; Go owns connections and wire framing. + +The host selects a channel, invokes the request-building hook, validates the +returned URL against the channel host, performs HTTP, and gives the decoded +response to the matching parse hook. It owns persistence, retries, polling, +billing, and settlement. Plugins only transform data and report usage facts. +See [v1.d.ts](./v1.d.ts) for signatures and +[v1.schema.json](./v1.schema.json) for machine-readable shapes. + +Plugins that expose task outputs export `listArtifacts(task)` and +`buildContentRequest(ctx)` together. Artifacts are projected on explicit reads +from persisted `Task.Data`; they are never stored as a second source of truth. +The list contains only stable `key`, `type`, and optional `mimeType` fields. +The content hook receives the selected key, raw decoded task data, the explicit +private upstream task id, the producer plugin version, channel authentication, +and a safe client Range/conditional-header subset. Its URL and headers exist +only for that proxy request. + +When a Responses observation reaches persisted `SUCCESS`, the host also runs +the pinned plugin's `listArtifacts` and injects a read-only +`ctx.artifacts[key] = {key, type, mimeType?, url}` map into `renderEvents` or +`renderFinal`. Each `url` is a long-lived host-signed capability URL, never the +provider URL from `Task.Data`. Nonterminal and failed tasks receive no artifact +map. Capability construction or rendering failure fails only that Responses +observation; it cannot change the task, billing settlement, or refunds. +The absolute URL uses `TaskPublicAddress`, falling back only to +`ServerAddress`; multi-node deployments must share the effective +`CRYPTO_SECRET`. +Dashboard artifact reads return each `content_url` (or the legacy +`legacy_content_url`) directly, without a temporary URL exchange. Capability +generation and verification are stateless and have no expiry; after +verification the host still loads the task, owner, and plugin needed to serve +the artifact. Rotating `CRYPTO_SECRET` invalidates issued URLs. The `access` +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. + +Provider-authenticated content URLs must use the channel base host or a +plugin-declared `meta.allowedHosts` entry. A public dynamic CDN URL may instead +set `credentialless: true`; the host then permits only GET/HEAD with no +plugin-supplied headers or body and applies SSRF checks to the initial URL and +every redirect. + +Registry publication is generation-atomic. A request pins one plugin generation +for its full lifetime, while background polling may use a later active plugin +version. New versions must continue parsing responses for in-flight tasks. +Root administrators can inspect the local node with +`GET /api/plugin/task/runtime/status`. The response includes the node-local generation, +a deterministic revision of the active database overrides, the latest rebuild +outcome, and plugin-level compile or routing errors. Generation numbers are +local to a node; compare database revisions when diagnosing rollout lag between +nodes. If the database snapshot is temporarily unavailable, the endpoint keeps +serving node-local state and the last known revision with `database_error` set. + +For live diagnosis, start the process with `DEBUG=true` and filter logs on +`task_plugin`. Plugin registry, routing, endpoint ownership, channel selection, +submit durability, polling adapters, and protocol observation emit safe +key/value lifecycle events. Request-context events carry the request id; +scheduled, background, and context-less work is labeled `SYSTEM`. Plugin +`console.log` output is also forwarded in DEBUG mode. Hook-time output is +prefixed with plugin key/version; module-initialization output may have an empty +identity during initial upload validation. Do not print credentials, headers, +request bodies, upstream payloads, or private URLs from plugin code; free-form +console output cannot be redacted by the host. + +## Fixtures and dry runs + +A fixture case is `{name?, hook, member?, args, expected?, expectedError?}`. +Keep deterministic cases for every exported hook, its main error branch, batch +behavior, renderers, usage, and content requests. Run a fixture locally with: + +```sh +new-api plugin lint plugin.js +new-api plugin test plugin.js --fixture golden.json +``` + +Root administrators can open the plugin detail Sandbox tab, choose a hook, and +submit an `args` JSON array. `POST /api/plugin/task/:key/dryrun` compiles the +active database source or factory source in a temporary registry and invokes +only that synchronous function. Dry runs never execute a request descriptor and +therefore never contact an upstream service. + +## Upload and release + +Upload from the root-only task plugin page or `POST /api/plugin/task` with +`{"source":"...","remark":"..."}`. The server compiles the module, validates +v1 metadata and required exports, and rejects invalid source before saving it. +Use semantic plugin versions. Reusing a key/version with different source is +rejected; activate or roll back a stored version through the management page. + +For a third-party platform, create a channel of type `Task Plugin`, select the +plugin key, provide an explicit base URL, and configure models. Clients may use +the plugin's declared native routes. The generic management surface remains +`POST /v1/tasks/:pluginKey`, `GET /v1/tasks/:taskId`, and +`GET /v1/tasks/:taskId/artifacts` plus +`GET|HEAD /v1/tasks/:taskId/artifacts/:key/content`. + +## Security boundary + +Plugins have no `fetch`, filesystem, `require`, imports, async functions, or +environment access. The host limits execution time, concurrency, input size, +allowed request hosts, and resolves OAuth credentials outside JavaScript. +Multipart files enter JavaScript only as opaque references. + +This is not a hard memory-isolation boundary. A plugin sees data needed for the +current request and can influence an authenticated upstream request. Uploading a +plugin is an administrator-level trust decision equivalent to configuring a +channel credential. Review source and version diffs before activation. Never run +untrusted plugins merely because they compile. + +Usage hooks may return facts such as seconds, resolution, or upstream units, but +must never calculate prices or attempt quota settlement. The host owns all +pricing and clamps billing conversions. diff --git a/docs/plugin-api/v1.d.ts b/docs/plugin-api/v1.d.ts new file mode 100644 index 000000000000..3dc1919f6438 --- /dev/null +++ b/docs/plugin-api/v1.d.ts @@ -0,0 +1,51 @@ +export type JSONValue = null | boolean | number | string | readonly JSONValue[] | {readonly [key: string]: JSONValue}; +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 = + | Readonly<{kind: "json"; value: JSONValue}> + | Readonly<{kind: "form"; fields: Readonly>}> + | Readonly<{kind: "multipart"; fields: Readonly>; files: readonly FileReference[]}> + | Readonly<{kind: "none"}>; + +export interface NativeDecodeContext {method: string; path: string; params: Readonly>; query: Readonly>; body: DecodedBody} +export interface ProtocolDecodeContext extends NativeDecodeContext {protocol: "openai_responses" | "openai_video"; operation: string; model: string; stream: boolean} +export type SubmitIntent = {kind: "submit"; model: string; action?: string; requestBody?: unknown; originTaskIds?: readonly string[]}; +export type QueryIntent = {kind: "query"; taskIds: readonly string[]}; +export type TaskIntent = SubmitIntent | QueryIntent; +export interface NativeRoute {method: "GET" | "POST" | "PUT" | "PATCH" | "DELETE"; path: string; type: "submit" | "query" | "dynamic"; action?: string; taskIdParam?: string; decode?: string; render: string; models?: readonly string[]} +export type ProtocolName = "openai_responses" | "openai_video"; +export type ResponsesMode = "stream" | "sync" | "background"; +export type ProtocolClaim = + | "openai_video" + | {name: "openai_responses"; supports: readonly ResponsesMode[]; models?: readonly string[]} + | {name: "openai_video"; models?: readonly string[]}; +export type LocalizedText = string | ({ en: string } & Record); +export type UsageFieldSchema = {type: "number"; unit: "second" | "count" | "token" | "credit"; description?: LocalizedText} | {type: "boolean"; description?: LocalizedText} | {enum: readonly string[]; description?: LocalizedText}; +export type UsageExample = {label: string; facts: Readonly>}; +export interface Meta {apiVersion: 1; key: string; name: string; icon?: string; description?: LocalizedText; version: string; author: {name: string; url?: string}; channelTypes?: readonly number[]; models: readonly string[]; fetchMode: "per_task" | "batch"; allowedHosts?: readonly string[]; routes?: readonly NativeRoute[]; protocols?: readonly ProtocolClaim[]; usageSchema?: Readonly>; usageExamples?: readonly UsageExample[]; 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} +export interface DriverContext {requestBody: unknown; requestHeaders: Readonly>; 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>} +export interface RequestDescriptor {url: string; method?: string; headers?: Record; /** 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>; 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} +export declare const meta: Meta; +export declare const native: Record TaskIntent) | ((ctx: NativeDecodeContext, task: TaskView | readonly TaskView[]) => unknown)> & {error?: (ctx: NativeDecodeContext, error: {code: string; message: string; httpStatus: number; retryable: boolean}) => unknown}; +export declare const protocols: { + openai_responses?: {decodeRequest(ctx: ProtocolDecodeContext): SubmitIntent; renderEvents?(ctx: unknown, task: TaskView, previousState: unknown): unknown; renderFinal?(ctx: unknown, task: TaskView): unknown}; + openai_video?: {decodeRequest(ctx: ProtocolDecodeContext): SubmitIntent; render(ctx: unknown, task: TaskView): unknown}; +}; +export declare function buildSubmitRequest(ctx: DriverContext): RequestDescriptor; +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; +export declare function parseTaskResult(ctx: TaskQueryContext, body: unknown, response: HookHTTPResponse): NormalizedTaskResult; +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> | null; +export declare function extractUsageOnSubmit(ctx: DriverContext, taskData: unknown): Readonly> | null; +export declare function extractUsageOnComplete(task: TaskView, result: NormalizedTaskResult, data: unknown): Readonly> | 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>}}): RequestDescriptor; diff --git a/docs/plugin-api/v1.md b/docs/plugin-api/v1.md new file mode 100644 index 000000000000..1c8137811aaa --- /dev/null +++ b/docs/plugin-api/v1.md @@ -0,0 +1,177 @@ +# Task Plugin API v1 + +Task Plugin v1 has two independent entry surfaces. `meta.routes` registers plugin-owned native URLs; `meta.protocols` claims host-owned protocols without registering or copying their URLs. `apiVersion` remains `1`. + +## Manifest + +```js +export const meta = { + apiVersion: 1, + key: "vendor", + name: "Vendor", + version: "1.0.0", + author: {name: "Author"}, + description: {en: "Video generation via the vendor API", zh: "通过厂商接口生成视频"}, + models: ["vendor-model"], + fetchMode: "per_task", + routes: [ + {method: "POST", path: "/vendor/v1/jobs", type: "submit", decode: "createJob", render: "jobCreated"}, + {method: "GET", path: "/vendor/v1/jobs/:task_id", type: "query", render: "jobStatus"}, + ], + protocols: [{name: "openai_responses", supports: ["stream", "sync", "background"]}, "openai_video"], +}; +``` + +`submit` and `dynamic` routes require `decode` and `render`. `query` routes require `render`, prohibit `decode`, and use `taskIdParam` (default `task_id`). Names refer only to callable members of `native`. Route conflicts use method plus normalized path shape and the route index is published atomically with its plugin generation. + +`routes[].models` optionally restricts a `submit` or `dynamic` route to a subset of `meta.models`. The host matches the canonical top-level `model` body field before any JS hook runs; a missing, non-string, or unlisted model is rejected with 400 (plugin routes are exclusive paths, so there is no relay fallback). Declare it only when the model lives at the body top level — vendor formats that nest the model elsewhere should omit `models` and filter inside `decode`. `query` routes have no body and reject the field. + +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. + +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. + +`endpoints`, `routes[].renderer`, global `resolveRequest`, global `renderError`, and global `renderers` are rejected. `parseSubmitResponse` returns only `{taskId, taskData, immediate?, state?}`; `clientResponse` is rejected. + +`icon` is an optional LobeHub icon name string (for example `Sora.Color`). The values `text` and `text: