diff --git a/.github/workflows/pr-agent.yml b/.github/workflows/pr-agent.yml index 63472d98..0fb35e77 100644 --- a/.github/workflows/pr-agent.yml +++ b/.github/workflows/pr-agent.yml @@ -2,8 +2,7 @@ name: PR-Agent on: pull_request_target: - # PR-Agent 通过 base repo 上下文读取 PR diff 并发布 Review,不 checkout 或执行 PR 分支代码。 - # pull_request_target 允许 fork PR 使用仓库 secrets,因此 workflow 只运行固定 digest 的 PR-Agent 容器。 + # Fork 审查需要目标仓库凭据;该 job 仅通过 GitHub API 读取 PR 内容,不 checkout 或执行 PR 分支代码。 types: - opened - reopened @@ -11,17 +10,13 @@ on: - review_requested - synchronize issue_comment: - # 手动命令只在 PR 评论中有意义;编辑后的命令也可重新触发。 types: - created - edited permissions: - # 读取仓库内容和 PR diff。 contents: read - # 更新 PR 描述、发布 PR Review 或修改 PR 相关元数据。 pull-requests: write - # PR 评论在 GitHub API 中属于 issue comments,手动问答需要该权限。 issues: write jobs: @@ -29,21 +24,11 @@ jobs: if: >- github.event.sender.type != 'Bot' && ( - ( - github.event_name == 'pull_request_target' - ) || + github.event_name == 'pull_request_target' || ( github.event_name == 'issue_comment' && github.event.issue.pull_request != null && - contains(fromJSON('["OWNER", "MEMBER", "COLLABORATOR", "CONTRIBUTOR", "FIRST_TIME_CONTRIBUTOR"]'), github.event.comment.author_association) && - ( - github.event.comment.body == '/describe' || - startsWith(github.event.comment.body, '/describe ') || - github.event.comment.body == '/review' || - startsWith(github.event.comment.body, '/review ') || - github.event.comment.body == '/ask' || - startsWith(github.event.comment.body, '/ask ') - ) + contains(fromJSON('["OWNER", "MEMBER", "COLLABORATOR", "CONTRIBUTOR", "FIRST_TIME_CONTRIBUTOR"]'), github.event.comment.author_association) ) ) concurrency: @@ -53,484 +38,15 @@ jobs: timeout-minutes: 20 steps: - - name: Capture PR context - id: pr_context - env: - GH_TOKEN: ${{ github.token }} - REPO: ${{ github.repository }} - PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }} - run: | - set -euo pipefail - pull="$(mktemp)" - gh api "repos/${REPO}/pulls/${PR_NUMBER}" > "${pull}" - python3 - "${pull}" >> "${GITHUB_OUTPUT}" <<'PY' - import json - import re - import sys - from pathlib import Path - - pull = json.loads(Path(sys.argv[1]).read_text(encoding="utf-8")) - title = pull.get("title") or "" - body = pull.get("body") or "" - labels = {item.get("name", "") for item in pull.get("labels") or []} - skip_pr_agent = "true" if "skip pr-agent" in labels or re.search(r"^(?:\[Auto\]|Auto)", title) else "false" - body = re.sub( - r".*?", - " ", - body, - flags=re.DOTALL, - ) - text = f"{title}\n{body}" - cjk_count = len(re.findall(r"[\u4e00-\u9fff]", text)) - latin_words = len(re.findall(r"\b[A-Za-z][A-Za-z]{2,}\b", text)) - if cjk_count >= 4 or latin_words < 8: - response_language = "zh-CN" - summary_heading = "PR-Agent 摘要" - else: - response_language = "en-US" - summary_heading = "PR-Agent Summary" - print(f"head_sha={pull['head']['sha']}") - print(f"changed_files={pull.get('changed_files') or 0}") - print(f"response_language={response_language}") - print(f"summary_heading={summary_heading}") - print(f"skip_pr_agent={skip_pr_agent}") - PY - - - name: Prepare PR-Agent description markers - if: >- - steps.pr_context.outputs.skip_pr_agent != 'true' && - ( - github.event_name == 'pull_request_target' || - github.event.comment.body == '/describe' || - startsWith(github.event.comment.body, '/describe ') - ) - env: - GH_TOKEN: ${{ github.token }} - REPO: ${{ github.repository }} - PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }} - SUMMARY_HEADING: ${{ steps.pr_context.outputs.summary_heading }} - CHANGED_FILES: ${{ steps.pr_context.outputs.changed_files }} - run: | - set -euo pipefail - payload="$(mktemp)" - python3 - "${payload}" <<'PY' - import json - import os - import re - import subprocess - import sys - - pull = json.loads(subprocess.check_output( - ["gh", "api", f"repos/{os.environ['REPO']}/pulls/{os.environ['PR_NUMBER']}"], - text=True, - )) - body = pull.get("body") or "" - start = "" - end = "" - heading = os.environ.get("SUMMARY_HEADING") or "PR-Agent 摘要" - changed_files = int(os.environ.get("CHANGED_FILES") or 0) - block = f"## {heading}\n\n{start}\npr_agent:summary\n{end}" - start_index = body.find(start) - end_index = body.find(end, start_index + len(start)) if start_index >= 0 else -1 - owned_block = re.compile( - r"(?ims)^##\s+(?:PR-Agent\s+摘要|PR-Agent\s+Summary)\s*\n\s*" - r".*?\s*" - ) - if changed_files == 0: - updated = owned_block.sub("", body) - if updated != body: - updated = updated.rstrip() - elif start_index >= 0 and end_index >= 0: - body = re.sub( - r"(?im)^##\s+(PR-Agent\s+摘要|PR-Agent\s+Summary)\s*\n\s*(?=)", - f"## {heading}\n\n", - body, - ) - start_index = body.find(start) - end_index = body.find(end, start_index + len(start)) - updated = body[:start_index] + f"{start}\npr_agent:summary\n{end}" + body[end_index + len(end):] - elif start_index >= 0 or end in body: - updated = body - elif body.strip(): - updated = f"{body.rstrip()}\n\n{block}\n" - else: - updated = f"{block}\n" - if updated != body: - with open(sys.argv[1], "w", encoding="utf-8") as handle: - json.dump({"body": updated}, handle, ensure_ascii=False) - PY - if [ -s "${payload}" ]; then - gh api --method PATCH "repos/${REPO}/pulls/${PR_NUMBER}" --input "${payload}" >/dev/null - fi - - - name: Update PR description - if: >- - steps.pr_context.outputs.skip_pr_agent != 'true' && - steps.pr_context.outputs.changed_files != '0' && - ( - github.event_name == 'pull_request_target' || - github.event.comment.body == '/describe' || - startsWith(github.event.comment.body, '/describe ') - ) - uses: docker://pragent/pr-agent:0.39.0-github_action@sha256:b253845caa8c7ff5ce8be78f32996647982bdd4890826a962b78eff2e385a825 + - name: Run PR Review + uses: docker://ghcr.io/infinitypacer/pr-review-runner:latest env: GITHUB_TOKEN: ${{ github.token }} OPENAI_KEY: ${{ secrets.OPENAI_KEY }} OPENAI.API_BASE: ${{ secrets.OPENAI_API_BASE }} - config.model: 'gpt-5.6-terra' - config.fallback_models: '["gpt-5.5", "gpt-5.4"]' - config.custom_model_max_tokens: '1050000' - config.reasoning_effort: 'medium' - config.ai_timeout: '900' - config.response_language: ${{ steps.pr_context.outputs.response_language }} - config.large_patch_policy: 'clip' - config.ignore_pr_title: '["^\\[Auto\\]", "^Auto"]' - config.ignore_pr_labels: '["skip pr-agent"]' - github_action_config.auto_review: 'false' - github_action_config.auto_describe: 'true' - github_action_config.auto_improve: 'false' - github_action_config.pr_actions: '["opened", "reopened", "ready_for_review", "review_requested"]' - github_action_config.handle_push_trigger: 'true' - github_action_config.push_commands: '["/describe"]' - pr_description.generate_ai_title: 'false' - pr_description.publish_labels: 'false' - pr_description.publish_description_as_comment: 'false' - pr_description.publish_description_as_comment_persistent: 'false' - pr_description.enable_pr_diagram: 'false' - pr_description.enable_pr_type: 'false' - pr_description.enable_help_text: 'false' - pr_description.enable_help_comment: 'false' - pr_description.enable_semantic_files_types: 'false' - pr_description.collapsible_file_list: 'adaptive' - pr_description.add_original_user_description: 'true' - pr_description.use_description_markers: 'true' - pr_description.final_update_message: 'false' - pr_description.extra_instructions: | - Match the configured response language. - Summarize the change goal, key implementation details, compatibility impact, tests, and notable risks. - Use 2-4 bullets for small pull requests and 4-8 bullets for larger changes. - Avoid file lists and local command transcripts. - - - name: Remove unfilled PR-Agent description marker - if: >- - always() && - steps.pr_context.outputs.skip_pr_agent != 'true' && - ( - github.event_name == 'pull_request_target' || - github.event.comment.body == '/describe' || - startsWith(github.event.comment.body, '/describe ') - ) - env: - GH_TOKEN: ${{ github.token }} - REPO: ${{ github.repository }} - PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }} - run: | - set -euo pipefail - payload="$(mktemp)" - python3 - "${payload}" <<'PY' - import json - import os - import re - import subprocess - import sys - - pull = json.loads(subprocess.check_output( - ["gh", "api", f"repos/{os.environ['REPO']}/pulls/{os.environ['PR_NUMBER']}"], - text=True, - )) - body = pull.get("body") or "" - placeholder = "\npr_agent:summary\n" - owned_block = re.compile( - r"(?ims)^##\s+(?:PR-Agent\s+摘要|PR-Agent\s+Summary)\s*\n\s*" - r"\s*pr_agent:summary\s*\s*" - ) - if placeholder in body: - updated = owned_block.sub("", body).rstrip() - if updated != body: - with open(sys.argv[1], "w", encoding="utf-8") as handle: - json.dump({"body": updated}, handle, ensure_ascii=False) - PY - if [ -s "${payload}" ]; then - gh api --method PATCH "repos/${REPO}/pulls/${PR_NUMBER}" --input "${payload}" >/dev/null - fi - - - name: Analyze PR review - id: review_analysis - if: >- - steps.pr_context.outputs.skip_pr_agent != 'true' && - ( - github.event_name == 'pull_request_target' || - github.event.comment.body == '/review' || - startsWith(github.event.comment.body, '/review ') - ) - uses: docker://pragent/pr-agent:0.39.0-github_action@sha256:b253845caa8c7ff5ce8be78f32996647982bdd4890826a962b78eff2e385a825 - env: - GITHUB_TOKEN: ${{ github.token }} - OPENAI_KEY: ${{ secrets.OPENAI_KEY }} - OPENAI.API_BASE: ${{ secrets.OPENAI_API_BASE }} - config.model: ${{ github.event_name == 'issue_comment' && 'gpt-5.6-sol' || 'gpt-5.6-terra' }} - config.fallback_models: '["gpt-5.5", "gpt-5.4"]' - config.custom_model_max_tokens: '1050000' - config.reasoning_effort: 'xhigh' - config.ai_timeout: '900' - config.response_language: ${{ steps.pr_context.outputs.response_language }} - config.large_patch_policy: 'clip' - config.ignore_pr_title: '["^\\[Auto\\]", "^Auto"]' - config.ignore_pr_labels: '["skip pr-agent"]' - config.publish_output: 'false' - github_action_config.auto_review: 'true' - github_action_config.auto_describe: 'false' - github_action_config.auto_improve: 'false' - github_action_config.pr_actions: '["opened", "reopened", "ready_for_review", "review_requested"]' - github_action_config.handle_push_trigger: 'true' - github_action_config.push_commands: '["/review"]' - github_action_config.enable_output: 'true' - pr_reviewer.num_max_findings: '4' - pr_reviewer.require_score_review: 'false' - pr_reviewer.require_tests_review: 'false' - pr_reviewer.require_security_review: 'false' - pr_reviewer.require_estimate_effort_to_review: 'false' - pr_reviewer.require_estimate_contribution_time_cost: 'false' - pr_reviewer.require_can_be_split_review: 'false' - pr_reviewer.require_todo_scan: 'false' - pr_reviewer.require_ticket_analysis_review: 'false' - pr_reviewer.enable_review_labels_effort: 'false' - pr_reviewer.enable_review_labels_security: 'false' - pr_reviewer.extra_instructions: | - Return key_issues_to_review only for concrete behavior defects introduced by this pull request. - Each finding must identify the affected behavior, a reachable trigger, and the existing contract or invariant it violates. - Use issue_content to state the smallest correction boundary, not a code patch. - Do not report style preferences, comments, refactors, architecture alternatives, speculative races, extra hardening, optional tests, or hypothetical concerns. - Return no findings when the evidence is incomplete. - - - name: Publish review comments and summary - if: >- - steps.pr_context.outputs.skip_pr_agent != 'true' && - steps.review_analysis.outcome == 'success' && - ( - github.event_name == 'pull_request_target' || - github.event.comment.body == '/review' || - startsWith(github.event.comment.body, '/review ') - ) - env: - GH_TOKEN: ${{ github.token }} - REPO: ${{ github.repository }} - PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }} - REVIEWED_HEAD_SHA: ${{ steps.pr_context.outputs.head_sha }} - CHANGED_FILES: ${{ steps.pr_context.outputs.changed_files }} - RESPONSE_LANGUAGE: ${{ steps.pr_context.outputs.response_language }} - REVIEW_JSON: ${{ steps.review_analysis.outputs.review }} - run: | - set -euo pipefail - current_head_sha="$(gh api "repos/${REPO}/pulls/${PR_NUMBER}" --jq '.head.sha')" - if [ "${current_head_sha}" != "${REVIEWED_HEAD_SHA}" ]; then - echo "PR head changed during analysis; skip stale review publication." - exit 0 - fi - - review_payload="$(mktemp)" - python3 - "${review_payload}" <<'PY' - import hashlib - import json - import os - import re - import subprocess - import sys - from urllib.parse import quote - - review_raw = os.environ.get("REVIEW_JSON") or "{}" - review = json.loads(review_raw) - if not review_raw.strip() or review == {}: - if int(os.environ.get("CHANGED_FILES") or 0): - raise SystemExit("Review analysis produced no structured output for a non-empty PR.") - review = {} - - repo = os.environ["REPO"] - number = os.environ["PR_NUMBER"] - head_sha = os.environ["REVIEWED_HEAD_SHA"] - language = os.environ.get("RESPONSE_LANGUAGE") or "zh-CN" - - def paged(endpoint): - result = json.loads(subprocess.check_output( - ["gh", "api", "--paginate", "--slurp", endpoint], text=True - )) - if result and all(isinstance(page, list) for page in result): - return [item for page in result for item in page] - return result - - files = paged(f"repos/{repo}/pulls/{number}/files?per_page=100") - comments = paged(f"repos/{repo}/pulls/{number}/comments?per_page=100") - reviews = paged(f"repos/{repo}/pulls/{number}/reviews?per_page=100") - - hunk_pattern = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,\d+)? @@") - changed_lines = {} - for file_data in files: - path = str(file_data.get("filename") or "") - line = None - lines = set() - for patch_line in (file_data.get("patch") or "").splitlines(): - hunk = hunk_pattern.match(patch_line) - if hunk: - line = int(hunk.group(1)) - continue - if line is None or patch_line.startswith("\\"): - continue - if patch_line.startswith("+") and not patch_line.startswith("+++"): - lines.add(line) - line += 1 - elif patch_line.startswith("-") and not patch_line.startswith("---"): - continue - else: - line += 1 - changed_lines[path] = lines - - def fingerprint(path, line): - normalized = "\n".join((path, str(line))) - return hashlib.sha256(normalized.encode("utf-8")).hexdigest()[:16] - - current_fingerprints = set() - current_locations = set() - marker_pattern = re.compile(r"") - for comment in comments: - if comment.get("user", {}).get("login") != "github-actions[bot]": - continue - if comment.get("line") is None: - continue - match = marker_pattern.search(str(comment.get("body") or "")) - if match: - current_fingerprints.add(match.group(1)) - path = str(comment.get("path") or "") - try: - line = int(comment.get("line") or 0) - except (TypeError, ValueError): - line = 0 - if path and line > 0: - current_locations.add((path, line)) - - def code_url(path, line): - return f"https://github.com/{repo}/blob/{head_sha}/{quote(path, safe='/')}#L{line}" - - issues = review.get("key_issues_to_review") or [] - findings = [] - seen = set() - for issue in issues: - if not isinstance(issue, dict): - continue - path = str(issue.get("relevant_file") or "").strip() - header = str(issue.get("issue_header") or "").strip() - content = str(issue.get("issue_content") or "").strip() - try: - line = int(issue.get("start_line") or 0) - except (TypeError, ValueError): - line = 0 - if not path or not header or not content or line < 1: - continue - finding_key = (path, line, header.lower(), " ".join(content.split()).lower()) - if finding_key in seen: - continue - seen.add(finding_key) - findings.append({ - "path": path, - "line": line, - "header": header, - "content": content, - "fingerprint": fingerprint(path, line), - }) - - new_comments = [] - for finding in findings: - if finding["line"] not in changed_lines.get(finding["path"], set()): - continue - if finding["fingerprint"] in current_fingerprints or (finding["path"], finding["line"]) in current_locations: - continue - new_comments.append({ - "path": finding["path"], - "line": finding["line"], - "side": "RIGHT", - "body": "\n".join([ - f"", - f"**{finding['header']}**", - "", - finding["content"], - ]), - }) - - marker = "" - short_sha = head_sha[:7] - commit_url = f"https://github.com/{repo}/commit/{head_sha}" - chinese = language == "zh-CN" - lines = [marker, "## PR-Agent Code Review", ""] - if findings: - for finding in findings: - location = f"{finding['path']}:{finding['line']}" - concise = " ".join(finding["content"].split())[:360] - separator = ":" if chinese else ":" - lines.append(f"- [{location}]({code_url(finding['path'], finding['line'])}){separator} **{finding['header']}** - {concise}") - elif chinese: - lines.append("本次变更无需提出审查意见,暂无其他反馈。") - else: - lines.append("There are no review comments for the current changes. I have no additional feedback to provide.") - lines.extend([ - "", - f"审查提交:[{short_sha}]({commit_url})" if chinese else f"Reviewed commit: [{short_sha}]({commit_url})", - "", - ]) - payload = { - "body": "\n".join(lines), - "commit_id": head_sha, - "event": "COMMENT", - } - if new_comments: - payload["comments"] = new_comments - has_matching_summary = not new_comments and any( - existing.get("user", {}).get("login") == "github-actions[bot]" - and existing.get("commit_id") == head_sha - and str(existing.get("body") or "") == payload["body"] - for existing in reviews - ) - # 同一提交的手工重审仍会完成分析;完全相同的结果不重复发布 Review。 - if not has_matching_summary: - with open(sys.argv[1], "w", encoding="utf-8") as handle: - json.dump(payload, handle, ensure_ascii=False) - PY - - latest_head_sha="$(gh api "repos/${REPO}/pulls/${PR_NUMBER}" --jq '.head.sha')" - if [ "${latest_head_sha}" != "${REVIEWED_HEAD_SHA}" ]; then - echo "PR head changed while rendering review; skip stale review publication." - exit 0 - fi - - if [ -s "${review_payload}" ]; then - gh api --method POST "repos/${REPO}/pulls/${PR_NUMBER}/reviews" --input "${review_payload}" >/dev/null - fi - old_summary_ids="$(gh api --paginate "repos/${REPO}/issues/${PR_NUMBER}/comments?per_page=100" --jq ".[] | select(.user.login == \"github-actions[bot]\" and ((.body | startswith(\"\")) or (.body | startswith(\"\")) or (.body | startswith(\"\")) or (.body | startswith(\"\")))) | .id")" - while IFS= read -r comment_id; do - [ -z "${comment_id}" ] && continue - gh api --method DELETE "repos/${REPO}/issues/comments/${comment_id}" >/dev/null - done <<< "${old_summary_ids}" - - - name: Answer PR question - if: >- - steps.pr_context.outputs.skip_pr_agent != 'true' && - github.event_name == 'issue_comment' && - ( - github.event.comment.body == '/ask' || - startsWith(github.event.comment.body, '/ask ') - ) - uses: docker://pragent/pr-agent:0.39.0-github_action@sha256:b253845caa8c7ff5ce8be78f32996647982bdd4890826a962b78eff2e385a825 - env: - GITHUB_TOKEN: ${{ github.token }} - OPENAI_KEY: ${{ secrets.OPENAI_KEY }} - OPENAI.API_BASE: ${{ secrets.OPENAI_API_BASE }} - config.model: 'gpt-5.6-terra' - config.fallback_models: '["gpt-5.5", "gpt-5.4"]' - config.custom_model_max_tokens: '1050000' - config.reasoning_effort: 'high' - config.ai_timeout: '900' - config.response_language: ${{ steps.pr_context.outputs.response_language }} - config.large_patch_policy: 'clip' - config.ignore_pr_title: '["^\\[Auto\\]", "^Auto"]' - config.ignore_pr_labels: '["skip pr-agent"]' - github_action_config.auto_review: 'false' - github_action_config.auto_describe: 'false' - github_action_config.auto_improve: 'false' + PRR_AUTO_REVIEW_SCOPE: all + PRR_ALLOWED_ASSOCIATIONS: '["OWNER", "MEMBER", "COLLABORATOR", "CONTRIBUTOR", "FIRST_TIME_CONTRIBUTOR"]' + PRR_DISABLED_COMMANDS: '["/improve"]' + PRR_SKIP_LABEL: skip pr-agent + PRR_SKIP_TITLE_PATTERN: '^(?:\[Auto\]|Auto)' + config.response_language: zh-CN diff --git a/app/agent/llm/helper.py b/app/agent/llm/helper.py index 292708de..04accebc 100644 --- a/app/agent/llm/helper.py +++ b/app/agent/llm/helper.py @@ -1058,6 +1058,29 @@ class LLMHelper: http_async_client=_build_httpx_client(llm_proxy, async_client=True), **thinking_kwargs, ) + elif runtime["runtime"] == "bedrock": + from langchain_aws import ChatBedrockConverse + + from app.agent.llm.provider import LLMProviderManager + + aws_region = runtime.get("aws_region") or "us-east-1" + aws_auth = runtime.get("aws_auth") or {} + # Bearer 认证需要跳过 SigV4 签名并注入 Authorization 头,SigV4 认证 + # 直接以 AK/SK 签名;两种方式统一由 provider 管理器构造 boto3 客户端。 + bedrock_client = LLMProviderManager().create_bedrock_client( + "bedrock-runtime", + region=aws_region, + credentials=aws_auth, + base_url=runtime.get("base_url"), + use_proxy=use_proxy, + read_timeout=settings.LLM_TOOL_TIMEOUT, + ) + model = ChatBedrockConverse( + model_id=model_name, + client=bedrock_client, + temperature=temperature_value, + disable_streaming=not streaming, + ) elif runtime["runtime"] in {"anthropic_compatible", "copilot_anthropic"}: from langchain_anthropic import ChatAnthropic @@ -1107,7 +1130,11 @@ class LLMHelper: # 优先使用 provider / models.dev 目录中的上下文上限,减少用户手填成本。 model_profile = getattr(model, "profile", None) if model_profile: - logger.debug(f"使用LLM模型: {model.model},Profile: {model.profile}") + # ChatBedrockConverse 等模型类没有 model 属性,模型名存放在 model_id。 + logged_model_name = getattr(model, "model", None) or getattr( + model, "model_id", model_name + ) + logger.debug(f"使用LLM模型: {logged_model_name},Profile: {model_profile}") else: model_record = runtime.get("model_record") or {} model_metadata = runtime.get("model_metadata") or {} diff --git a/app/agent/llm/provider.py b/app/agent/llm/provider.py index 70bf98b6..d6afcd7f 100644 --- a/app/agent/llm/provider.py +++ b/app/agent/llm/provider.py @@ -7,13 +7,14 @@ import base64 import copy import hashlib import json +import re import secrets import threading import time from dataclasses import dataclass, field from pathlib import Path from typing import Any, Dict, Optional, Tuple -from urllib.parse import urlencode +from urllib.parse import urlencode, urlsplit import aiofiles import httpx @@ -106,6 +107,90 @@ class LLMProviderManager(metaclass=Singleton): _MODELS_DEV_BUNDLED_PATH = Path(__file__).with_name("models.json") _MODELS_DEV_CACHE_TTL = 7 * 24 * 60 * 60 _AUTH_SESSION_DONE_RETENTION = 300 + _BEDROCK_DEFAULT_REGION = "us-east-1" + _BEDROCK_API_KEY_PREFIX = "bedrock-api-key-" + _BEDROCK_GPT_OSS_BASE_REGIONS = ( + "ap-northeast-1", + "ap-south-1", + "ap-southeast-2", + "eu-central-1", + "eu-north-1", + "eu-west-1", + "eu-west-2", + "sa-east-1", + "us-east-1", + "us-east-2", + "us-west-2", + ) + _BEDROCK_GPT_OSS_SAFEGUARD_REGIONS = ( + "ap-northeast-1", + "ap-south-1", + "ap-southeast-2", + "eu-west-1", + "eu-west-2", + "sa-east-1", + "us-east-1", + "us-east-2", + "us-west-2", + ) + _BEDROCK_ON_DEMAND_MODEL_REGIONS = { + "openai.gpt-oss-120b-1:0": _BEDROCK_GPT_OSS_BASE_REGIONS, + "openai.gpt-oss-20b-1:0": _BEDROCK_GPT_OSS_BASE_REGIONS, + "openai.gpt-oss-safeguard-120b": _BEDROCK_GPT_OSS_SAFEGUARD_REGIONS, + "openai.gpt-oss-safeguard-20b": _BEDROCK_GPT_OSS_SAFEGUARD_REGIONS, + "amazon.nova-lite-v1:0": ( + "ap-northeast-1", + "ap-southeast-2", + "eu-west-2", + "us-east-1", + "us-gov-west-1", + ), + "amazon.nova-micro-v1:0": ( + "ap-southeast-2", + "eu-west-2", + "us-east-1", + "us-gov-west-1", + ), + "amazon.nova-pro-v1:0": ( + "ap-southeast-2", + "eu-west-2", + "us-east-1", + "us-gov-west-1", + ), + "anthropic.claude-3-5-haiku-20241022-v1:0": ( + "us-west-2", + ), + "anthropic.claude-3-5-sonnet-20240620-v1:0": ( + "ap-northeast-1", + "ap-northeast-2", + "ap-southeast-1", + "eu-central-1", + "eu-central-2", + "us-east-1", + "us-gov-west-1", + "us-west-2", + ), + "anthropic.claude-3-5-sonnet-20241022-v2:0": ( + "ap-southeast-2", + "us-west-2", + ), + "anthropic.claude-3-7-sonnet-20250219-v1:0": ( + "eu-west-2", + "us-gov-west-1", + ), + "anthropic.claude-3-haiku-20240307-v1:0": ( + "ap-northeast-1", + "ap-northeast-2", + "ap-south-1", + "ap-southeast-2", + "eu-central-1", + "eu-west-1", + "eu-west-3", + "us-east-1", + "us-gov-west-1", + "us-west-2", + ), + } _CHATGPT_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann" _CHATGPT_ISSUER = "https://auth.openai.com" _CHATGPT_CODEX_BASE_URL = "https://chatgpt.com/backend-api/codex" @@ -367,6 +452,50 @@ class LLMProviderManager(metaclass=Singleton): api_key_hint="填写 Anthropic API Key。", description="Anthropic Claude 官方端点。", ), + ProviderSpec( + id="amazon-bedrock", + name="Amazon Bedrock", + runtime="bedrock", + models_dev_provider_id="amazon-bedrock", + default_base_url="https://bedrock-runtime.us-east-1.amazonaws.com", + base_url_presets=( + url_preset( + id="bedrock-us-east-1", + label="美东(弗吉尼亚北部)us-east-1", + value="https://bedrock-runtime.us-east-1.amazonaws.com", + ), + url_preset( + id="bedrock-us-west-2", + label="美西(俄勒冈)us-west-2", + value="https://bedrock-runtime.us-west-2.amazonaws.com", + ), + url_preset( + id="bedrock-eu-central-1", + label="欧洲(法兰克福)eu-central-1", + value="https://bedrock-runtime.eu-central-1.amazonaws.com", + ), + url_preset( + id="bedrock-ap-northeast-1", + label="亚太(东京)ap-northeast-1", + value="https://bedrock-runtime.ap-northeast-1.amazonaws.com", + ), + url_preset( + id="bedrock-ap-southeast-1", + label="亚太(新加坡)ap-southeast-1", + value="https://bedrock-runtime.ap-southeast-1.amazonaws.com", + ), + ), + base_url_editable=True, + api_key_label="Bedrock API Key / AK:SK", + api_key_hint=( + "支持两种认证方式:填写 Amazon Bedrock API Key(bedrock-api-key- 开头," + "Bearer 认证);或填写 Access Key ID:Secret Access Key(可选追加 :Session Token," + "SigV4 认证)。Base URL 决定 AWS Region。" + ), + model_list_strategy="bedrock", + description="Amazon Bedrock 托管模型服务,支持 Bedrock API Key 与 AK/SK 双认证。", + sort_order=35, + ), ProviderSpec( id="deepseek", name="DeepSeek", @@ -1743,6 +1872,112 @@ class LLMProviderManager(metaclass=Singleton): return normalized[:-3] return normalized + @classmethod + def _extract_bedrock_region(cls, base_url: Optional[str]) -> str: + """ + 从 Bedrock 运行时端点 URL 中提取 AWS Region + + 兼容标准端点、FIPS 端点与 PrivateLink(VPCE)端点等主机名形态, + 从中识别 Region 段。 + + :param base_url: 形如 https://bedrock-runtime.us-east-1.amazonaws.com 的端点地址 + :return: 提取到的 Region,无法识别时回退 us-east-1 + """ + hostname = urlsplit((base_url or "").strip().lower()).hostname or "" + match = re.search( + r"(?:^|\.)(?:bedrock(?:-runtime)?(?:-fips)?)" + r"\.([a-z0-9-]+-\d+)(?:\.|$)", + hostname, + ) + if match: + return match.group(1) + return cls._BEDROCK_DEFAULT_REGION + + # Inference Profile 的地理前缀与可用 Region 的对应关系,用于降级目录按 + # 当前 Region 过滤掉不可调用的 Profile 条目。 + _BEDROCK_GEO_PREFIXES: dict[str, tuple[str, ...]] = { + "us": ("us-east-", "us-west-"), + "eu": ("eu-",), + "apac": ("ap-",), + "au": ("ap-southeast-2", "ap-southeast-4"), + "jp": ("ap-northeast-1", "ap-northeast-3"), + "ca": ("ca-",), + } + _BEDROCK_NON_COMMERCIAL_REGION_PREFIXES = ( + "cn-", + "eu-isoe-", + "us-gov-", + "us-iso-", + "us-isob-", + "us-isof-", + ) + + @classmethod + def _bedrock_model_matches_region(cls, model_id: str, region: str) -> bool: + """ + 判断目录中的模型 ID 在指定 Region 是否可调用 + + models.dev 目录同时收录裸模型 ID(直连调用)与带地理前缀的 + Inference Profile ID(us./eu./apac./global. 等)。带前缀的条目只在 + 对应地理分区和 AWS 分区的 Region 可用;global Profile 仅允许商业 + AWS 分区。裸 ID 仅在明确记录的 ON_DEMAND Region 可用,未知条目 + 按不可直连处理。 + + :param model_id: 目录中的模型 ID + :param region: 当前 Base URL 对应的 AWS Region + :return: 该模型在当前 Region 可调用时返回 True + """ + prefix = model_id.split(".", 1)[0] + if prefix == "global": + return not region.startswith(cls._BEDROCK_NON_COMMERCIAL_REGION_PREFIXES) + region_prefixes = cls._BEDROCK_GEO_PREFIXES.get(prefix) + if region_prefixes is not None: + return ( + not region.startswith(cls._BEDROCK_NON_COMMERCIAL_REGION_PREFIXES) + and region.startswith(region_prefixes) + ) + on_demand_regions = cls._BEDROCK_ON_DEMAND_MODEL_REGIONS.get(model_id) + return on_demand_regions is not None and region in on_demand_regions + + @classmethod + def _parse_bedrock_credentials(cls, api_key: Optional[str]) -> dict[str, Any]: + """ + 解析 Bedrock 凭证字符串,识别 Bearer 与 SigV4 两种认证方式 + + - Bedrock API Key(bedrock-api-key- 开头的长期 Key,或控制台生成的短期 + Token)走 Bearer 认证; + - `AccessKeyId:SecretAccessKey` 或 `AccessKeyId:SecretAccessKey:SessionToken` + 走 SigV4 认证,AWS Access Key ID 均以 "AKIA"/"ASIA" 开头。 + + :param api_key: 用户在 API Key 输入框填写的凭证内容 + :return: 含 auth_scheme 及对应凭证字段的字典 + """ + normalized = str(api_key or "").strip() + if not normalized: + raise LLMProviderAuthError( + "Amazon Bedrock 需要填写 Bedrock API Key 或 Access Key ID:Secret Access Key" + ) + + if not normalized.startswith(cls._BEDROCK_API_KEY_PREFIX): + parts = [part.strip() for part in normalized.split(":")] + if len(parts) in {2, 3} and all(parts): + credentials = { + "auth_scheme": "sigv4", + "access_key_id": parts[0], + "secret_access_key": parts[1], + } + if len(parts) == 3: + credentials["session_token"] = parts[2] + return credentials + if ":" in normalized: + raise LLMProviderAuthError( + "Amazon Bedrock AK/SK 凭证格式不正确," + "请按 AccessKeyId:SecretAccessKey 或 " + "AccessKeyId:SecretAccessKey:SessionToken 填写" + ) + + return {"auth_scheme": "bearer", "bearer_token": normalized} + async def _list_models_from_google( self, api_key: str, @@ -1857,6 +2092,235 @@ class LLMProviderManager(metaclass=Singleton): ) return sorted(results, key=lambda item: item["name"].lower()) + def _build_bedrock_boto3_config( + self, + use_proxy: Optional[bool] = None, + ) -> Any: + """ + 构造 Bedrock boto3 客户端配置,统一超时、重试与代理策略 + + :param use_proxy: 是否使用系统代理,None 时读取 LLM_USE_PROXY 配置 + :return: botocore Config 实例 + """ + from botocore.config import Config + + should_use_proxy = settings.LLM_USE_PROXY if use_proxy is None else use_proxy + proxies = None + if should_use_proxy and settings.PROXY_HOST: + proxies = {"http": settings.PROXY_HOST, "https": settings.PROXY_HOST} + return Config( + connect_timeout=10, + read_timeout=60, + retries={"max_attempts": 3, "mode": "standard"}, + proxies=proxies, + ) + + @staticmethod + def _bedrock_endpoint_url( + service_name: str, base_url: Optional[str] + ) -> Optional[str]: + """ + 解析应传给 boto3 客户端的自定义端点 URL + + 标准公有端点交由 boto3 按 Region 自行推导;用户填写 PrivateLink、 + FIPS 等非标准端点时才显式透传,保证所选网络路径实际生效。 + + :param service_name: boto3 服务名(bedrock 或 bedrock-runtime) + :param base_url: 用户配置的 Base URL + :return: 需要显式指定端点时返回 URL,否则返回 None + """ + normalized = (base_url or "").strip().rstrip("/") + if not normalized: + return None + if re.fullmatch( + rf"https://{service_name}\.[a-z0-9-]+\.amazonaws\.com", + normalized, + ): + return None + return normalized + + def create_bedrock_client( + self, + service_name: str, + region: str, + credentials: dict[str, Any], + base_url: Optional[str] = None, + use_proxy: Optional[bool] = None, + read_timeout: Optional[int] = None, + ) -> Any: + """ + 按解析后的凭证创建 Bedrock boto3 客户端,Bearer 方式注入 Authorization 头 + + :param service_name: boto3 服务名(bedrock 或 bedrock-runtime) + :param region: AWS Region + :param credentials: `_parse_bedrock_credentials` 的解析结果 + :param base_url: 用户配置的 Base URL,非标准端点(PrivateLink/FIPS 等)时透传给 boto3 + :param use_proxy: 是否使用系统代理 + :param read_timeout: 读取超时秒数,None 时使用默认值 + :return: boto3 客户端实例 + """ + import boto3 + from botocore import UNSIGNED + + config = self._build_bedrock_boto3_config(use_proxy) + if read_timeout: + config = config.merge(type(config)(read_timeout=read_timeout)) + endpoint_kwargs: dict[str, Any] = {} + endpoint_url = self._bedrock_endpoint_url(service_name, base_url) + if endpoint_url: + endpoint_kwargs["endpoint_url"] = endpoint_url + + if credentials["auth_scheme"] == "sigv4": + return boto3.client( + service_name, + region_name=region, + aws_access_key_id=credentials["access_key_id"], + aws_secret_access_key=credentials["secret_access_key"], + aws_session_token=credentials.get("session_token"), + config=config, + **endpoint_kwargs, + ) + + # Bearer 认证:以 UNSIGNED 跳过 SigV4 签名,再把 API Key 注入 Authorization 头。 + bearer_token = credentials["bearer_token"] + config = config.merge(type(config)(signature_version=UNSIGNED)) + client = boto3.client( + service_name, + region_name=region, + aws_access_key_id="unsigned", + aws_secret_access_key="unsigned", + config=config, + **endpoint_kwargs, + ) + + def _inject_bearer(request: Any, **_kwargs: Any) -> None: + request.headers["Authorization"] = f"Bearer {bearer_token}" + + client.meta.events.register( + f"request-created.{service_name}", + _inject_bearer, + ) + return client + + async def _list_models_from_bedrock_fallback( + self, + region: str, + use_proxy: Optional[bool] = None, + ) -> list[dict[str, Any]]: + """ + 从 models.dev 目录筛选当前 Region 可调用的 Bedrock 模型 + + :param region: 当前 Base URL 对应的 AWS Region + :param use_proxy: 是否使用系统代理 + :return: 过滤后的标准化模型记录列表 + """ + models = await self._list_models_from_models_dev_only( + provider_id="amazon-bedrock", + use_proxy=use_proxy, + ) + return [ + model + for model in models + if self._bedrock_model_matches_region(model["id"], region) + ] + + async def _list_models_from_bedrock( + self, + api_key: str, + base_url: Optional[str], + use_proxy: Optional[bool] = None, + ) -> list[dict[str, Any]]: + """ + 从 Bedrock 控制面拉取模型目录,聚合跨区 Inference Profile 与直连模型 + + Bedrock 多数新模型仅允许通过 Inference Profile(us./eu./apac./global. 前缀) + 调用,因此优先列出 Profile,再补充支持 ON_DEMAND 直连的基础模型。 + + :param api_key: 用户填写的凭证内容(Bedrock API Key 或 AK/SK) + :param base_url: Bedrock 运行时端点,决定 Region + :param use_proxy: 是否使用系统代理 + :return: 标准化后的模型记录列表 + """ + credentials = self._parse_bedrock_credentials(api_key) + region = self._extract_bedrock_region(base_url) + # runtime VPCE 无法安全推导对应的控制面 VPCE;FIPS 端点也不能绕回 + # 公有非 FIPS 控制面,因此直接使用本地目录。 + if self._bedrock_endpoint_url("bedrock-runtime", base_url): + return await self._list_models_from_bedrock_fallback(region, use_proxy) + client = self.create_bedrock_client( + "bedrock", + region=region, + credentials=credentials, + use_proxy=use_proxy, + ) + + def _fetch() -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + profiles: list[dict[str, Any]] = [] + paginator = client.get_paginator("list_inference_profiles") + for page in paginator.paginate(typeEquals="SYSTEM_DEFINED"): + profiles.extend(page.get("inferenceProfileSummaries") or []) + foundation = client.list_foundation_models( + byOutputModality="TEXT", + byInferenceType="ON_DEMAND", + ).get("modelSummaries") or [] + return profiles, foundation + + try: + profile_summaries, foundation_summaries = await asyncio.to_thread(_fetch) + except Exception as err: + # 部分 Bedrock API Key 的授权范围仅覆盖 bedrock-runtime 推理接口, + # 控制面查询被拒时降级到 models.dev 目录,保证仍能选择模型。 + logger.warning( + f"获取 Amazon Bedrock 控制面模型列表失败,降级 models.dev 目录: {err}" + ) + return await self._list_models_from_bedrock_fallback(region, use_proxy) + finally: + await asyncio.to_thread(client.close) + + results: list[dict[str, Any]] = [] + seen_ids: set[str] = set() + + def _append_record(model_id: str, display_name: Optional[str]) -> None: + if not model_id or model_id in seen_ids: + return + seen_ids.add(model_id) + # Inference Profile 带区域前缀,models.dev 目录按基础模型 ID 收录, + # 去掉首个前缀段再查一次元数据。 + metadata = self._cached_models_dev_model("amazon-bedrock", model_id) + if not metadata and "." in model_id: + metadata = self._cached_models_dev_model( + "amazon-bedrock", + model_id.split(".", 1)[1], + ) + results.append( + self._normalize_model_record( + model_id=model_id, + display_name=display_name or (metadata or {}).get("name") or model_id, + metadata=metadata or {}, + source="provider", + ) + ) + + for profile in profile_summaries: + if (profile.get("status") or "ACTIVE") != "ACTIVE": + continue + _append_record( + str(profile.get("inferenceProfileId") or "").strip(), + profile.get("inferenceProfileName"), + ) + # 控制面已按当前 Region 和 ON_DEMAND 筛选,不能复用仅面向 + # models.dev 降级目录的静态白名单,否则 AWS 新增模型会被遗漏。 + for summary in foundation_summaries: + lifecycle = (summary.get("modelLifecycle") or {}).get("status") or "ACTIVE" + if lifecycle != "ACTIVE": + continue + _append_record( + str(summary.get("modelId") or "").strip(), + summary.get("modelName"), + ) + + return sorted(results, key=lambda item: item["name"].lower()) + @staticmethod def _copilot_headers( token: Optional[str] = None, include_auth: bool = True @@ -2064,6 +2528,13 @@ class LLMProviderManager(metaclass=Singleton): use_proxy=use_proxy, ) + if resolved_model_list_strategy == "bedrock": + return await self._list_models_from_bedrock( + api_key=runtime["api_key"], + base_url=runtime.get("base_url"), + use_proxy=use_proxy, + ) + if resolved_model_list_strategy == "anthropic_compatible": return await self._list_models_from_models_dev_only( provider_id=provider_id, @@ -2731,6 +3202,22 @@ class LLMProviderManager(metaclass=Singleton): ) return result + if resolved_runtime == "bedrock": + effective_base_url = normalized_base_url or self._default_base_url_for_provider( + spec + ) + credentials = self._parse_bedrock_credentials(normalized_api_key) + result.update( + { + "api_key": normalized_api_key, + "base_url": effective_base_url, + "aws_region": self._extract_bedrock_region(effective_base_url), + "aws_auth": credentials, + "auth_mode": "api_key", + } + ) + return result + if resolved_runtime == "anthropic_compatible": effective_base_url = normalized_base_url or self._default_base_url_for_provider( spec diff --git a/app/api/endpoints/system.py b/app/api/endpoints/system.py index 0fd70b3c..98b44244 100644 --- a/app/api/endpoints/system.py +++ b/app/api/endpoints/system.py @@ -1339,12 +1339,7 @@ def restart_system(_: User = Depends(get_current_active_superuser)): """ if not SystemHelper.can_restart(): return schemas.Response(success=False, message="当前运行环境不支持重启操作!") - # 标识停止事件 - global_vars.stop_system() - # 执行重启 ret, msg = SystemHelper.restart() - if not ret: - global_vars.resume_system() return schemas.Response(success=ret, message=msg) @@ -1362,11 +1357,7 @@ def upgrade_system( if not SystemHelper.can_restart(): return schemas.Response(success=False, message="当前运行环境不支持升级操作!") - # 标识停止事件 - global_vars.stop_system() ret, msg = SystemHelper.upgrade(mode=mode or "release") - if not ret: - global_vars.resume_system() return schemas.Response(success=ret, message=msg) diff --git a/app/chain/mediaserver.py b/app/chain/mediaserver.py index cb0d5f82..8f3dcf7d 100644 --- a/app/chain/mediaserver.py +++ b/app/chain/mediaserver.py @@ -1,6 +1,6 @@ import threading from datetime import datetime -from typing import Callable, List, Union, Optional, Generator, Any +from typing import Callable, Dict, List, Union, Optional, Generator, Any from app.chain import ChainBase from app.core.config import global_vars @@ -210,6 +210,24 @@ class MediaServerChain(ChainBase): """ return self.run_module("mediaserver_play_url", server=server, item_id=item_id) + def get_season_episode_ids(self, server: str, item_id: Union[str, int], + season: int) -> Dict[int, str]: + """ + 获取指定季的集号到媒体服务器条目 ID 映射 + + :param server: 媒体服务器名称 + :param item_id: 剧集在媒体服务器中的条目 ID + :param season: 季号 + :return: 集号到条目 ID 的映射,无数据时返回空字典 + """ + result = self.run_module( + "mediaserver_season_episode_ids", + server=server, + item_id=item_id, + season=season, + ) + return result or {} + def get_image_cookies( self, server: Optional[str], image_url: str ) -> Optional[str | dict]: diff --git a/app/chain/subscribe.py b/app/chain/subscribe.py index dc60f313..7b4b7d3c 100644 --- a/app/chain/subscribe.py +++ b/app/chain/subscribe.py @@ -11,6 +11,7 @@ from app import schemas from app.chain import ChainBase from app.chain.download import DownloadChain from app.chain.media import MediaChain +from app.chain.mediaserver import MediaServerChain from app.chain.search import SearchChain from app.chain.tmdb import TmdbChain from app.chain.torrents import TorrentsChain @@ -34,6 +35,7 @@ from app.helper.interaction import ( supports_markdown, update_or_post_message, ) +from app.helper.mediaserver import MediaServerHelper from app.helper.server import MoviePilotServerHelper from app.helper.torrent import TorrentHelper from app.log import logger @@ -3642,6 +3644,84 @@ class SubscribeChain(ChainBase): else: episodes[0].library.append(file_info) + # 合并所有媒体服务器已存在条目(逐台查询,不只取第一个命中) + mediaserver_chain = MediaServerChain() + server_names = list(MediaServerHelper().get_services().keys()) + + def _has_server_entry(library_list: List[schemas.SubscribeLibraryFileInfo], + server_name: Optional[str], + server_type: Optional[str]) -> bool: + for info in library_list or []: + if info.server and server_name and info.server == server_name: + return True + if info.server_type and server_type and info.server_type == server_type \ + and info.server == server_name \ + and (not info.file_path or str(info.file_path).startswith(("http://", "https://"))): + return True + return False + + for server_name in server_names: + exists_media = self.media_exists(mediainfo=mediainfo, server=server_name) + # 仅合并真实媒体服务器结果,跳过本地 FileManager 兜底(已由 media_files 覆盖) + if not exists_media or not (exists_media.server or exists_media.server_type): + continue + + resolved_server = exists_media.server or server_name + server_storage = exists_media.server_type or resolved_server + server_itemid = str(exists_media.itemid) if exists_media.itemid is not None else None + series_detail_url = None + if resolved_server and exists_media.itemid is not None: + series_detail_url = mediaserver_chain.get_play_url( + server=resolved_server, + item_id=exists_media.itemid, + ) + + if subscribe.type == MediaType.TV.value: + season_number = subscribe.season if subscribe.season is not None else 1 + exist_episodes = (exists_media.seasons or {}).get(season_number) or [] + episode_item_ids: Dict[int, str] = {} + if resolved_server and exists_media.itemid is not None: + episode_item_ids = mediaserver_chain.get_season_episode_ids( + server=resolved_server, + item_id=exists_media.itemid, + season=season_number, + ) + for episode_number in exist_episodes: + episode_info = episodes.get(episode_number) + if not episode_info: + continue + if _has_server_entry(episode_info.library, resolved_server, exists_media.server_type): + continue + episode_itemid = episode_item_ids.get(episode_number) or server_itemid + detail_url = series_detail_url + if resolved_server and episode_item_ids.get(episode_number): + detail_url = mediaserver_chain.get_play_url( + server=resolved_server, + item_id=episode_itemid, + ) or series_detail_url + episode_info.library.append( + schemas.SubscribeLibraryFileInfo( + storage=server_storage, + file_path=detail_url, + server=resolved_server, + server_type=exists_media.server_type, + itemid=str(episode_itemid) if episode_itemid is not None else None, + ) + ) + else: + episode_info = episodes.get(0) + if episode_info and not _has_server_entry( + episode_info.library, resolved_server, exists_media.server_type): + episode_info.library.append( + schemas.SubscribeLibraryFileInfo( + storage=server_storage, + file_path=series_detail_url, + server=resolved_server, + server_type=exists_media.server_type, + itemid=server_itemid, + ) + ) + # 更新订阅信息 subscribe_info.subscribe = Subscribe(**subscribe.to_dict()) subscribe_info.episodes = episodes diff --git a/app/chain/system.py b/app/chain/system.py index c3bc5cf6..59abaab9 100644 --- a/app/chain/system.py +++ b/app/chain/system.py @@ -38,8 +38,6 @@ class SystemChain(ChainBase): """ 重启系统 """ - from app.core.config import global_vars - if channel and userid: self.post_message(Notification( channel=channel, @@ -54,8 +52,6 @@ class SystemChain(ChainBase): }, self._restart_file) # 主动备份一次插件 self.backup_plugins() - # 设置停止标志,通知所有模块准备停止 - global_vars.stop_system() # 重启 SystemHelper.restart() diff --git a/app/core/config.py b/app/core/config.py index 3d7a5f6d..06f52260 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -1211,12 +1211,6 @@ class GlobalVar(object): """ self.STOP_EVENT.set() - def resume_system(self): - """ - 恢复系统运行标记。 - """ - self.STOP_EVENT.clear() - @property def is_system_stopped(self): """ diff --git a/app/main.py b/app/main.py index 9cc53f29..400cc75e 100644 --- a/app/main.py +++ b/app/main.py @@ -30,16 +30,31 @@ elif SystemUtils.is_frozen(): sys.stderr = open(os.devnull, 'w') from app.factory import app -from app.core.config import settings +from app.core.config import global_vars, settings from app.db.init import init_db, update_db # 设置进程名 setproctitle.setproctitle(settings.PROJECT_NAME) + +class MoviePilotServer(uvicorn.Server): + """在 Uvicorn 开始优雅退出前发布应用协作停止标志""" + + def handle_exit(self, sig, frame) -> None: + global_vars.stop_system() + super().handle_exit(sig, frame) + + # uvicorn服务 -Server = uvicorn.Server(Config(app, host=settings.HOST, port=settings.PORT, - reload=settings.DEV, workers=multiprocessing.cpu_count() * 2 + 1, - timeout_graceful_shutdown=60)) +Server = MoviePilotServer(Config(app, host=settings.HOST, port=settings.PORT, + reload=settings.DEV, workers=multiprocessing.cpu_count() * 2 + 1, + timeout_graceful_shutdown=60)) + + +def request_shutdown() -> None: + """发布协作停止标志并请求 Uvicorn 退出""" + global_vars.stop_system() + Server.should_exit = True def start_tray(): @@ -64,8 +79,8 @@ def start_tray(): """ 退出程序 """ + request_shutdown() TrayIcon.stop() - Server.should_exit = True import pystray @@ -93,10 +108,11 @@ def signal_handler(signum, frame): 信号处理函数,用于优雅停止服务 """ print(f"收到信号 {signum},开始优雅停止服务...") - Server.should_exit = True + request_shutdown() -if __name__ == '__main__': +def run_application() -> None: + """初始化进程并启动 API 服务""" # 注册信号处理器 signal.signal(signal.SIGTERM, signal_handler) signal.signal(signal.SIGINT, signal_handler) @@ -109,3 +125,7 @@ if __name__ == '__main__': update_db() # 启动API服务 Server.run() + + +if __name__ == '__main__': + run_application() diff --git a/app/modules/emby/__init__.py b/app/modules/emby/__init__.py index 665a4d6e..1aa36104 100644 --- a/app/modules/emby/__init__.py +++ b/app/modules/emby/__init__.py @@ -1,4 +1,4 @@ -from typing import Any, Generator, List, Optional, Tuple, Union +from typing import Any, Dict, Generator, List, Optional, Tuple, Union from app import schemas from app.core.context import MediaInfo @@ -300,6 +300,21 @@ class EmbyModule(_ModuleBase, _MediaServerBase[Emby]): return None return server_obj.get_play_url(item_id) + def mediaserver_season_episode_ids(self, server: str, item_id: Union[str, int], + season: int) -> Optional[Dict[int, str]]: + """ + 获取指定季的集号到条目 ID 映射 + + :param server: Emby 媒体服务器名称 + :param item_id: 剧集在 Emby 中的条目 ID + :param season: 季号 + :return: 集号到条目 ID 的映射,服务器不可用或无数据时返回 None + """ + server_obj: Emby = self.get_instance(server) + if not server_obj: + return None + return server_obj.get_season_episode_ids(str(item_id), season) + def mediaserver_latest(self, server: Optional[str] = None, count: Optional[int] = 20, username: Optional[str] = None) -> List[schemas.MediaServerPlayItem]: """ diff --git a/app/modules/emby/emby.py b/app/modules/emby/emby.py index 68bba8b7..5b25f4a7 100644 --- a/app/modules/emby/emby.py +++ b/app/modules/emby/emby.py @@ -474,6 +474,37 @@ class Emby: return None, None return None, {} + def get_season_episode_ids(self, item_id: str, season: int) -> Dict[int, str]: + """ + 获取指定季的集号到媒体服务器条目 ID 映射 + :param item_id: 剧集在 Emby 中的 ID + :param season: 季号 + :return: {集号: episode_item_id} + """ + if not item_id or not self._host or not self._apikey: + return {} + try: + url = f"{self._host}emby/Shows/{item_id}/Episodes" + params = { + "Season": season, + "IsMissing": "false", + "api_key": self._apikey + } + res_json = RequestUtils().get_res(url, params) + if not res_json: + return {} + episode_ids: Dict[int, str] = {} + for res_item in res_json.json().get("Items") or []: + episode_index = res_item.get("IndexNumber") + episode_id = res_item.get("Id") + if episode_index is None or not episode_id: + continue + episode_ids[int(episode_index)] = str(episode_id) + return episode_ids + except Exception as e: + logger.error(f"获取 Emby 季集条目 ID 出错:{str(e)}") + return {} + def get_remote_image_by_id(self, item_id: str, image_type: str) -> Optional[str]: """ 根据ItemId从Emby查询TMDB的图片地址 diff --git a/app/modules/filemanager/__init__.py b/app/modules/filemanager/__init__.py index b6e7b76c..078bc7a3 100644 --- a/app/modules/filemanager/__init__.py +++ b/app/modules/filemanager/__init__.py @@ -606,8 +606,12 @@ class FileManagerModule(_ModuleBase): """ 判断媒体文件是否存在于文件系统(网盘或本地文件),只支持标准媒体库结构 :param mediainfo: 识别的媒体信息 + :param server: 指定媒体服务器名称时跳过本地文件系统检查 :return: 如不存在返回None,存在时返回信息,包括每季已存在所有集{type: movie/tv, seasons: {season: [episodes]}} """ + if kwargs.get("server"): + return None + if not settings.LOCAL_EXISTS_SEARCH: return None diff --git a/app/modules/jellyfin/__init__.py b/app/modules/jellyfin/__init__.py index 5b10f384..b0da28b3 100644 --- a/app/modules/jellyfin/__init__.py +++ b/app/modules/jellyfin/__init__.py @@ -1,4 +1,4 @@ -from typing import Any, Generator, List, Optional, Tuple, Union +from typing import Any, Dict, Generator, List, Optional, Tuple, Union from app import schemas from app.core.context import MediaInfo @@ -299,6 +299,21 @@ class JellyfinModule(_ModuleBase, _MediaServerBase[Jellyfin]): return None return server_obj.get_play_url(item_id) + def mediaserver_season_episode_ids(self, server: str, item_id: Union[str, int], + season: int) -> Optional[Dict[int, str]]: + """ + 获取指定季的集号到条目 ID 映射 + + :param server: Jellyfin 媒体服务器名称 + :param item_id: 剧集在 Jellyfin 中的条目 ID + :param season: 季号 + :return: 集号到条目 ID 的映射,服务器不可用或无数据时返回 None + """ + server_obj: Jellyfin = self.get_instance(server) + if not server_obj: + return None + return server_obj.get_season_episode_ids(str(item_id), season) + def mediaserver_latest(self, server: Optional[str] = None, count: Optional[int] = 20, username: Optional[str] = None) -> List[schemas.MediaServerPlayItem]: """ diff --git a/app/modules/jellyfin/jellyfin.py b/app/modules/jellyfin/jellyfin.py index 3a6f324f..005e0244 100644 --- a/app/modules/jellyfin/jellyfin.py +++ b/app/modules/jellyfin/jellyfin.py @@ -523,6 +523,38 @@ class Jellyfin: return None, None return None, {} + def get_season_episode_ids(self, item_id: str, season: int) -> Dict[int, str]: + """ + 获取指定季的集号到媒体服务器条目 ID 映射 + :param item_id: 剧集在 Jellyfin 中的 ID + :param season: 季号 + :return: {集号: episode_item_id} + """ + if not item_id or not self._host or not self._apikey or not self.user: + return {} + try: + url = f"{self._host}Shows/{item_id}/Episodes" + params = { + "season": season, + "userId": self.user, + "isMissing": "false", + "api_key": self._apikey + } + res_json = RequestUtils().get_res(url, params) + if not res_json: + return {} + episode_ids: Dict[int, str] = {} + for res_item in res_json.json().get("Items") or []: + episode_index = res_item.get("IndexNumber") + episode_id = res_item.get("Id") + if episode_index is None or not episode_id: + continue + episode_ids[int(episode_index)] = str(episode_id) + return episode_ids + except Exception as e: + logger.error(f"获取 Jellyfin 季集条目 ID 出错:{str(e)}") + return {} + def get_remote_image_by_id(self, item_id: str, image_type: str) -> Optional[str]: """ 根据ItemId从Jellyfin查询TMDB图片地址 diff --git a/app/modules/plex/__init__.py b/app/modules/plex/__init__.py index f386658f..46cf2ebe 100644 --- a/app/modules/plex/__init__.py +++ b/app/modules/plex/__init__.py @@ -1,4 +1,4 @@ -from typing import Optional, Tuple, Union, Any, List, Generator +from typing import Optional, Tuple, Union, Any, List, Generator, Dict from app import schemas from app.core.context import MediaInfo @@ -349,3 +349,18 @@ class PlexModule(_ModuleBase, _MediaServerBase[Plex]): if not server_obj: return None return server_obj.get_play_url(item_id) + + def mediaserver_season_episode_ids(self, server: str, item_id: Union[str, int], + season: int) -> Optional[Dict[int, str]]: + """ + 获取指定季的集号到条目 ID 映射 + + :param server: Plex 媒体服务器名称 + :param item_id: 剧集在 Plex 中的条目 ID / key + :param season: 季号 + :return: 集号到条目 ID 的映射,服务器不可用或无数据时返回 None + """ + server_obj: Plex = self.get_instance(server) + if not server_obj: + return None + return server_obj.get_season_episode_ids(str(item_id), season) diff --git a/app/modules/plex/plex.py b/app/modules/plex/plex.py index c051ee95..f0a38cc5 100644 --- a/app/modules/plex/plex.py +++ b/app/modules/plex/plex.py @@ -293,6 +293,31 @@ class Plex: season_episodes[episode.seasonNumber].append(episode.index) return videos.key, season_episodes + def get_season_episode_ids(self, item_id: str, season: int) -> Dict[int, str]: + """ + 获取指定季的集号到媒体服务器条目 ID 映射 + :param item_id: 剧集在 Plex 中的 ID / key + :param season: 季号 + :return: {集号: episode_item_key} + """ + if not self._plex or not item_id: + return {} + try: + videos = self.__fetch_item(item_id) + if not videos: + return {} + episode_ids: Dict[int, str] = {} + for episode in videos.episodes(): + if episode.seasonNumber != int(season): + continue + if episode.index is None or not episode.key: + continue + episode_ids[int(episode.index)] = str(episode.key) + return episode_ids + except Exception as e: + logger.error(f"获取 Plex 季集条目 ID 出错:{str(e)}") + return {} + def __search_show(self, title: Optional[str] = None, original_title: Optional[str] = None, diff --git a/app/modules/thetvdb/__init__.py b/app/modules/thetvdb/__init__.py index de4e2711..d7d9baa5 100644 --- a/app/modules/thetvdb/__init__.py +++ b/app/modules/thetvdb/__init__.py @@ -114,7 +114,6 @@ class TheTvDbModule(_ModuleBase): return 4 def stop(self): - logger.info("TheTvDbModule 停止。正在清除 TVDB 会话。") with self.__auth_lock: self.tvdb = None diff --git a/app/schemas/subscribe.py b/app/schemas/subscribe.py index a16338ed..68dcc354 100644 --- a/app/schemas/subscribe.py +++ b/app/schemas/subscribe.py @@ -217,6 +217,12 @@ class SubscribeLibraryFileInfo(BaseModel): storage: Optional[str] = "local" # 文件路径 file_path: Optional[str] = None + # 媒体服务器名称 + server: Optional[str] = None + # 媒体服务器类型:emby、jellyfin、plex 等 + server_type: Optional[str] = None + # 媒体服务器条目 ID + itemid: Optional[str] = None class SubscribeEpisodeInfo(BaseModel): diff --git a/app/startup/lifecycle.py b/app/startup/lifecycle.py index a40d1525..3d114271 100644 --- a/app/startup/lifecycle.py +++ b/app/startup/lifecycle.py @@ -1,5 +1,7 @@ import asyncio +import inspect from contextlib import asynccontextmanager +from typing import Callable from fastapi import FastAPI @@ -20,7 +22,7 @@ from app.chain.system import SystemChain from app.core.config import global_vars, settings from app.helper.server import MoviePilotServerHelper from app.helper.system import SystemHelper -from app.log import LoggerManager +from app.log import logger, LoggerManager from app.startup.command_initializer import init_command, stop_command, restart_command from app.startup.modules_initializer import init_modules, stop_modules from app.startup.monitor_initializer import stop_monitor, init_monitor @@ -56,6 +58,16 @@ async def init_extra(): await MoviePilotServerHelper.async_report_usage() +async def run_shutdown_step(name: str, callback: Callable[[], object]) -> None: + """隔离单个关闭阶段的异常,确保后续资源仍有机会释放""" + try: + result = callback() + if inspect.isawaitable(result): + await result + except Exception as err: + logger.error(f"关闭{name}失败:{err}") + + @asynccontextmanager async def lifespan(app: FastAPI): """ @@ -90,6 +102,7 @@ async def lifespan(app: FastAPI): yield finally: print("Shutting down...") + global_vars.stop_system() # 取消同步插件任务 try: sync_plugins_task.cancel() @@ -100,22 +113,19 @@ async def lifespan(app: FastAPI): print(str(e)) try: if not settings.MOVIEPILOT_SAFE_MODE: - # 备份插件 - SystemChain().backup_plugins() - # 停止工作流 - stop_workflow() - # 停止命令 - stop_command() - # 停止监控器 - stop_monitor() - # 停止定时器 - stop_scheduler() - # 停止插件 - stop_plugins() - # 停止模块 - await stop_modules() - # 关闭共享的异步 HTTP 连接池,释放底层连接资源 - await aclose_shared_async_transports() + await run_shutdown_step( + "插件备份", lambda: SystemChain().backup_plugins() + ) + await run_shutdown_step("工作流", stop_workflow) + await run_shutdown_step("命令服务", stop_command) + await run_shutdown_step("监控器", stop_monitor) + await run_shutdown_step("定时器", stop_scheduler) + await run_shutdown_step("插件", stop_plugins) + await run_shutdown_step("模块服务", stop_modules) + await run_shutdown_step( + "共享异步 HTTP 连接池", + aclose_shared_async_transports, + ) finally: # 日志最后关闭,确保其他组件的收尾信息已写入文件 LoggerManager.shutdown() diff --git a/app/startup/modules_initializer.py b/app/startup/modules_initializer.py index 840cfd3d..e8a66f54 100644 --- a/app/startup/modules_initializer.py +++ b/app/startup/modules_initializer.py @@ -1,4 +1,6 @@ +import inspect import sys +from typing import Callable from app.helper.redis import RedisHelper, AsyncRedisHelper @@ -129,29 +131,27 @@ async def stop_modules(): """ 服务关闭 """ - # 停止AI智能体 - await stop_agent() - # 停止模块 - ModuleManager().stop() - # 停止事件消费 - EventManager().stop() - # 停止虚拟显示 - DisplayHelper().stop() - # 停止 DoH 服务 - DohHelper().shutdown() - # 停止线程池 - ThreadHelper().shutdown() - # 停止消息服务 - stop_message() - # 关闭Redis缓存连接 - RedisHelper().close() - await AsyncRedisHelper().close() - # 停止数据库连接 - await close_database() - # 停止前端服务 - stop_frontend() - # 清理临时文件 - clear_temp() + async def run_step(name: str, callback: Callable[[], object]) -> None: + """单个模块资源关闭失败时继续执行后续阶段""" + try: + result = callback() + if inspect.isawaitable(result): + await result + except Exception as err: + logger.error(f"关闭{name}失败:{err}") + + await run_step("AI智能体", stop_agent) + await run_step("模块", lambda: ModuleManager().stop()) + await run_step("事件消费", lambda: EventManager().stop()) + await run_step("虚拟显示", lambda: DisplayHelper().stop()) + await run_step("DoH服务", lambda: DohHelper().shutdown()) + await run_step("线程池", lambda: ThreadHelper().shutdown()) + await run_step("消息服务", stop_message) + await run_step("Redis缓存连接", lambda: RedisHelper().close()) + await run_step("异步Redis缓存连接", lambda: AsyncRedisHelper().close()) + await run_step("数据库连接", close_database) + await run_step("前端服务", stop_frontend) + await run_step("临时文件", clear_temp) def init_modules(): diff --git a/app/utils/http.py b/app/utils/http.py index dbf8ddf2..12662d5e 100644 --- a/app/utils/http.py +++ b/app/utils/http.py @@ -76,6 +76,14 @@ _REQUESTS_RETRY_IDEMPOTENT_METHODS = ("GET", "HEAD", "OPTIONS") _pending_eviction_tasks: set[asyncio.Task] = set() +def _discard_pending_eviction_task(task: asyncio.Task) -> None: + """从跨线程共享集合移除已完成的 transport 关闭任务""" + with _shared_async_transports_lock: + _pending_eviction_tasks.discard(task) + if not task.cancelled() and (error := task.exception()): + logger.debug(f"LRU 淘汰共享 transport 时关闭失败: {error!r}") + + def _get_shared_async_transport( proxy: Optional[str], verify: Union[bool, str], @@ -140,8 +148,9 @@ def _get_shared_async_transport( try: task = loop.create_task(evicted_transport.aclose()) # 强引用避免 task 仅被 loop 弱持有而触发 "Task was destroyed but pending" - _pending_eviction_tasks.add(task) - task.add_done_callback(_pending_eviction_tasks.discard) + with _shared_async_transports_lock: + _pending_eviction_tasks.add(task) + task.add_done_callback(_discard_pending_eviction_task) except Exception as e: # pragma: no cover - 防御性 logger.debug(f"LRU 淘汰共享 transport 时调度关闭失败: {e!r}") @@ -160,17 +169,27 @@ async def aclose_shared_async_transports() -> None: # 弹出而非 get+clear,避免外层 dict 残留空 OrderedDict 占位 with _shared_async_transports_lock: per_loop = _shared_async_transports.pop(loop, None) - if not per_loop: + pending_evictions = [ + task + for task in _pending_eviction_tasks + if task.get_loop() is loop + ] + transports = list(per_loop.values()) if per_loop else [] + if per_loop: + per_loop.clear() + if not transports and not pending_evictions: return - transports = list(per_loop.values()) - per_loop.clear() # 并行关闭:每个 transport 的 TLS close_notify 各占一个 RTT, # 顺序等待会线性放大 shutdown 耗时;return_exceptions 让单点失败 # 不影响其他 transport 的释放 results = await asyncio.gather( - *(t.aclose() for t in transports), return_exceptions=True + *pending_evictions, + *(t.aclose() for t in transports), + return_exceptions=True, ) - for result in results: + with _shared_async_transports_lock: + _pending_eviction_tasks.difference_update(pending_evictions) + for result in results[len(pending_evictions):]: if isinstance(result, BaseException): logger.debug(f"关闭共享 AsyncHTTPTransport 失败: {result!r}") diff --git a/docs/pr-agent.md b/docs/pr-agent.md index a13def77..be64ae36 100644 --- a/docs/pr-agent.md +++ b/docs/pr-agent.md @@ -1,6 +1,6 @@ # PR-Agent 使用说明 -本仓库通过 GitHub Actions 运行 PR-Agent,帮助贡献者维护 PR 摘要、获取代码审查结果和提出 PR 相关问题。 +本仓库通过 PR Review Runner 运行 PR-Agent,帮助贡献者维护 PR 摘要、获取代码审查结果和提出 PR 相关问题。 ## 自动执行 @@ -25,26 +25,24 @@ PR 带有 `skip pr-agent` 标签,或标题以 `[Auto]`、`Auto` 开头时, /ask 这次改动有没有遗漏权限校验? ``` -- `/describe`:更新 PR Body 内按语言显示的 `PR-Agent 摘要` 或 `PR-Agent Summary`,并保留贡献者原有的 PR 描述。 +- `/describe`:更新 PR Body 内的 `PR-Agent 摘要`,并保留贡献者原有的 PR 描述。 - `/review`:发起一次代码审查。 - `/ask ...`:就当前 PR 提问,回复会发布在普通 PR 评论中。 +本仓库禁用 `/improve` 及其等价别名。其他命令是否可用由 runner 所包含的 PR-Agent 能力决定。 + 手工命令仅允许以下 GitHub 身份关联的用户使用:`OWNER`、`MEMBER`、`COLLABORATOR`、`CONTRIBUTOR`、`FIRST_TIME_CONTRIBUTOR`。 新建的合法命令评论会触发执行;编辑后仍为合法命令的评论也会触发。编辑普通讨论评论不会调用模型。 ## 审查结果 -`/describe` 的结果位于 PR Body 中按语言显示的 `PR-Agent 摘要` 或 `PR-Agent Summary` 区域,用于概览本次变更。 +本仓库固定使用中文生成 PR-Agent 内容。`/describe` 的结果位于 PR Body 的 `PR-Agent 摘要` 区域,用于概览本次变更。 -`/review` 和自动审查会通过原生 GitHub Review 发布,结果位于 Review 页签,标题固定为 `PR-Agent Code Review`。审查摘要包含可点击的 `文件:行号` 链接;可定位到本次变更的具体问题会在对应代码行以行内评论呈现,无法行内定位的问题仍通过摘要中的链接呈现。 +`/review` 和自动审查会通过原生 GitHub Review 发布,结果位于 Review 页签,标题固定为 `PR-Agent Code Review`。可定位到本次变更的问题会在对应代码行以行内评论呈现,并使用 high、medium 或 low 风险标识。 -审查不会额外创建专用的摘要评论。未发现需要处理的问题时,Review 会显示: - -> 本次变更无需提出审查意见,暂无其他反馈。 +Review 摘要会自然概括本次变更和整体审查结论,不会复制行内评论。未发现需要处理的问题时,摘要会概括变更并自然说明暂无其他反馈。审查不会额外创建专用的 issue comment 摘要。 ## 安全边界 -workflow 使用固定 digest 的 PR-Agent 容器镜像,不使用浮动标签。自动审查通过 `pull_request_target` 在目标仓库上下文中读取 PR 信息,但不会 checkout 或执行 PR 分支代码。 - -权限保持最小化:只授予读取仓库内容所需的 `contents: read`,以及更新 PR Body、发布 Review 和回复 PR 评论所需的写权限;不会向仓库推送代码或创建提交。 +自动审查通过 `pull_request_target` 在目标仓库上下文中读取 PR 信息,并使用共享 runner 的 `latest` 镜像,但不会 checkout 或执行 PR 分支代码。权限保持最小化:只授予读取仓库内容所需的 `contents: read`,以及更新 PR Body、发布 Review 和回复 PR 评论所需的写权限;不会向仓库推送代码或创建提交。 diff --git a/requirements.in b/requirements.in index fbeaedc9..3f598b72 100644 --- a/requirements.in +++ b/requirements.in @@ -1,4 +1,4 @@ -moviepilot-rust~=0.2.2 +moviepilot-rust~=0.2.3 pydantic>=2.13.4,<3.0.0 pydantic-settings>=2.14.1,<3.0.0 SQLAlchemy~=2.0.50 @@ -80,6 +80,8 @@ langchain~=1.3.9 langchain-core~=1.4.7 langchain-community~=0.4.2 langchain-anthropic~=1.4.6 +langchain-aws~=1.6.2 +boto3~=1.42.42 langchain-openai~=1.3.2 langchain-google-genai~=4.2.5 langchain-deepseek~=1.1.0 diff --git a/tests/test_lifecycle_shutdown.py b/tests/test_lifecycle_shutdown.py index c984a3f5..34b9c7d4 100644 --- a/tests/test_lifecycle_shutdown.py +++ b/tests/test_lifecycle_shutdown.py @@ -1,16 +1,28 @@ import asyncio +import signal +import threading from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import FastAPI -from app.startup import lifecycle +from app.startup import lifecycle, modules_initializer +from app.utils import http as http_utils -def test_lifespan_closes_logger_when_early_shutdown_step_fails(monkeypatch): - """前置关闭步骤失败时仍应关闭 Logger""" +def _assert_completed_once(mock: MagicMock) -> None: + if isinstance(mock, AsyncMock): + mock.assert_awaited_once_with() + else: + mock.assert_called_once_with() + + +def _patch_lifespan(monkeypatch, *, failing_step: str | None = None) -> dict: + """隔离 lifespan 的外部依赖,并按名称注入一个关闭失败""" monkeypatch.setattr(lifecycle.settings, "MOVIEPILOT_SAFE_MODE", False) monkeypatch.setattr(lifecycle.global_vars, "set_loop", MagicMock()) + monkeypatch.setattr(lifecycle.global_vars, "stop_system", MagicMock()) + for name in ( "init_routers", "init_modules", @@ -19,28 +31,375 @@ def test_lifespan_closes_logger_when_early_shutdown_step_fails(monkeypatch): "init_monitor", "init_command", "init_workflow", + ): + monkeypatch.setattr(lifecycle, name, MagicMock()) + + system_chain = MagicMock() + monkeypatch.setattr(lifecycle, "SystemChain", MagicMock(return_value=system_chain)) + monkeypatch.setattr(lifecycle, "init_extra", AsyncMock()) + + shutdown_steps = { + "backup_plugins": system_chain.backup_plugins, + "stop_workflow": MagicMock(), + "stop_command": MagicMock(), + "stop_monitor": MagicMock(), + "stop_scheduler": MagicMock(), + "stop_plugins": MagicMock(), + "stop_modules": AsyncMock(), + "close_http": AsyncMock(), + } + for name in ( "stop_workflow", "stop_command", "stop_monitor", "stop_scheduler", "stop_plugins", ): - monkeypatch.setattr(lifecycle, name, MagicMock()) + monkeypatch.setattr(lifecycle, name, shutdown_steps[name]) + monkeypatch.setattr(lifecycle, "stop_modules", shutdown_steps["stop_modules"]) + monkeypatch.setattr( + lifecycle, + "aclose_shared_async_transports", + shutdown_steps["close_http"], + ) + + if failing_step: + shutdown_steps[failing_step].side_effect = RuntimeError( + f"{failing_step} failed" + ) - system_chain = MagicMock() - system_chain.backup_plugins.side_effect = RuntimeError("backup failed") - monkeypatch.setattr(lifecycle, "SystemChain", MagicMock(return_value=system_chain)) - monkeypatch.setattr(lifecycle, "init_extra", AsyncMock()) - monkeypatch.setattr(lifecycle, "stop_modules", AsyncMock()) - monkeypatch.setattr(lifecycle, "aclose_shared_async_transports", AsyncMock()) logger_shutdown = MagicMock() monkeypatch.setattr(lifecycle.LoggerManager, "shutdown", logger_shutdown) + shutdown_steps["logger"] = logger_shutdown + return shutdown_steps + + +@pytest.mark.parametrize( + "failing_step", + [ + "backup_plugins", + "stop_workflow", + "stop_command", + "stop_monitor", + "stop_scheduler", + "stop_plugins", + "stop_modules", + "close_http", + ], +) +def test_lifespan_continues_after_each_shutdown_owner_failure( + monkeypatch, + failing_step, +): + """任一关闭阶段失败都不能跳过后续资源所有者""" + shutdown_steps = _patch_lifespan(monkeypatch, failing_step=failing_step) async def run_lifespan(): - with pytest.raises(RuntimeError, match="backup failed"): - async with lifecycle.lifespan(FastAPI()): - pass + async with lifecycle.lifespan(FastAPI()): + pass asyncio.run(run_lifespan()) - logger_shutdown.assert_called_once_with() + lifecycle.global_vars.stop_system.assert_called_once_with() + for step in shutdown_steps.values(): + _assert_completed_once(step) + + +def test_uvicorn_signal_publishes_stop_before_server_exit(monkeypatch): + """Uvicorn 接管系统信号时必须先发布协作停止标志""" + from app import main + + calls = [] + monkeypatch.setattr(main.global_vars, "stop_system", lambda: calls.append("stop")) + monkeypatch.setattr( + main.uvicorn.Server, + "handle_exit", + lambda _self, _sig, _frame: calls.append("uvicorn"), + ) + + server = object.__new__(main.MoviePilotServer) + server.handle_exit(signal.SIGTERM, None) + + assert calls == ["stop", "uvicorn"] + + +def test_application_preserves_stop_requested_before_startup(monkeypatch): + """启动流程不能清除初始化前已经发布的退出请求""" + from app import main + + stop_event = threading.Event() + stop_event.set() + monkeypatch.setattr(main.global_vars, "STOP_EVENT", stop_event) + calls = [] + monkeypatch.setattr( + main.signal, + "signal", + lambda *_args: calls.append("signal"), + ) + monkeypatch.setattr(main, "start_tray", lambda: calls.append("tray")) + monkeypatch.setattr(main, "init_db", lambda: calls.append("init_db")) + monkeypatch.setattr(main, "update_db", lambda: calls.append("update_db")) + monkeypatch.setattr(main.Server, "run", lambda: calls.append("server")) + + main.run_application() + + assert stop_event.is_set() + assert calls == [ + "signal", + "signal", + "tray", + "init_db", + "update_db", + "server", + ] + + +def test_uvicorn_preserves_stop_requested_before_serve(monkeypatch): + """Uvicorn 启动不能清除数据库初始化阶段已经发布的停止请求""" + from app import main + + stop_event = threading.Event() + monkeypatch.setattr(main.global_vars, "STOP_EVENT", stop_event) + main.global_vars.stop_system() + + async def serve(_self, sockets=None): + assert main.global_vars.is_system_stopped + + monkeypatch.setattr(main.uvicorn.Server, "serve", serve) + server = object.__new__(main.MoviePilotServer) + asyncio.run(server.serve()) + + +@pytest.mark.parametrize("endpoint_name", ["restart_system", "upgrade_system"]) +@pytest.mark.parametrize( + "initially_stopped", + [False, True], + ids=["running", "stopping"], +) +def test_restart_endpoint_failure_preserves_stop_state( + monkeypatch, + endpoint_name, + initially_stopped, +): + """重启或升级失败不能发布或撤销停止请求""" + from app.api.endpoints import system + + stop_event = threading.Event() + if initially_stopped: + stop_event.set() + monkeypatch.setattr(system.global_vars, "STOP_EVENT", stop_event) + monkeypatch.setattr(system.SystemHelper, "can_restart", MagicMock(return_value=True)) + monkeypatch.setattr( + system.SystemHelper, + "restart" if endpoint_name == "restart_system" else "upgrade", + MagicMock(return_value=(False, "restart failed")), + ) + + if endpoint_name == "restart_system": + response = system.restart_system(None) + else: + response = system.upgrade_system(None, None) + + assert not response.success + assert stop_event.is_set() is initially_stopped + + +def test_command_restart_failure_does_not_publish_stop_request(monkeypatch): + """命令重启失败时进程仍在运行,不能提前发布停止请求""" + from app.chain.system import SystemChain + from app.core.config import global_vars + + stop_event = threading.Event() + monkeypatch.setattr(global_vars, "STOP_EVENT", stop_event) + monkeypatch.setattr(SystemChain, "backup_plugins", MagicMock()) + restart = MagicMock(return_value=(False, "restart failed")) + monkeypatch.setattr("app.chain.system.SystemHelper.restart", restart) + + chain = object.__new__(SystemChain) + chain.restart(channel=None, userid=None) + + restart.assert_called_once_with() + assert not stop_event.is_set() + + +def test_stop_modules_continues_after_internal_owner_failures(monkeypatch): + """模块关闭编排中的多个失败不能阻断其余清理""" + stop_agent = AsyncMock(side_effect=RuntimeError("agent failed")) + monkeypatch.setattr(modules_initializer, "stop_agent", stop_agent) + dependencies = _patch_module_shutdown_dependencies(monkeypatch) + dependencies["module"].side_effect = RuntimeError("module failed") + + asyncio.run(modules_initializer.stop_modules()) + + stop_agent.assert_awaited_once_with() + for dependency in dependencies.values(): + _assert_completed_once(dependency) + + +def _patch_module_shutdown_dependencies(monkeypatch) -> dict: + """替换 stop_modules 的资源所有者,避免测试启动真实后台服务""" + dependencies = {} + for name, method_name in ( + ("ModuleManager", "stop"), + ("EventManager", "stop"), + ("DisplayHelper", "stop"), + ("DohHelper", "shutdown"), + ("ThreadHelper", "shutdown"), + ("RedisHelper", "close"), + ): + instance = MagicMock() + setattr(instance, method_name, MagicMock()) + monkeypatch.setattr( + modules_initializer, + name, + MagicMock(return_value=instance), + ) + key = name.removesuffix("Helper").removesuffix("Manager").lower() + dependencies[key] = getattr(instance, method_name) + + for name in ("stop_message", "stop_frontend", "clear_temp"): + dependency = MagicMock() + monkeypatch.setattr(modules_initializer, name, dependency) + dependencies[name] = dependency + + async_redis = MagicMock() + async_redis.close = AsyncMock() + monkeypatch.setattr( + modules_initializer, + "AsyncRedisHelper", + MagicMock(return_value=async_redis), + ) + dependencies["async_redis"] = async_redis.close + close_database = AsyncMock() + monkeypatch.setattr(modules_initializer, "close_database", close_database) + dependencies["close_database"] = close_database + return dependencies + + +def test_shared_http_close_waits_for_real_lru_eviction(monkeypatch): + """最终 HTTP 关闭必须等待真实 LRU 淘汰任务并消费其异常""" + + class FakeTransport: + created = [] + + def __init__(self, **_kwargs): + self.close_started = asyncio.Event() + self.release_close = asyncio.Event() + self.closed = False + self.fail_on_close = not self.created + if not self.fail_on_close: + self.release_close.set() + self.created.append(self) + + async def aclose(self): + self.close_started.set() + await self.release_close.wait() + self.closed = True + if self.fail_on_close: + raise RuntimeError("eviction close failed") + + monkeypatch.setattr(http_utils, "_MAX_SHARED_TRANSPORTS_PER_LOOP", 1) + monkeypatch.setattr(http_utils.httpx, "AsyncHTTPTransport", FakeTransport) + debug = MagicMock() + monkeypatch.setattr(http_utils.logger, "debug", debug) + + async def run_test(): + transport_kwargs = { + "proxy": None, + "verify": True, + "http2": False, + "max_keepalive_connections": 1, + "max_connections": 1, + } + evicted_transport = http_utils._get_shared_async_transport( + **transport_kwargs, + keepalive_expiry=1, + ) + active_transport = http_utils._get_shared_async_transport( + **transport_kwargs, + keepalive_expiry=2, + ) + await asyncio.wait_for(evicted_transport.close_started.wait(), timeout=1) + + loop = asyncio.get_running_loop() + with http_utils._shared_async_transports_lock: + eviction_tasks = [ + task + for task in http_utils._pending_eviction_tasks + if task.get_loop() is loop + ] + assert len(eviction_tasks) == 1 + + close_task = asyncio.create_task(http_utils.aclose_shared_async_transports()) + await asyncio.sleep(0) + try: + assert not close_task.done() + evicted_transport.release_close.set() + await close_task + await asyncio.sleep(0) + assert eviction_tasks[0].done() + assert evicted_transport.closed + assert active_transport.closed + with http_utils._shared_async_transports_lock: + assert not any( + task.get_loop() is loop + for task in http_utils._pending_eviction_tasks + ) + finally: + evicted_transport.release_close.set() + active_transport.release_close.set() + await asyncio.gather(close_task, return_exceptions=True) + await http_utils.aclose_shared_async_transports() + + asyncio.run(run_test()) + + debug.assert_any_call( + "LRU 淘汰共享 transport 时关闭失败: " + "RuntimeError('eviction close failed')" + ) + + +def test_shared_http_close_ignores_eviction_from_other_loop(): + """当前事件循环关闭不能等待其他循环持有的淘汰任务""" + ready = threading.Event() + release = threading.Event() + failures = [] + state = {} + + def run_foreign_loop(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + async def delayed_close(): + while not release.is_set(): + await asyncio.sleep(0.01) + + task = loop.create_task(delayed_close()) + state["task"] = task + with http_utils._shared_async_transports_lock: + http_utils._pending_eviction_tasks.add(task) + task.add_done_callback(http_utils._discard_pending_eviction_task) + ready.set() + try: + loop.run_until_complete(task) + loop.run_until_complete(asyncio.sleep(0)) + except BaseException as err: + failures.append(err) + finally: + with http_utils._shared_async_transports_lock: + http_utils._pending_eviction_tasks.discard(task) + loop.close() + + thread = threading.Thread(target=run_foreign_loop) + thread.start() + try: + assert ready.wait(timeout=2) + asyncio.run(http_utils.aclose_shared_async_transports()) + assert thread.is_alive() + assert not state["task"].done() + finally: + release.set() + thread.join(timeout=2) + + assert not thread.is_alive() + assert not failures diff --git a/tests/test_llm_provider_bedrock.py b/tests/test_llm_provider_bedrock.py new file mode 100644 index 00000000..c3a6f265 --- /dev/null +++ b/tests/test_llm_provider_bedrock.py @@ -0,0 +1,392 @@ +"""Amazon Bedrock provider 的凭证解析、Region 提取与运行时解析测试""" + +import asyncio +import time +from unittest.mock import MagicMock, patch + +import pytest + +from app.agent.llm.provider import ( + LLMProviderAuthError, + LLMProviderManager, +) + + +@pytest.fixture(autouse=True) +def _reset_manager_singleton(): + """每个用例前后清理 LLMProviderManager 单例,避免缓存互相污染""" + LLMProviderManager._instances.clear() + yield + LLMProviderManager._instances.clear() + + +def test_bedrock_provider_registered(): + manager = LLMProviderManager() + spec = manager.get_provider("amazon-bedrock") + + assert spec.runtime == "bedrock" + assert spec.model_list_strategy == "bedrock" + assert spec.base_url_editable is True + assert spec.default_base_url == "https://bedrock-runtime.us-east-1.amazonaws.com" + preset_ids = {preset.id for preset in spec.base_url_presets} + assert "bedrock-us-east-1" in preset_ids + assert "bedrock-ap-northeast-1" in preset_ids + + +def test_parse_bedrock_credentials_bearer_api_key(): + credentials = LLMProviderManager._parse_bedrock_credentials( + "bedrock-api-key-abcdef123456" + ) + + assert credentials["auth_scheme"] == "bearer" + assert credentials["bearer_token"] == "bedrock-api-key-abcdef123456" + + +def test_parse_bedrock_credentials_sigv4_ak_sk(): + credentials = LLMProviderManager._parse_bedrock_credentials( + "AKIAIOSFODNN7EXAMPLE:wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" + ) + + assert credentials["auth_scheme"] == "sigv4" + assert credentials["access_key_id"] == "AKIAIOSFODNN7EXAMPLE" + assert credentials["secret_access_key"] == "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY" + assert "session_token" not in credentials + + +def test_parse_bedrock_credentials_sigv4_with_session_token(): + credentials = LLMProviderManager._parse_bedrock_credentials( + "ASIAIOSFODNN7EXAMPLE:secret/key:session-token-value" + ) + + assert credentials["auth_scheme"] == "sigv4" + assert credentials["session_token"] == "session-token-value" + + +def test_parse_bedrock_credentials_empty_rejected(): + with pytest.raises(LLMProviderAuthError): + LLMProviderManager._parse_bedrock_credentials("") + + +def test_parse_bedrock_credentials_malformed_colon_rejected(): + with pytest.raises(LLMProviderAuthError): + LLMProviderManager._parse_bedrock_credentials("AKIA123:") + + +def test_extract_bedrock_region_from_base_url(): + """应从标准、FIPS 与 PrivateLink Bedrock 端点提取 Region""" + extract = LLMProviderManager._extract_bedrock_region + + assert extract("https://bedrock-runtime.us-east-1.amazonaws.com") == "us-east-1" + assert extract("https://bedrock-runtime.ap-northeast-1.amazonaws.com/") == "ap-northeast-1" + assert extract("https://bedrock-runtime.mx-central-1.amazonaws.com") == "mx-central-1" + assert extract("https://bedrock.eu-central-1.amazonaws.com") == "eu-central-1" + # FIPS 与 PrivateLink(VPCE)端点同样能识别 Region + assert extract("https://bedrock-runtime-fips.us-east-1.amazonaws.com") == "us-east-1" + assert ( + extract("https://vpce-0abc123-xyz.bedrock-runtime.us-west-2.vpce.amazonaws.com") + == "us-west-2" + ) + # 无法识别时回退默认 Region + assert extract("https://example.com/us-west-2") == "us-east-1" + assert extract("https://example.com?region=.us-west-2.") == "us-east-1" + assert extract("https://example.com") == "us-east-1" + assert extract(None) == "us-east-1" + assert extract("") == "us-east-1" + + +def test_bedrock_endpoint_url_passthrough(): + """自定义 Bedrock 端点应透传,标准端点交由 boto3 推导""" + resolve = LLMProviderManager._bedrock_endpoint_url + + # 标准公有端点交由 boto3 推导,不显式透传 + assert resolve("bedrock-runtime", "https://bedrock-runtime.us-east-1.amazonaws.com") is None + assert resolve("bedrock", "https://bedrock.eu-central-1.amazonaws.com") is None + assert resolve("bedrock-runtime", None) is None + assert resolve("bedrock-runtime", "") is None + # FIPS / PrivateLink 等非标准端点需要显式生效 + assert ( + resolve("bedrock-runtime", "https://bedrock-runtime-fips.us-east-1.amazonaws.com") + == "https://bedrock-runtime-fips.us-east-1.amazonaws.com" + ) + assert ( + resolve( + "bedrock-runtime", + "https://vpce-0abc123-xyz.bedrock-runtime.us-west-2.vpce.amazonaws.com/", + ) + == "https://vpce-0abc123-xyz.bedrock-runtime.us-west-2.vpce.amazonaws.com" + ) + # runtime 端点填给控制面服务名时不匹配标准形态,同样透传 + assert ( + resolve("bedrock", "https://bedrock-runtime.us-east-1.amazonaws.com") + == "https://bedrock-runtime.us-east-1.amazonaws.com" + ) + + +def test_create_bedrock_client_uses_custom_endpoint(): + """创建 Bedrock 客户端时应把 PrivateLink 地址传给 boto3""" + manager = LLMProviderManager() + endpoint_url = ( + "https://vpce-0abc123-xyz.bedrock-runtime.us-west-2.vpce.amazonaws.com" + ) + client = MagicMock() + + with patch("boto3.client", return_value=client) as create_client: + result = manager.create_bedrock_client( + service_name="bedrock-runtime", + region="us-west-2", + credentials={ + "auth_scheme": "sigv4", + "access_key_id": "AKIAIOSFODNN7EXAMPLE", + "secret_access_key": "secret", + }, + base_url=endpoint_url, + use_proxy=False, + ) + + assert result is client + assert create_client.call_args.kwargs["endpoint_url"] == endpoint_url + + +def test_resolve_runtime_bedrock_bearer(): + manager = LLMProviderManager() + runtime = asyncio.run( + manager.resolve_runtime( + provider_id="amazon-bedrock", + model="global.anthropic.claude-haiku-4-5-20251001-v1:0", + api_key="bedrock-api-key-abc123", + base_url="https://bedrock-runtime.ap-northeast-1.amazonaws.com", + ) + ) + + assert runtime["runtime"] == "bedrock" + assert runtime["aws_region"] == "ap-northeast-1" + assert runtime["aws_auth"]["auth_scheme"] == "bearer" + + +def test_resolve_runtime_bedrock_sigv4_default_region(): + manager = LLMProviderManager() + runtime = asyncio.run( + manager.resolve_runtime( + provider_id="amazon-bedrock", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + api_key="AKIAIOSFODNN7EXAMPLE:wJalrXUtnFEMI/K7MDENG", + ) + ) + + assert runtime["runtime"] == "bedrock" + assert runtime["aws_region"] == "us-east-1" + assert runtime["aws_auth"]["auth_scheme"] == "sigv4" + assert runtime["aws_auth"]["access_key_id"] == "AKIAIOSFODNN7EXAMPLE" + + +def test_resolve_runtime_bedrock_missing_credentials_rejected(): + manager = LLMProviderManager() + with pytest.raises(LLMProviderAuthError): + asyncio.run( + manager.resolve_runtime( + provider_id="amazon-bedrock", + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + api_key=None, + ) + ) + + +def test_bedrock_model_matches_region(): + """目录模型应按 Profile 分区及裸模型 ON_DEMAND Region 过滤""" + matches = LLMProviderManager._bedrock_model_matches_region + + # 已知裸模型 ID 仅在其支持 ON_DEMAND 的 Region 保留 + assert matches("anthropic.claude-3-5-sonnet-20241022-v2:0", "us-west-2") + assert matches("anthropic.claude-3-5-sonnet-20241022-v2:0", "ap-southeast-2") + assert not matches("anthropic.claude-3-5-sonnet-20241022-v2:0", "ap-northeast-1") + assert not matches("anthropic.claude-sonnet-4-5-20250929-v1:0", "us-west-2") + assert not matches("amazon.nova-premier-v1:0", "ap-northeast-1") + assert not matches("meta.llama4-maverick-17b-instruct-v1:0", "ap-northeast-1") + # 已确认支持 ON_DEMAND 的裸模型与 global Profile 维持可用 + assert matches("amazon.nova-lite-v1:0", "ap-northeast-1") + assert matches("openai.gpt-oss-20b-1:0", "ap-northeast-1") + assert not matches("openai.gpt-oss-20b-1:0", "ap-southeast-1") + assert matches("global.anthropic.claude-sonnet-4-5-20250929-v1:0", "ap-northeast-1") + assert not matches( + "global.anthropic.claude-sonnet-4-5-20250929-v1:0", + "us-gov-west-1", + ) + # 地理前缀只在对应分区 Region 可调用 + assert matches("us.anthropic.claude-haiku-4-5-20251001-v1:0", "us-west-2") + assert not matches("us.anthropic.claude-haiku-4-5-20251001-v1:0", "ap-northeast-1") + assert not matches("us.anthropic.claude-haiku-4-5-20251001-v1:0", "us-gov-west-1") + assert matches("apac.amazon.nova-micro-v1:0", "ap-southeast-1") + assert not matches("apac.amazon.nova-micro-v1:0", "eu-central-1") + assert matches("eu.anthropic.claude-haiku-4-5-20251001-v1:0", "eu-central-1") + assert not matches("eu.anthropic.claude-haiku-4-5-20251001-v1:0", "us-east-1") + assert not matches("eu.anthropic.claude-haiku-4-5-20251001-v1:0", "eu-isoe-west-1") + + +def test_bedrock_au_profile_matches_melbourne_region(): + """AU Inference Profile 应允许从悉尼和墨尔本 Region 调用""" + matches = LLMProviderManager._bedrock_model_matches_region + + assert matches("au.amazon.nova-lite-v1:0", "ap-southeast-2") + assert matches("au.amazon.nova-lite-v1:0", "ap-southeast-4") + + +def test_list_models_bedrock_custom_endpoint_skips_control_plane(): + """自定义 runtime 端点刷新模型时应直接使用离线目录""" + manager = LLMProviderManager() + manager._models_dev_data = { + "amazon-bedrock": { + "id": "amazon-bedrock", + "name": "Amazon Bedrock", + "models": { + "global.anthropic.claude-sonnet-4-5-20250929-v1:0": { + "name": "Claude Sonnet 4.5 (Global)", + "limit": {"context": 200000, "output": 64000}, + }, + }, + } + } + manager._models_dev_loaded_at = time.time() + + with patch.object( + LLMProviderManager, + "create_bedrock_client", + side_effect=AssertionError("不应访问控制面"), + ): + models = asyncio.run( + manager._list_models_from_bedrock( + api_key="bedrock-api-key-runtime-only", + base_url="https://bedrock-runtime-fips.us-east-1.amazonaws.com", + use_proxy=False, + ) + ) + + assert [model["id"] for model in models] == [ + "global.anthropic.claude-sonnet-4-5-20250929-v1:0" + ] + + +def test_list_models_bedrock_keeps_control_plane_on_demand_models(): + """控制面返回的 ON_DEMAND 基础模型不应被静态降级规则遗漏""" + manager = LLMProviderManager() + client = MagicMock() + client.get_paginator.return_value.paginate.return_value = [ + { + "inferenceProfileSummaries": [ + { + "inferenceProfileId": ( + "global.anthropic.claude-sonnet-4-5-20250929-v1:0" + ), + "inferenceProfileName": "Claude Sonnet 4.5 (Global)", + "status": "ACTIVE", + } + ] + } + ] + client.list_foundation_models.return_value = { + "modelSummaries": [ + { + "modelId": "openai.gpt-oss-20b-1:0", + "modelName": "GPT OSS 20B", + "modelLifecycle": {"status": "ACTIVE"}, + }, + { + "modelId": "amazon.nova-lite-v1:0", + "modelName": "Nova Lite", + "modelLifecycle": {"status": "ACTIVE"}, + }, + ] + } + + with patch.object( + LLMProviderManager, "create_bedrock_client", return_value=client + ): + models = asyncio.run( + manager._list_models_from_bedrock( + api_key="bedrock-api-key-runtime-only", + base_url="https://bedrock-runtime.ap-northeast-1.amazonaws.com", + use_proxy=False, + ) + ) + + assert {model["id"] for model in models} == { + "amazon.nova-lite-v1:0", + "global.anthropic.claude-sonnet-4-5-20250929-v1:0", + "openai.gpt-oss-20b-1:0", + } + client.close.assert_called_once() + + +def test_list_models_bedrock_falls_back_to_models_dev_on_control_plane_denial(): + """控制面被拒(如 API Key 仅授权 bedrock-runtime)时降级 models.dev 目录""" + manager = LLMProviderManager() + # 预填 models.dev 内存缓存,降级路径不触发真实网络请求 + manager._models_dev_data = { + "amazon-bedrock": { + "id": "amazon-bedrock", + "name": "Amazon Bedrock", + "models": { + "anthropic.claude-3-5-sonnet-20241022-v2:0": { + "name": "Claude Sonnet 3.5 v2", + "limit": {"context": 200000, "output": 8192}, + }, + "amazon.nova-lite-v1:0": { + "name": "Nova Lite", + "limit": {"context": 300000, "output": 5000}, + }, + "openai.gpt-oss-20b-1:0": { + "name": "GPT OSS 20B", + "limit": {"context": 131072, "output": 16384}, + }, + "apac.amazon.nova-lite-v1:0": { + "name": "Nova Lite (APAC)", + "limit": {"context": 300000, "output": 5000}, + }, + "meta.llama4-maverick-17b-instruct-v1:0": { + "name": "Llama 4 Maverick", + "limit": {"context": 1000000, "output": 8192}, + }, + "anthropic.claude-sonnet-4-5-20250929-v1:0": { + "name": "Claude Sonnet 4.5", + "limit": {"context": 200000, "output": 64000}, + }, + "global.anthropic.claude-sonnet-4-5-20250929-v1:0": { + "name": "Claude Sonnet 4.5 (Global)", + "limit": {"context": 200000, "output": 64000}, + }, + "us.anthropic.claude-haiku-4-5-20251001-v1:0": { + "name": "Claude Haiku 4.5 (US)", + "limit": {"context": 200000, "output": 64000}, + }, + }, + } + } + manager._models_dev_loaded_at = time.time() + + denied_client = MagicMock() + denied_client.get_paginator.side_effect = Exception( + "AccessDeniedException: not authorized to perform bedrock:ListInferenceProfiles" + ) + + with patch.object( + LLMProviderManager, "create_bedrock_client", return_value=denied_client + ): + models = asyncio.run( + manager._list_models_from_bedrock( + api_key="bedrock-api-key-runtime-only", + base_url="https://bedrock-runtime.ap-northeast-1.amazonaws.com", + use_proxy=False, + ) + ) + + # 降级后仅保留东京 Region 可调用的裸模型与 Profile + model_ids = {m["id"] for m in models} + assert "anthropic.claude-3-5-sonnet-20241022-v2:0" not in model_ids + assert "amazon.nova-lite-v1:0" in model_ids + assert "openai.gpt-oss-20b-1:0" in model_ids + assert "apac.amazon.nova-lite-v1:0" in model_ids + assert "meta.llama4-maverick-17b-instruct-v1:0" not in model_ids + assert "global.anthropic.claude-sonnet-4-5-20250929-v1:0" in model_ids + assert "us.anthropic.claude-haiku-4-5-20251001-v1:0" not in model_ids + assert "anthropic.claude-sonnet-4-5-20250929-v1:0" not in model_ids + assert all(m["source"] == "models.dev" for m in models) + denied_client.close.assert_called_once() diff --git a/tests/test_subscribe_files_info.py b/tests/test_subscribe_files_info.py new file mode 100644 index 00000000..d672f1db --- /dev/null +++ b/tests/test_subscribe_files_info.py @@ -0,0 +1,143 @@ +"""订阅文件统计相关测试""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from app.chain.subscribe import SubscribeChain +from app.modules.filemanager import FileManagerModule +from app.schemas.mediaserver import ExistMediaInfo +from app.schemas.types import MediaType + + +def _build_subscribe(**overrides): + data = { + "id": 1, + "name": "Test Show", + "year": "2026", + "type": MediaType.TV.value, + "season": 1, + "tmdbid": None, + "doubanid": None, + "imdbid": None, + "tvdbid": None, + "bangumiid": None, + "episode_group": None, + "start_episode": 1, + "total_episode": 2, + } + data.update(overrides) + subscribe = SimpleNamespace(**data) + subscribe.to_dict = lambda: dict(data) + return subscribe + + +def _build_mediainfo(): + return SimpleNamespace( + type=MediaType.TV, + title="Test Show", + title_year="Test Show (2026)", + year="2026", + tmdb_id=None, + douban_id=None, + ) + + +def test_filemanager_media_exists_skips_local_when_server_specified(): + module = FileManagerModule() + mediainfo = _build_mediainfo() + + with patch.object(module, "media_files", return_value=[SimpleNamespace(path="/media/test.mkv")]) as media_files: + result = module.media_exists(mediainfo, server="Emby1") + + assert result is None + media_files.assert_not_called() + + +def test_subscribe_files_info_merges_multiple_mediaservers(): + subscribe = _build_subscribe(season=1, total_episode=2) + mediainfo = _build_mediainfo() + + def _media_exists_side_effect(*, mediainfo, server=None, **kwargs): + if server == "Emby1": + return ExistMediaInfo( + type=MediaType.TV, + seasons={1: [1]}, + server_type="emby", + server="Emby1", + itemid="emby-series", + ) + if server == "Jellyfin1": + return ExistMediaInfo( + type=MediaType.TV, + seasons={1: [1]}, + server_type="jellyfin", + server="Jellyfin1", + itemid="jf-series", + ) + return None + + helper = MagicMock() + helper.get_services.return_value = {"Emby1": object(), "Jellyfin1": object()} + + mediaserver_chain = MagicMock() + mediaserver_chain.get_play_url.side_effect = lambda server, item_id: f"https://{server}/item/{item_id}" + mediaserver_chain.get_season_episode_ids.side_effect = lambda server, item_id, season: {1: f"{item_id}-ep1"} + + chain = SubscribeChain() + with patch("app.chain.subscribe.DownloadHistoryOper") as download_oper, \ + patch.object(chain, "recognize_media", return_value=mediainfo), \ + patch.object(chain, "media_files", return_value=None), \ + patch.object(chain, "media_exists", side_effect=_media_exists_side_effect), \ + patch("app.chain.subscribe.MediaServerHelper", return_value=helper), \ + patch("app.chain.subscribe.MediaServerChain", return_value=mediaserver_chain), \ + patch("app.chain.subscribe.Subscribe", side_effect=lambda **kwargs: SimpleNamespace(**kwargs)): + download_oper.return_value.get_by_mediaid.return_value = [] + result = chain.subscribe_files_info(subscribe) + + library = result.episodes[1].library + servers = {item.server for item in library} + assert servers == {"Emby1", "Jellyfin1"} + assert all(str(item.file_path).startswith("https://") for item in library) + + +def test_subscribe_files_info_uses_season_zero_for_tv(): + subscribe = _build_subscribe(season=0, total_episode=1, start_episode=1) + mediainfo = _build_mediainfo() + captured_seasons = [] + + def _media_exists_side_effect(*, mediainfo, server=None, **kwargs): + if server == "Emby1": + return ExistMediaInfo( + type=MediaType.TV, + seasons={0: [1]}, + server_type="emby", + server="Emby1", + itemid="emby-special", + ) + return None + + def _get_season_episode_ids(server, item_id, season): + captured_seasons.append(season) + return {1: f"{item_id}-ep1"} + + helper = MagicMock() + helper.get_services.return_value = {"Emby1": object()} + + mediaserver_chain = MagicMock() + mediaserver_chain.get_play_url.return_value = "https://emby/item/1" + mediaserver_chain.get_season_episode_ids.side_effect = _get_season_episode_ids + + chain = SubscribeChain() + with patch("app.chain.subscribe.DownloadHistoryOper") as download_oper, \ + patch.object(chain, "recognize_media", return_value=mediainfo), \ + patch.object(chain, "media_files", return_value=None), \ + patch.object(chain, "media_exists", side_effect=_media_exists_side_effect), \ + patch("app.chain.subscribe.MediaServerHelper", return_value=helper), \ + patch("app.chain.subscribe.MediaServerChain", return_value=mediaserver_chain), \ + patch("app.chain.subscribe.Subscribe", side_effect=lambda **kwargs: SimpleNamespace(**kwargs)): + download_oper.return_value.get_by_mediaid.return_value = [] + result = chain.subscribe_files_info(subscribe) + + assert captured_seasons == [0] + assert len(result.episodes[1].library) == 1 + assert result.episodes[1].library[0].server == "Emby1" diff --git a/version.py b/version.py index 35c1b00c..79b83442 100644 --- a/version.py +++ b/version.py @@ -1,2 +1,2 @@ -APP_VERSION = 'v2.14.3' -FRONTEND_VERSION = 'v2.14.3' +APP_VERSION = 'v2.14.4' +FRONTEND_VERSION = 'v2.14.4'