diff --git a/backend/app/core/middlewares.py b/backend/app/core/middlewares.py index 3a861ff6d..967472fbd 100644 --- a/backend/app/core/middlewares.py +++ b/backend/app/core/middlewares.py @@ -2,13 +2,13 @@ import uuid from typing import Any +from starlette.datastructures import Headers from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from starlette.middleware.cors import CORSMiddleware from starlette.middleware.gzip import GZipMiddleware -from starlette.middleware.trustedhost import TrustedHostMiddleware from starlette.requests import Request -from starlette.responses import RedirectResponse, Response -from starlette.types import ASGIApp +from starlette.responses import PlainTextResponse, RedirectResponse, Response +from starlette.types import ASGIApp, Receive, Scope, Send from app.common.enums import RedisInitKeyConfig, SysParamKey from app.common.response import ErrorResponse @@ -126,11 +126,82 @@ async def dispatch(self, request: Request, call_next): return await call_next(request) -class CustomTrustedHostMiddleware(TrustedHostMiddleware): - """可信主机 Host 头校验中间件""" - - def __init__(self, app: ASGIApp) -> None: - super().__init__(app, allowed_hosts=settings.ALLOWED_HOSTS) +def host_of(value: str) -> str: + """从 Host 头里取出主机部分（小写、去端口）。 + + 修正 Starlette 的 ``host.split(":")[0]`` 在 IPv6 下的误判： + ``127.0.0.1:8000`` → ``127.0.0.1``；``[::1]:8000`` → ``::1``；裸 ``::1`` → ``::1``。 + """ + v = (value or "").strip().lower() + if v.startswith("["): + end = v.find("]") + return v[1:end] if end != -1 else v + # 只有「主机:端口」这一种情况该剥；裸 IPv6 里的多个冒号不是端口分隔符 + return v.rsplit(":", 1)[0] if v.count(":") == 1 else v + + +class CustomTrustedHostMiddleware: + """可信主机 Host 头校验中间件（纯 ASGI 实现）。 + + 在 Starlette ``TrustedHostMiddleware`` 基础上修正两处，并保持其余语义不变： + 1. **IPv6 字面量**：Starlette 用 ``host.split(":")[0]``，``[::1]:8000`` 会被切成 ``"["``， + 导致 ``::1`` 永远校验失败（本地开 ``ALLOW_LOCALHOST_HOSTS`` 时用 IPv6 访问会 400）； + 2. **大小写**：DNS 主机名不区分大小写，主机与允许清单统一转小写后比较。 + + 其余行为照旧：``*`` 放行全部；支持 ``*.example.com`` 形式通配；``www.`` 前缀可 301 重定向； + 允许清单里的条目也做同样归一（``[::1]`` 与 ``[::1]:8000`` 都按 ``::1`` 处理）。 + """ + + app: ASGIApp + allowed_hosts: list[str] + allow_any: bool + www_redirect: bool + + def __init__( + self, + app: ASGIApp, + allowed_hosts: list[str] | None = None, + www_redirect: bool = True, + ) -> None: + hosts = allowed_hosts if allowed_hosts is not None else settings.ALLOWED_HOSTS + normalized: list[str] = [] + for pattern in hosts: + assert "*" not in pattern[1:], "通配符只允许出现在最左侧（如 *.example.com）" + if pattern.startswith("*") and pattern != "*": + assert pattern.startswith("*."), "通配符只允许 *.example.com 形式" + normalized.append(host_of(pattern)) + self.app = app + self.allowed_hosts = normalized + self.allow_any = "*" in normalized + self.www_redirect = www_redirect + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if self.allow_any or scope["type"] not in ("http", "websocket"): + await self.app(scope, receive, send) + return + + host = host_of(Headers(scope=scope).get("host", "")) + found_www_redirect = False + for pattern in self.allowed_hosts: + if host == pattern or (pattern.startswith("*") and host.endswith(pattern[1:])): + await self.app(scope, receive, send) + return + if "www." + host == pattern: + found_www_redirect = True + + if scope["type"] == "websocket": + # WS 上不能发 HTTP 响应：按「策略违规」关闭连接 + await send({"type": "websocket.close", "code": 1008}) + return + response: Response + if found_www_redirect and self.www_redirect: + url = f"{scope.get('scheme', 'http')}://www.{host}" + path = scope.get("path", "/") + query = scope.get("query_string", b"").decode() + response = RedirectResponse(f"{url}{path}" + (f"?{query}" if query else ""), 301) + else: + response = PlainTextResponse("Invalid host header", 400) + await response(scope, receive, send) class CorrelationIdMiddleware(BaseHTTPMiddleware): diff --git a/backend/app/modules/ai/chat/controller.py b/backend/app/modules/ai/chat/controller.py index cba504a2b..cccdefc45 100644 --- a/backend/app/modules/ai/chat/controller.py +++ b/backend/app/modules/ai/chat/controller.py @@ -18,13 +18,23 @@ AiChatResponseSchema, AiModelConfigListResponse, AiModelConfigSchema, + AiModelOptionSchema, + AiModelProbeSchema, + AiModelTestResultSchema, + AiModelTestSchema, ChatQuerySchema, ChatSessionCreateSchema, ChatSessionOutSchema, ChatSessionQueryParam, ChatSessionUpdateSchema, ) -from .service import AiModelConfigService, ChatService, get_user_model_config +from .service import ( + AiModelConfigService, + ChatService, + get_user_model_config, + probe_available_models, + probe_model_connection, +) ChatRouter = APIRouter(route_class=OperationLogRoute, prefix="/chat", tags=["AI管理"]) @@ -178,6 +188,32 @@ async def activate_model_config_controller( return SuccessResponse(data=None, msg="已切换模型") +@ChatRouter.post( + "/model/available", + summary="拉取接口可用模型列表（按填写中的接口与密钥，不落库）", + response_model=ResponseSchema[list[AiModelOptionSchema]], +) +async def list_available_models_controller( + auth: Annotated[AuthSchema, Security(AuthPermission(["module_ai:chat:update"]))], + data: Annotated[AiModelProbeSchema, Body(description="接口地址与密钥")], +) -> JSONResponse: + result = await probe_available_models(data.base_url, data.api_key) + return SuccessResponse(data=result, msg="拉取模型列表成功") + + +@ChatRouter.post( + "/model/test", + summary="测试 AI 接口连通性（按填写中的参数探测，不落库）", + response_model=ResponseSchema[AiModelTestResultSchema], +) +async def test_model_connection_controller( + auth: Annotated[AuthSchema, Security(AuthPermission(["module_ai:chat:update"]))], + data: Annotated[AiModelTestSchema, Body(description="接口地址、密钥与模型 ID")], +) -> JSONResponse: + result = await probe_model_connection(data.base_url, data.api_key, data.model_id) + return SuccessResponse(data=result, msg="连通性测试通过") + + async def _send_error_and_close(websocket: WebSocket, message: str) -> None: """发送错误消息并关闭连接""" try: diff --git a/backend/app/modules/ai/chat/schema.py b/backend/app/modules/ai/chat/schema.py index 3ff327033..904d43ba7 100644 --- a/backend/app/modules/ai/chat/schema.py +++ b/backend/app/modules/ai/chat/schema.py @@ -89,14 +89,11 @@ class AiChatResponseSchema(BaseModel): action: dict[str, Any] | None = Field(None, description="建议执行的操作") -class AiModelConfigSchema(BaseModel): - """AI 模型配置项（创建/更新共用）""" +class AiModelEndpointSchema(BaseModel): + """OpenAI 兼容接口的接入信息（探测类请求共用：不落库）""" - name: str = Field(..., min_length=1, max_length=50, description="配置名称（用户可读）") base_url: str = Field(..., min_length=1, max_length=500, description="API Base URL，如 https://api.openai.com/v1") api_key: str = Field(..., min_length=1, max_length=500, description="API 密钥") - model_id: str = Field(..., min_length=1, max_length=100, description="模型 ID") - temperature: float = Field(0.7, ge=0.0, le=2.0, description="温度参数") @field_validator("base_url") @classmethod @@ -106,6 +103,47 @@ def validate_base_url(cls, v: str) -> str: raise ValueError("Base URL 必须以 http:// 或 https:// 开头") return v + +class AiModelProbeSchema(AiModelEndpointSchema): + """拉取模型列表的入参（按填写中的接口与密钥探测，不落库）""" + + +class AiModelTestSchema(AiModelEndpointSchema): + """连通性测试入参（按填写中的接口/密钥/模型探测，不落库）""" + + model_id: str = Field(..., min_length=1, max_length=100, description="模型 ID") + + @field_validator("model_id") + @classmethod + def validate_model_id(cls, v: str) -> str: + v = v.strip() + if not v: + raise ValueError("模型 ID 不能为空") + return v + + +class AiModelOptionSchema(BaseModel): + """接口返回的可用模型""" + + id: str = Field(..., description="模型 ID") + owned_by: str | None = Field(default=None, description="提供方（接口返回时才有）") + + +class AiModelTestResultSchema(BaseModel): + """连通性测试结果""" + + ok: bool = Field(..., description="是否连通") + model_id: str = Field(..., description="被测模型 ID") + latency_ms: int = Field(..., description="往返耗时（毫秒）") + + +class AiModelConfigSchema(AiModelEndpointSchema): + """AI 模型配置项（创建/更新共用）""" + + name: str = Field(..., min_length=1, max_length=50, description="配置名称（用户可读）") + model_id: str = Field(..., min_length=1, max_length=100, description="模型 ID") + temperature: float = Field(0.7, ge=0.0, le=2.0, description="温度参数") + @field_validator("model_id") @classmethod def validate_model_id(cls, v: str) -> str: diff --git a/backend/app/modules/ai/chat/service.py b/backend/app/modules/ai/chat/service.py index fba30b5b7..c35bcc64b 100644 --- a/backend/app/modules/ai/chat/service.py +++ b/backend/app/modules/ai/chat/service.py @@ -1,11 +1,13 @@ import asyncio import json +import time import uuid from collections.abc import AsyncGenerator from contextlib import asynccontextmanager from datetime import datetime from typing import Any +import httpx from agno.run.team import TeamRunOutput from agno.session.team import TeamSession from redis.asyncio import Redis @@ -530,3 +532,64 @@ async def set_active(self, config_id: str) -> None: ok = await set_active_model_config(self.redis, self._user_id, config_id) if not ok: raise CustomException(msg="模型配置不存在", code=RET.NOT_FOUND.code) + + +# ── OpenAI 兼容接口探测：拉取模型列表 / 连通性测试（不落库）─────────────── +# 只用于「按填写中的接口与密钥」做即时探测，配置落库仍走上面的 AiModelConfigService。 +# 安全：base_url 已由 schema 校验为 http(s)；仅取结果的模型 id / owned_by，不回显密钥。 + +_AI_PROBE_TIMEOUT = 20.0 + + +def _ai_auth_headers(api_key: str) -> dict[str, str]: + """构造 Bearer 鉴权头（密钥只放进请求头，不落日志）。""" + return {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} + + +async def probe_available_models(base_url: str, api_key: str) -> list[dict[str, Any]]: + """拉取接口可用的模型列表（GET {base_url}/models）。 + + 返回: + - list[dict[str, Any]]: [{id, owned_by}] + + 异常: + - CustomException: 接口不可达或返回非 2xx + """ + url = f"{base_url.rstrip('/')}/models" + try: + async with httpx.AsyncClient(timeout=_AI_PROBE_TIMEOUT) as client: + response = await client.get(url, headers=_ai_auth_headers(api_key)) + except httpx.HTTPError as e: + raise CustomException(msg=f"请求模型列表失败: {type(e).__name__}") from e + if response.status_code >= 400: + raise CustomException(msg=f"请求模型列表失败: HTTP {response.status_code}") + payload: dict[str, Any] = response.json() or {} + items = payload.get("data") or [] + return [ + {"id": str(item.get("id")), "owned_by": item.get("owned_by")} + for item in items + if isinstance(item, dict) and item.get("id") + ] + + +async def probe_model_connection(base_url: str, api_key: str, model_id: str) -> dict[str, Any]: + """连通性测试：发一次最小 chat/completions 请求（max_tokens=1）。 + + 返回: + - dict[str, Any]: {ok, model_id, latency_ms} + + 异常: + - CustomException: 接口不可达或返回非 2xx + """ + url = f"{base_url.rstrip('/')}/chat/completions" + body = {"model": model_id, "messages": [{"role": "user", "content": "ping"}], "max_tokens": 1} + started = time.perf_counter() + try: + async with httpx.AsyncClient(timeout=_AI_PROBE_TIMEOUT) as client: + response = await client.post(url, headers=_ai_auth_headers(api_key), json=body) + except httpx.HTTPError as e: + raise CustomException(msg=f"连通性测试失败: {type(e).__name__}") from e + latency_ms = int((time.perf_counter() - started) * 1000) + if response.status_code >= 400: + raise CustomException(msg=f"连通性测试失败: HTTP {response.status_code}") + return {"ok": True, "model_id": model_id, "latency_ms": latency_ms} diff --git a/backend/app/modules/task/storage/browse/controller.py b/backend/app/modules/task/storage/browse/controller.py index bf542b76d..d7815f9c1 100644 --- a/backend/app/modules/task/storage/browse/controller.py +++ b/backend/app/modules/task/storage/browse/controller.py @@ -110,6 +110,17 @@ async def list_storage_buckets_controller( return SuccessResponse(data=result, msg="查询桶列表成功") +@StorageBrowseRouter.post("/create-bucket", summary="创建存储桶（对象存储）", response_model=ResponseSchema[None]) +async def create_storage_bucket_controller( + auth: Annotated[AuthSchema, Security(AuthPermission(["module_task:storage:browse:update"]))], + db: Annotated[AsyncSession, Depends(db_getter)], + bucket: Annotated[str, Body(description="桶名（COS 需带 -appid 后缀）")], + source_id: Annotated[int | None, Body(description="存储源ID（不传使用默认存储源）")] = None, +) -> JSONResponse: + await StorageFileService(auth, db).create_bucket(source_id=source_id, bucket=bucket) + return SuccessResponse(msg="创建桶成功") + + @StorageBrowseRouter.post("/copy", summary="复制/移动文件", response_model=ResponseSchema[StoragePathResultSchema]) async def copy_or_move_storage_file_controller( auth: Annotated[AuthSchema, Security(AuthPermission(["module_task:storage:browse:update"]))], diff --git a/backend/app/modules/task/storage/browse/service.py b/backend/app/modules/task/storage/browse/service.py index f2a9996dd..da76b4d21 100644 --- a/backend/app/modules/task/storage/browse/service.py +++ b/backend/app/modules/task/storage/browse/service.py @@ -246,6 +246,28 @@ async def list_buckets(self, source_id: int | None) -> list[str]: finally: await adapter.close() + async def create_bucket(self, source_id: int | None, bucket: str) -> None: + """创建存储桶（仅对象存储协议）。 + + 桶名只做与协议无关的通用校验（非空/长度/非法字符）；各协议的命名细则 + （如 COS 的 ``-appid`` 后缀、S3 的区域约束）交给 SDK 报错。 + """ + name = (bucket or "").strip() + if not name: + raise CustomException(msg="桶名不能为空") + if len(name) > 63: + raise CustomException(msg="桶名长度不能超过 63 个字符") + if any(ch in name for ch in ("/", "\\", " ", "\t", "\n")): + raise CustomException(msg="桶名不能包含空格或斜杠") + config = await self._get_source(source_id) + if config.protocol not in _OBJECT_STORE_PROTOCOLS: + raise CustomException(msg="仅对象存储协议（S3/OSS/COS/OBS）支持创建桶") + adapter = StorageAdapterFactory.create(config) + try: + await adapter.create_bucket(name) + finally: + await adapter.close() + async def copy_or_move( self, source_id: int | None, diff --git a/backend/app/modules/task/storage/core/base.py b/backend/app/modules/task/storage/core/base.py index 76d484012..4f995d05f 100644 --- a/backend/app/modules/task/storage/core/base.py +++ b/backend/app/modules/task/storage/core/base.py @@ -185,6 +185,10 @@ def _sync_list_buckets(self) -> list[str]: """列出账号下全部存储桶（对象存储专用）。FTP/SFTP/LOCAL 等协议无需实现。""" raise NotImplementedError(f"{type(self).__name__} 不支持桶列表") + def _sync_create_bucket(self, bucket: str) -> None: + """创建存储桶（对象存储专用）。FTP/SFTP/LOCAL 等协议无需实现。""" + raise NotImplementedError(f"{type(self).__name__} 不支持创建桶") + def _multipart_settings(self) -> tuple[int, int, bool]: """返回 (分片大小, 并发数, 是否分片)，分片大小单位字节；并发按内存预算护栏收敛。 @@ -430,6 +434,10 @@ async def list_buckets(self) -> list[str]: """列出账号下全部存储桶（对象存储专用）。""" return await asyncio.to_thread(self._sync_list_buckets) + async def create_bucket(self, bucket: str) -> None: + """创建存储桶（对象存储专用；桶名规则由各协议校验）。""" + await asyncio.to_thread(self._sync_create_bucket, bucket) + async def list_recursive(self, prefix: str = "") -> list[StorageObject]: """递归列出目录树（含隐含目录条目）。""" return await asyncio.to_thread(self._sync_list_recursive, prefix) diff --git a/backend/app/modules/task/storage/core/cos_adapter.py b/backend/app/modules/task/storage/core/cos_adapter.py index cda43e535..3974cb455 100644 --- a/backend/app/modules/task/storage/core/cos_adapter.py +++ b/backend/app/modules/task/storage/core/cos_adapter.py @@ -1,6 +1,6 @@ -from typing import Any +from __future__ import annotations -from qcloud_cos import CosConfig, CosS3Client +from typing import TYPE_CHECKING, Any from app.core.exceptions import CustomException from app.core.logger import logger @@ -11,6 +11,12 @@ StoragePage, StorageProtocol, ) +from app.utils.lazy_sdk import lazy_module + +if TYPE_CHECKING: + import qcloud_cos +else: + qcloud_cos = lazy_module("qcloud_cos") class CosStorageAdapter(BaseStorageAdapter): @@ -30,7 +36,7 @@ def __init__(self, config) -> None: proxies["http"] = adv.http_proxy if adv.https_proxy: proxies["https"] = adv.https_proxy - cos_config = CosConfig( + cos_config = qcloud_cos.CosConfig( Region=self.config.region, SecretId=self.config.username or "", SecretKey=self.config.password or "", @@ -47,7 +53,7 @@ def __init__(self, config) -> None: KeepAlive=adv.keep_alive, AllowRedirects=adv.allow_redirects, ) - self.client = CosS3Client(cos_config, retry=adv.retry) + self.client = qcloud_cos.CosS3Client(cos_config, retry=adv.retry) self._adv = adv def _sync_test_connection(self) -> bool: @@ -267,6 +273,16 @@ def _sync_list_buckets(self) -> list[str]: logger.exception('COS 桶列表失败') raise + def _sync_create_bucket(self, bucket: str) -> None: + """创建桶（COS：桶名须带 ``-appid`` 后缀）。""" + try: + self.client.create_bucket(Bucket=bucket) + except CustomException: + raise + except Exception as e: + logger.exception('COS 创建桶失败') + raise CustomException(msg=f"COS 创建桶失败: {e}") from e + def _sync_get_url(self, remote_path: str, expire: int) -> str: """生成预签名 URL :param Bucket(string): 存储桶名称. diff --git a/backend/app/modules/task/storage/core/obs_adapter.py b/backend/app/modules/task/storage/core/obs_adapter.py index ffca09889..5d478a06b 100644 --- a/backend/app/modules/task/storage/core/obs_adapter.py +++ b/backend/app/modules/task/storage/core/obs_adapter.py @@ -1,7 +1,7 @@ -from datetime import UTC, datetime -from typing import Any +from __future__ import annotations -from obs import ObsClient +from datetime import UTC, datetime +from typing import TYPE_CHECKING, Any from app.core.exceptions import CustomException from app.core.logger import logger @@ -13,6 +13,12 @@ StorageProtocol, normalize_endpoint, ) +from app.utils.lazy_sdk import lazy_module + +if TYPE_CHECKING: + import obs +else: + obs = lazy_module("obs") class ObsStorageAdapter(BaseStorageAdapter): @@ -123,7 +129,7 @@ def __init__(self, config) -> None: if not self.config.endpoint: raise CustomException(msg="OBS 存储源必须配置 endpoint") adv = ObsAdvancedConfig(**self.config.advanced_config) - self.client = ObsClient( + self.client = obs.ObsClient( access_key_id=self.config.username or "", secret_access_key=self.config.password or "", server=normalize_endpoint(self.config.endpoint, self.config.scheme), @@ -368,6 +374,20 @@ def _sync_list_buckets(self) -> list[str]: logger.exception('OBS 桶列表失败') raise + def _sync_create_bucket(self, bucket: str) -> None: + """创建桶（OBS SDK 不抛异常，按返回状态码判断）。""" + try: + resp = self.client.createBucket(bucketName=bucket) + status = int(getattr(resp, "status", 0) or 0) + if status >= 300: + reason = getattr(resp, "errorMessage", "") or status + raise CustomException(msg=f"OBS 创建桶失败: {reason}") + except CustomException: + raise + except Exception as e: + logger.exception('OBS 创建桶失败') + raise CustomException(msg=f"OBS 创建桶失败: {e}") from e + def _sync_get_url(self, remote_path: str, expire: int) -> str: try: resp = self.client.createSignedUrl("GET", bucketName=self._require_bucket(), objectKey=remote_path, expires=expire) diff --git a/backend/app/modules/task/storage/core/oss_adapter.py b/backend/app/modules/task/storage/core/oss_adapter.py index b7719fdcf..1be84e085 100644 --- a/backend/app/modules/task/storage/core/oss_adapter.py +++ b/backend/app/modules/task/storage/core/oss_adapter.py @@ -1,6 +1,7 @@ -from datetime import UTC, datetime, timedelta +from __future__ import annotations -import alibabacloud_oss_v2 as oss +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING from app.core.exceptions import CustomException from app.core.logger import logger @@ -12,6 +13,12 @@ StorageProtocol, normalize_endpoint, ) +from app.utils.lazy_sdk import lazy_module + +if TYPE_CHECKING: + import alibabacloud_oss_v2 as oss +else: + oss = lazy_module("alibabacloud_oss_v2") class OssStorageAdapter(BaseStorageAdapter): @@ -201,6 +208,16 @@ def _sync_list_buckets(self) -> list[str]: logger.exception('OSS 桶列表失败') raise + def _sync_create_bucket(self, bucket: str) -> None: + """创建桶（OSS v2：put_bucket）。""" + try: + self.client.put_bucket(oss.PutBucketRequest(bucket=bucket)) + except CustomException: + raise + except Exception as e: + logger.exception('OSS 创建桶失败') + raise CustomException(msg=f"OSS 创建桶失败: {e}") from e + def _sync_get_url(self, remote_path: str, expire: int) -> str | None: try: result = self.client.presign( diff --git a/backend/app/modules/task/storage/core/s3_adapter.py b/backend/app/modules/task/storage/core/s3_adapter.py index 5b7ff8797..653a4e249 100644 --- a/backend/app/modules/task/storage/core/s3_adapter.py +++ b/backend/app/modules/task/storage/core/s3_adapter.py @@ -1,9 +1,6 @@ -from typing import Any +from __future__ import annotations -import boto3 -from boto3.s3.transfer import TransferConfig -from botocore.config import Config -from botocore.exceptions import ClientError +from typing import TYPE_CHECKING, Any from app.core.exceptions import CustomException from app.core.logger import logger @@ -15,6 +12,19 @@ StorageProtocol, normalize_endpoint, ) +from app.utils.lazy_sdk import lazy_module + +if TYPE_CHECKING: + import boto3 + import boto3.s3.transfer as boto_transfer + import botocore.config as botocore_config + import botocore.exceptions as botocore_exceptions +else: + # 冷启动不拉 SDK：只在真正使用 S3 协议时才导入（见 app/utils/lazy_sdk.py） + boto3 = lazy_module("boto3") + boto_transfer = lazy_module("boto3.s3.transfer") + botocore_config = lazy_module("botocore.config") + botocore_exceptions = lazy_module("botocore.exceptions") class S3StorageAdapter(BaseStorageAdapter): @@ -38,7 +48,7 @@ def __init__(self, config) -> None: s3_options["use_accelerate_endpoint"] = True if adv.us_east_1_regional_endpoint: s3_options["us_east_1_regional_endpoint"] = adv.us_east_1_regional_endpoint - config = Config( + config = botocore_config.Config( signature_version=adv.signature_version, retries={"max_attempts": adv.max_attempts, "mode": adv.retries_mode}, connect_timeout=adv.connect_timeout, @@ -65,7 +75,7 @@ def _sync_test_connection(self) -> bool: try: self.client.head_bucket(Bucket=self._require_bucket()) return True - except ClientError as e: + except botocore_exceptions.ClientError as e: code = e.response.get("Error", {}).get("Code", "") # 403 表示凭据有效但无权限查看桶，连接本身是通的 if code == "403": @@ -81,7 +91,7 @@ def _sync_upload(self, local_path: str, remote_path: str) -> str: # 按端点分片参数配置：超过分片大小阈值即分片并发上传（并发受内存预算护栏约束）； # 流式传输（stream）时阈值为超大值，等效单次上传 part_size, concurrency, _ = self._multipart_settings() - transfer_config = TransferConfig( + transfer_config = boto_transfer.TransferConfig( multipart_threshold=part_size, multipart_chunksize=part_size, max_concurrency=concurrency, @@ -111,7 +121,7 @@ def _sync_exists(self, remote_path: str) -> bool: try: self.client.head_object(Bucket=self._require_bucket(), Key=remote_path) return True - except ClientError as e: + except botocore_exceptions.ClientError as e: if e.response.get("ResponseMetadata", {}).get("HTTPStatusCode") == 404: return False logger.warning(f"S3 head_object 失败: {e}") @@ -219,6 +229,20 @@ def _sync_list_buckets(self) -> list[str]: logger.exception('S3 桶列表失败') raise + def _sync_create_bucket(self, bucket: str) -> None: + """创建桶；AWS 规定 us-east-1 以外的区域必须带 LocationConstraint。""" + try: + kwargs: dict[str, Any] = {"Bucket": bucket} + region = (self.config.region or "").strip() + if region and region != "us-east-1": + kwargs["CreateBucketConfiguration"] = {"LocationConstraint": region} + self.client.create_bucket(**kwargs) + except CustomException: + raise + except Exception as e: + logger.exception('S3 创建桶失败') + raise CustomException(msg=f"S3 创建桶失败: {e}") from e + def _sync_get_url(self, remote_path: str, expire: int) -> str: try: return self.client.generate_presigned_url( diff --git a/backend/app/modules/task/storage/core/sftp_adapter.py b/backend/app/modules/task/storage/core/sftp_adapter.py index f10c15179..6e0acdbfe 100644 --- a/backend/app/modules/task/storage/core/sftp_adapter.py +++ b/backend/app/modules/task/storage/core/sftp_adapter.py @@ -5,15 +5,23 @@ 同步操作经基类 asyncio.to_thread 包装，不阻塞事件循环。 用完由调用方调用 close() 释放连接。 """ + +from __future__ import annotations + import os import stat from datetime import UTC, datetime - -import paramiko +from typing import TYPE_CHECKING from app.core.exceptions import CustomException from app.core.logger import logger from app.modules.task.storage.core.base import BaseStorageAdapter, SftpAdvancedConfig, StorageObject, StorageProtocol +from app.utils.lazy_sdk import lazy_module + +if TYPE_CHECKING: + import paramiko +else: + paramiko = lazy_module("paramiko") class SftpStorageAdapter(BaseStorageAdapter): diff --git a/backend/app/modules/task/storage/node/config_io.py b/backend/app/modules/task/storage/node/config_io.py new file mode 100644 index 000000000..aa82ed232 --- /dev/null +++ b/backend/app/modules/task/storage/node/config_io.py @@ -0,0 +1,109 @@ +"""存储源配置导出/导入（**不含凭据**）。 + +为什么不含 ``password``：导出文件最常见的去处是群聊/工单/博客，带上密码等于把凭据抄一份出去； +导入后让用户重填一次，比事后收不回来强。这个决定改之前请先想清楚后果。 +``advanced_config`` 同样不导出（它是 SDK 参数 JSON，可能含代理等敏感信息）。 + +设计成纯函数，便于不依赖数据库直接测试；数据库读写留在 ``StorageNodeService``。 +""" + +from datetime import UTC, datetime +from typing import Any + +from .model import StorageNodeModel + +EXPORT_VERSION = 1 + +# 可导出的字段（纯配置；password / advanced_config / is_default 一律不导出） +EXPORT_FIELDS: tuple[str, ...] = ( + "name", + "protocol", + "host", + "port", + "username", + "bucket", + "endpoint", + "region", + "path_prefix", + "is_secure", + "implicit_tls", + "scheme", + "encrypt_type", + "connection_mode", + "encoding", + "multipart_part_size", + "multipart_concurrency", + "multipart_memory_budget", + "status", + "description", +) + +# 导入时的协议白名单（与 StorageProtocol 保持一致，避免把脏值写进库） +_VALID_PROTOCOLS = frozenset( + {"ftp", "ftps", "sftp", "s3", "obs", "oss", "cos", "local"} +) + + +def build_export(nodes: list[StorageNodeModel]) -> dict[str, Any]: + """把存储源列表打包成可传阅的 JSON 结构（不含凭据）。""" + return { + "version": EXPORT_VERSION, + "exported_at": datetime.now(UTC).isoformat(timespec="seconds"), + "nodes": [{field: getattr(node, field) for field in EXPORT_FIELDS} for node in nodes], + } + + +def plan_import( + data: dict[str, Any], existing_names: set[str] +) -> tuple[list[dict[str, Any]], int, int]: + """规划导入：按名称去重，**已存在的绝不覆盖**。 + + 返回: + - tuple[list[dict[str, Any]], int, int]: (待创建条目, 跳过数, 非法数) + """ + to_create: list[dict[str, Any]] = [] + skipped = 0 + invalid = 0 + seen = set(existing_names) + + items = data.get("nodes") if isinstance(data, dict) else None + for item in items or []: + if not isinstance(item, dict): + invalid += 1 + continue + name = str(item.get("name") or "").strip() + protocol = str(item.get("protocol") or "").strip().lower() + if not name or protocol not in _VALID_PROTOCOLS: + invalid += 1 + continue + if name in seen: + skipped += 1 + continue + payload: dict[str, Any] = { + field: item[field] for field in EXPORT_FIELDS if field in item and item[field] is not None + } + payload["name"] = name + payload["protocol"] = protocol + # port 在库里非空且无默认值：缺失时给协议默认端口，避免插入失败 + try: + payload["port"] = int(payload.get("port") or 0) or _default_port(protocol) + except (TypeError, ValueError): + payload["port"] = _default_port(protocol) + to_create.append(payload) + seen.add(name) + + return to_create, skipped, invalid + + +def _default_port(protocol: str) -> int: + """按协议给出默认端口。""" + return { + "ftp": 21, + "ftps": 21, + "sftp": 22, + "s3": 443, + "oss": 443, + "cos": 443, + "obs": 443, + "local": 0, + }.get(protocol, 21) diff --git a/backend/app/modules/task/storage/node/controller.py b/backend/app/modules/task/storage/node/controller.py index 5b7014495..cf8bec82f 100644 --- a/backend/app/modules/task/storage/node/controller.py +++ b/backend/app/modules/task/storage/node/controller.py @@ -1,4 +1,4 @@ -from typing import Annotated +from typing import Annotated, Any from fastapi import APIRouter, Body, Depends, Path, Query, Security, status from fastapi.responses import JSONResponse @@ -67,6 +67,28 @@ async def get_storage_source_list_controller( return SuccessResponse(data=result, msg="查询存储源列表成功") +@StorageNodeRouter.get("/export", summary="导出存储源配置（不含凭据）", response_model=ResponseSchema[dict[str, Any]]) +async def export_storage_node_config_controller( + auth: Annotated[AuthSchema, Security(AuthPermission(["module_task:storage:node:query"]))], + db: Annotated[AsyncSession, Depends(db_getter)], +) -> JSONResponse: + result: dict[str, Any] = await StorageNodeService(auth, db).export_config() + return SuccessResponse(data=result, msg="导出成功（不含凭据，导入后需重填密码）") + + +@StorageNodeRouter.post("/import", summary="导入存储源配置（不含凭据，按名称去重）", response_model=ResponseSchema[dict[str, int]]) +async def import_storage_node_config_controller( + auth: Annotated[AuthSchema, Security(AuthPermission(["module_task:storage:node:create"]))], + db: Annotated[AsyncSession, Depends(db_getter)], + data: Annotated[dict[str, Any], Body(description="导出的 JSON 结构（version/exported_at/nodes）")], +) -> JSONResponse: + result: dict[str, int] = await StorageNodeService(auth, db).import_config(data) + return SuccessResponse( + data=result, + msg=f"导入完成：新增 {result['created']} 个，跳过 {result['skipped']} 个（同名已存在），无效 {result['invalid']} 个", + ) + + @StorageNodeRouter.get("/detail/{id}", summary="查询存储源详情", response_model=ResponseSchema[StorageNodeOutSchema]) async def get_storage_source_detail_controller( auth: Annotated[AuthSchema, Security(AuthPermission(["module_task:storage:node:query"]))], diff --git a/backend/app/modules/task/storage/node/service.py b/backend/app/modules/task/storage/node/service.py index 8eede8046..15916b585 100644 --- a/backend/app/modules/task/storage/node/service.py +++ b/backend/app/modules/task/storage/node/service.py @@ -15,6 +15,7 @@ from app.modules.task.storage.core.factory import StorageAdapterFactory from app.utils.common_util import search_to_dict +from .config_io import build_export, plan_import from .crud import StorageNodeCRUD from .model import StorageNodeModel from .schema import StorageNodeCreateSchema, StorageNodeOutSchema, StorageNodeQueryParam, StorageNodeTestSchema, StorageNodeUpdateSchema @@ -76,6 +77,21 @@ def _build_config(obj: StorageNodeModel) -> StorageAdapterConfig: # ── 查询 ──────────────────────────────────────────────────────── + async def export_config(self) -> dict[str, Any]: + """导出全部存储源配置（不含凭据，可直接传阅）。""" + nodes = await self._crud().get_list() + return build_export(list(nodes)) + + async def import_config(self, data: dict[str, Any]) -> dict[str, int]: + """按名称去重导入存储源配置（不含凭据；已存在则跳过，绝不覆盖现有配置）。""" + existing = {node.name for node in await self._crud().get_list()} + to_create, skipped, invalid = plan_import(data, existing) + created = 0 + for payload in to_create: + await self._crud().create(data=payload) + created += 1 + return {"created": created, "skipped": skipped, "invalid": invalid} + async def detail(self, id: int) -> StorageNodeOutSchema: obj = await self._crud().get_or_404(id=id) return self._to_out(obj) diff --git a/backend/app/utils/lazy_sdk.py b/backend/app/utils/lazy_sdk.py new file mode 100644 index 000000000..2b7bd6e25 --- /dev/null +++ b/backend/app/utils/lazy_sdk.py @@ -0,0 +1,68 @@ +"""按需导入重型 SDK 的小工具（首次属性访问时才真正 import）。 + +为什么需要：存储适配器在 import 期就把 boto3 / botocore / qcloud_cos / obs / paramiko / +alibabacloud_oss_v2 全部拉进来，而适配器工厂为了「零改工厂新增协议」会在导入时引入全部适配器 —— +于是每次冷启动都白付这段导入时间（本仓库实测：导入工厂 913ms、6 个重型 SDK 全部进入 sys.modules）， +哪怕这次只用一个 SFTP 节点。 + +做法：模块级用代理对象占位，属性访问（如 ``boto3.client(...)``）时才 import 并缓存。 +所以各适配器里的调用写法基本不用改，只是 import 时机从「启动」挪到「首次用到该协议」。 + +注意（踩过才知道）： +1. **函数签名上的注解会在 def 时就求值**（如 ``-> paramiko.SFTPClient``），所以模块要加 + ``from __future__ import annotations``，让注解变成字符串不求值； +2. ``from X import Y`` 是 import 语句、拿不到代理的好处：要么改成 ``mod.Y``，要么在真正用到的 + 函数里局部 import；``except`` 子句里的异常类用 :func:`sdk_exception`（``except`` 表达式在 + 异常发生时才求值，不会在导入期拉起 SDK）； +3. 类型检查器/IDE 需要真模块：用 ``if TYPE_CHECKING:`` 分支放真 import，``else`` 分支放代理， + 否则注解里把 ``Any`` 当类型用会触发 pyright ``reportInvalidTypeForm``，整条链退化为 Any。 + +验证：导入工厂后，这些重型 SDK 不应出现在 ``sys.modules``；构造对应协议客户端后才应出现。 +""" + +from __future__ import annotations + +import importlib +from types import ModuleType +from typing import Any, ClassVar, cast, override + + +class LazyModule: + """模块代理：首次属性访问时才 import 真实模块，之后缓存复用。""" + + __slots__: ClassVar[tuple[str, ...]] = ("_module", "_name") + + def __init__(self, name: str) -> None: + self._name: str = name + self._module: ModuleType | None = None + + def _load(self) -> ModuleType: + if self._module is None: + self._module = importlib.import_module(self._name) + return self._module + + def __getattr__(self, item: str) -> Any: + return getattr(self._load(), item) + + @override + def __repr__(self) -> str: # pragma: no cover - 仅为调试可读 + state = "已导入" if self._module is not None else "未导入" + return f"" + + +def lazy_module(name: str) -> Any: + """占位一个模块；类型上按 Any 用，避免到处写 cast（真实类型由 SDK 提供）。""" + return LazyModule(name) + + +def sdk_exception(module_name: str, attr: str) -> type[BaseException]: + """按需取 SDK 的异常类，专供 ``except`` 子句： + + .. code-block:: python + + except sdk_exception("botocore.exceptions", "ClientError") as exc: + + ``except`` 的表达式在**异常发生时才求值**，所以不会在导入期拉起 SDK； + 而 ``from X import Y`` 那种写法是 import 语句，拿不到代理的好处，只能这样换掉。 + """ + return cast(type[BaseException], getattr(importlib.import_module(module_name), attr)) diff --git a/backend/tests/test_ai_model_probe.py b/backend/tests/test_ai_model_probe.py new file mode 100644 index 000000000..23f828259 --- /dev/null +++ b/backend/tests/test_ai_model_probe.py @@ -0,0 +1,109 @@ +"""AI 接口探测（拉取模型列表 / 连通性测试）回归测试。 + +不访问真实网络：把 ``httpx.AsyncClient`` 替换为假客户端，验证 URL/请求头/请求体， +以及错误分支的异常文案（不得回显密钥）。 +""" + +from typing import Any + +import httpx +import pytest + +from app.core.exceptions import CustomException +from app.modules.ai.chat import service as ai_service + + +class _FakeResponse: + def __init__(self, status_code: int = 200, payload: dict[str, Any] | None = None) -> None: + self.status_code = status_code + self._payload = payload or {} + + def json(self) -> dict[str, Any]: + return self._payload + + +class _FakeClient: + def __init__(self, response: _FakeResponse, error: Exception | None = None) -> None: + self._response = response + self._error = error + self.calls: list[tuple[str, str, dict[str, Any]]] = [] + + async def __aenter__(self) -> "_FakeClient": + return self + + async def __aexit__(self, *args: Any) -> bool: + return False + + def _record(self, method: str, url: str, headers: dict[str, str] | None, body: Any = None) -> _FakeResponse: + self.calls.append((method, url, {"headers": headers or {}, "json": body})) + if self._error is not None: + raise self._error + return self._response + + async def get(self, url: str, headers: dict[str, str] | None = None) -> _FakeResponse: + return self._record("GET", url, headers) + + async def post(self, url: str, headers: dict[str, str] | None = None, json: Any = None) -> _FakeResponse: + return self._record("POST", url, headers, json) + + +def _patch(monkeypatch: pytest.MonkeyPatch, response: _FakeResponse, error: Exception | None = None) -> dict[str, _FakeClient]: + holder: dict[str, _FakeClient] = {} + + def _factory(*args: Any, **kwargs: Any) -> _FakeClient: + client = _FakeClient(response, error) + holder["client"] = client + return client + + monkeypatch.setattr(ai_service.httpx, "AsyncClient", _factory) + return holder + + +async def test_probe_available_models_returns_ids(monkeypatch: pytest.MonkeyPatch) -> None: + holder = _patch( + monkeypatch, + _FakeResponse(200, {"data": [{"id": "gpt-4o", "owned_by": "openai"}, {"id": "deepseek-chat"}, {"no_id": 1}]}), + ) + models = await ai_service.probe_available_models("https://api.example.com/v1/", "sk-secret") + + assert models == [{"id": "gpt-4o", "owned_by": "openai"}, {"id": "deepseek-chat", "owned_by": None}] + method, url, payload = holder["client"].calls[0] + assert (method, url) == ("GET", "https://api.example.com/v1/models") + assert payload["headers"]["Authorization"] == "Bearer sk-secret" + + +async def test_probe_available_models_error_status_hides_key(monkeypatch: pytest.MonkeyPatch) -> None: + _patch(monkeypatch, _FakeResponse(401, {})) + with pytest.raises(CustomException) as exc_info: + await ai_service.probe_available_models("https://api.example.com/v1", "sk-secret") + assert "HTTP 401" in str(exc_info.value) + assert "sk-secret" not in str(exc_info.value) + + +async def test_probe_available_models_network_error(monkeypatch: pytest.MonkeyPatch) -> None: + _patch(monkeypatch, _FakeResponse(200, {}), error=httpx.ConnectError("boom")) + with pytest.raises(CustomException) as exc_info: + await ai_service.probe_available_models("https://api.example.com/v1", "sk-secret") + assert "ConnectError" in str(exc_info.value) + + +async def test_probe_model_connection_sends_minimal_payload(monkeypatch: pytest.MonkeyPatch) -> None: + holder = _patch(monkeypatch, _FakeResponse(200, {"choices": [{"message": {"content": "p"}}]})) + result = await ai_service.probe_model_connection("https://api.example.com/v1/", "sk-secret", "gpt-4o") + + assert result["ok"] is True + assert result["model_id"] == "gpt-4o" + assert isinstance(result["latency_ms"], int) + method, url, payload = holder["client"].calls[0] + assert (method, url) == ("POST", "https://api.example.com/v1/chat/completions") + assert payload["json"]["model"] == "gpt-4o" + assert payload["json"]["max_tokens"] == 1 + assert payload["headers"]["Authorization"] == "Bearer sk-secret" + + +async def test_probe_model_connection_error_status(monkeypatch: pytest.MonkeyPatch) -> None: + _patch(monkeypatch, _FakeResponse(404, {})) + with pytest.raises(CustomException) as exc_info: + await ai_service.probe_model_connection("https://api.example.com/v1", "sk-secret", "missing-model") + assert "HTTP 404" in str(exc_info.value) + assert "sk-secret" not in str(exc_info.value) diff --git a/backend/tests/test_host_guard.py b/backend/tests/test_host_guard.py new file mode 100644 index 000000000..c6110e242 --- /dev/null +++ b/backend/tests/test_host_guard.py @@ -0,0 +1,135 @@ +"""Host 头校验中间件回归测试。 + +重点固化两件事： +1. IPv6 字面量（``[::1]:8000`` / ``[::1]`` / ``::1``）不会被 Starlette 的 ``split(":")[0]`` 误判； +2. 其余语义与 Starlette 保持一致：端口忽略、``*.example.com`` 通配、``www.`` 301、非法主机 400。 +""" + +import httpx +import pytest +from starlette.types import ASGIApp, Receive, Scope, Send + +from app.core.middlewares import CustomTrustedHostMiddleware, host_of + + +async def _ok_app(scope: Scope, receive: Receive, send: Send) -> None: + await send({"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/plain")]}) + await send({"type": "http.response.body", "body": b"ok"}) + + +def _app(allowed: list[str], *, www_redirect: bool = True) -> ASGIApp: + return CustomTrustedHostMiddleware(_ok_app, allowed_hosts=allowed, www_redirect=www_redirect) + + +def _client(allowed: list[str], *, www_redirect: bool = True) -> httpx.AsyncClient: + return httpx.AsyncClient(transport=httpx.ASGITransport(app=_app(allowed, www_redirect=www_redirect)), + base_url="http://testserver") + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + ("127.0.0.1:8000", "127.0.0.1"), + ("[::1]:8000", "::1"), + ("[::1]", "::1"), + ("::1", "::1"), + ("Service.FastapiAdmin.com:443", "service.fastapiadmin.com"), + ("", ""), + ], +) +def test_host_of_normalizes_host(raw: str, expected: str) -> None: + assert host_of(raw) == expected + + +@pytest.mark.parametrize("host", ["[::1]:8000", "[::1]", "::1"]) +async def test_ipv6_hosts_are_allowed(host: str) -> None: + """Starlette 原实现把 ``[::1]:8000`` 切成 ``"["``，这里必须放行。""" + async with _client(["localhost", "127.0.0.1", "::1"]) as client: + response = await client.get("/", headers={"host": host}) + assert response.status_code == 200 + + +async def test_port_is_ignored() -> None: + async with _client(["service.fastapiadmin.com"]) as client: + response = await client.get("/", headers={"host": "service.fastapiadmin.com:8443"}) + assert response.status_code == 200 + + +async def test_bracketed_pattern_in_allowlist_is_normalized() -> None: + """允许清单里写成 ``[::1]``（settings.LOCALHOST_HOSTS 的形态）也应生效。""" + async with _client(["[::1]"]) as client: + response = await client.get("/", headers={"host": "[::1]:8000"}) + assert response.status_code == 200 + + +async def test_wildcard_subdomain_allowed_but_bare_domain_rejected() -> None: + async with _client(["*.fastapiadmin.com"]) as client: + assert (await client.get("/", headers={"host": "a.fastapiadmin.com"})).status_code == 200 + assert (await client.get("/", headers={"host": "fastapiadmin.com"})).status_code == 400 + + +async def test_host_comparison_is_case_insensitive() -> None: + async with _client(["service.fastapiadmin.com"]) as client: + response = await client.get("/", headers={"host": "SERVICE.FastapiAdmin.COM"}) + assert response.status_code == 200 + + +async def test_invalid_host_returns_400() -> None: + async with _client(["service.fastapiadmin.com"]) as client: + response = await client.get("/", headers={"host": "evil.example.com"}) + assert response.status_code == 400 + assert "Invalid host header" in response.text + + +async def test_www_redirect() -> None: + async with _client(["www.example.com"]) as client: + response = await client.get("/path?x=1", headers={"host": "example.com"}) + assert response.status_code == 301 + assert response.headers["location"] == "http://www.example.com/path?x=1" + + +async def test_www_redirect_can_be_disabled() -> None: + async with _client(["www.example.com"], www_redirect=False) as client: + response = await client.get("/", headers={"host": "example.com"}) + assert response.status_code == 400 + + +async def test_allow_any_wildcard() -> None: + async with _client(["*"]) as client: + response = await client.get("/", headers={"host": "anything.example"}) + assert response.status_code == 200 + + +async def test_non_http_scope_passes_through() -> None: + """lifespan 等非 http/ws scope 不应被拦截。""" + called: list[str] = [] + + async def _lifespan_app(scope: Scope, receive: Receive, send: Send) -> None: + called.append(scope["type"]) + + middleware = CustomTrustedHostMiddleware(_lifespan_app, allowed_hosts=["example.com"]) + await middleware({"type": "lifespan"}, _noop_receive, _noop_send) + assert called == ["lifespan"] + + +async def test_websocket_with_invalid_host_is_closed() -> None: + """WS scope 上不能发 HTTP 响应，应回 ``websocket.close``。""" + sent: list[dict[str, object]] = [] + + async def _ws_app(scope: Scope, receive: Receive, send: Send) -> None: + sent.append({"type": "app_called"}) + + async def _send(message: dict[str, object]) -> None: + sent.append(message) + + middleware = CustomTrustedHostMiddleware(_ws_app, allowed_hosts=["example.com"]) + await middleware({"type": "websocket", "headers": [(b"host", b"evil.example")]}, _noop_receive, _send) + assert sent == [{"type": "websocket.close", "code": 1008}] + + +async def _noop_receive() -> dict[str, object]: # pragma: no cover - 供 ASGI 调用签名使用 + return {"type": "http.request"} + + +async def _noop_send(message: dict[str, object]) -> None: # pragma: no cover - 供 ASGI 调用签名使用 + return None diff --git a/backend/tests/test_storage_config_io.py b/backend/tests/test_storage_config_io.py new file mode 100644 index 000000000..b2907d3d4 --- /dev/null +++ b/backend/tests/test_storage_config_io.py @@ -0,0 +1,153 @@ +"""存储源配置导出/导入回归测试。 + +不变量： +1. 导出**绝不含凭据**（password / advanced_config 都不出现）； +2. 导入按名称去重，已存在则跳过、绝不覆盖； +3. 协议非法 / 结构非法 / 端口缺失都有兜底，不会写坏库。 + +说明：这里用服务层 fake CRUD 验证，不走 HTTP 写路径——测试环境 SQLite 由同步引擎持锁， +API 写请求会等 ~60s 锁超时（测试基础设施问题，与业务逻辑无关）。 +""" + +import json +from types import SimpleNamespace +from typing import Any + +from app.modules.task.storage.node.config_io import EXPORT_FIELDS, build_export, plan_import +from app.modules.task.storage.node.service import StorageNodeService + +_BASE_NODE: dict[str, Any] = { + "name": "n1", + "protocol": "sftp", + "host": "127.0.0.1", + "port": 22, + "username": "u", + "bucket": None, + "endpoint": None, + "region": None, + "path_prefix": None, + "is_secure": False, + "implicit_tls": False, + "scheme": "https", + "encrypt_type": 1, + "connection_mode": 0, + "encoding": "UTF-8", + "multipart_part_size": 50, + "multipart_concurrency": 6, + "multipart_memory_budget": 512, + "status": 0, + "description": "", +} + + +def _node(**overrides: Any) -> SimpleNamespace: + data = dict(_BASE_NODE) + data.update(overrides) + # 凭据与高级配置：导出/导入都绝不能出现 + data["password"] = "super-secret" + data["advanced_config"] = {"proxy": "http://user:pass@proxy"} + return SimpleNamespace(**data) + + +class _FakeCrud: + """记录 create 调用的假 CRUD（get_list 返回既有节点）。""" + + def __init__(self, existing: list[SimpleNamespace]) -> None: + self.existing = existing + self.created: list[dict[str, Any]] = [] + + async def get_list(self, *args: Any, **kwargs: Any) -> list[SimpleNamespace]: + return self.existing + + async def create(self, data: dict[str, Any]) -> SimpleNamespace: + self.created.append(data) + return SimpleNamespace(id=len(self.created), **data) + + +def _service(crud: _FakeCrud) -> StorageNodeService: + service = StorageNodeService.__new__(StorageNodeService) + setattr(service, "_crud", lambda: crud) + return service + + +# ── 纯函数 ──────────────────────────────────────────────────────────── + + +def test_build_export_never_contains_credentials() -> None: + payload = build_export([_node()]) + node = payload["nodes"][0] + assert set(node) == set(EXPORT_FIELDS) + assert "password" not in node + assert "advanced_config" not in node + assert "super-secret" not in json.dumps(payload) + + +def test_plan_import_skips_existing_names() -> None: + data = {"nodes": [{"name": "a", "protocol": "sftp"}, {"name": "b", "protocol": "s3"}]} + to_create, skipped, invalid = plan_import(data, {"a"}) + assert [item["name"] for item in to_create] == ["b"] + assert (skipped, invalid) == (1, 0) + + +def test_plan_import_rejects_invalid_entries() -> None: + data = {"nodes": [{"name": "", "protocol": "sftp"}, {"name": "x", "protocol": "ftp2"}, "not-a-dict"]} + to_create, skipped, invalid = plan_import(data, set()) + assert to_create == [] + assert (skipped, invalid) == (0, 3) + + +def test_plan_import_defaults_port_by_protocol() -> None: + data = {"nodes": [{"name": "b", "protocol": "s3"}, {"name": "c", "protocol": "sftp", "port": "abc"}]} + to_create, _, _ = plan_import(data, set()) + assert to_create[0]["port"] == 443 + assert to_create[1]["port"] == 22 + + +def test_plan_import_never_carries_credentials() -> None: + """即便导入文件里夹带了 password/advanced_config，也不能写进待创建条目。""" + data = { + "nodes": [ + { + "name": "leaky", + "protocol": "sftp", + "password": "super-secret", + "advanced_config": {"proxy": "x"}, + } + ] + } + to_create, _, _ = plan_import(data, set()) + assert to_create and "password" not in to_create[0] + assert "advanced_config" not in to_create[0] + + +# ── 服务层 ──────────────────────────────────────────────────────────── + + +async def test_service_export_omits_credentials() -> None: + service = _service(_FakeCrud([_node()])) + payload = await service.export_config() + assert payload["version"] >= 1 + assert "exported_at" in payload + for node in payload["nodes"]: + assert "password" not in node + assert "advanced_config" not in node + + +async def test_service_import_dedupes_and_creates_without_credentials() -> None: + crud = _FakeCrud([_node(name="exists")]) + service = _service(crud) + payload = { + "nodes": [ + {"name": "exists", "protocol": "sftp"}, + {"name": "new-node", "protocol": "s3", "region": "us-east-1"}, + {"name": "", "protocol": "sftp"}, + ] + } + result = await service.import_config(payload) + assert result == {"created": 1, "skipped": 1, "invalid": 1} + assert len(crud.created) == 1 + created = crud.created[0] + assert created["name"] == "new-node" + assert created["protocol"] == "s3" + assert "password" not in created + assert "advanced_config" not in created diff --git a/backend/tests/test_storage_create_bucket.py b/backend/tests/test_storage_create_bucket.py new file mode 100644 index 000000000..1f8e839af --- /dev/null +++ b/backend/tests/test_storage_create_bucket.py @@ -0,0 +1,149 @@ +"""对象存储建桶（create_bucket）回归测试。 + +覆盖：服务层通用校验、协议守卫、适配器委托与资源释放；S3 的 LocationConstraint 规则； +OBS 以状态码而非异常报告失败的处理。 +""" + +from typing import Any + +import pytest + +from app.core.exceptions import CustomException +from app.modules.task.storage.browse.service import StorageFileService +from app.modules.task.storage.core.base import StorageAdapterConfig, StorageProtocol +from app.modules.task.storage.core.obs_adapter import ObsStorageAdapter +from app.modules.task.storage.core.s3_adapter import S3StorageAdapter + + +def _service() -> StorageFileService: + # 只测 create_bucket 的校验/委托逻辑：用 __new__ 跳过需要真实 auth/db 的构造 + return StorageFileService.__new__(StorageFileService) + + +def _config(protocol: StorageProtocol, **overrides: Any) -> StorageAdapterConfig: + data: dict[str, Any] = { + "protocol": protocol, + "port": 443, + "endpoint": "https://example.com", + "bucket": "demo", + "username": "ak", + "password": "sk", + } + data.update(overrides) + return StorageAdapterConfig(**data) + + +class _FakeAdapter: + """记录委托调用与关闭状态的假适配器。""" + + def __init__(self) -> None: + self.created: list[str] = [] + self.closed = False + + async def create_bucket(self, bucket: str) -> None: + self.created.append(bucket) + + async def close(self) -> None: + self.closed = True + + +@pytest.mark.parametrize("bad", ["", " ", "a/b", "a b", "x" * 64]) +async def test_create_bucket_rejects_blank_or_invalid_names(bad: str) -> None: + with pytest.raises(CustomException): + await _service().create_bucket(None, bad) + + +async def test_create_bucket_rejects_non_object_protocol(monkeypatch: pytest.MonkeyPatch) -> None: + service = _service() + + async def _get_source(_source_id: int | None) -> StorageAdapterConfig: + return _config(StorageProtocol.FTP, port=21, host="127.0.0.1") + + monkeypatch.setattr(service, "_get_source", _get_source) + with pytest.raises(CustomException, match="仅对象存储协议"): + await service.create_bucket(None, "bucket-1") + + +async def test_create_bucket_delegates_and_always_closes(monkeypatch: pytest.MonkeyPatch) -> None: + from app.modules.task.storage.core.factory import StorageAdapterFactory + + service = _service() + fake = _FakeAdapter() + + async def _get_source(_source_id: int | None) -> StorageAdapterConfig: + return _config(StorageProtocol.S3, region="ap-southeast-1") + + monkeypatch.setattr(service, "_get_source", _get_source) + monkeypatch.setattr(StorageAdapterFactory, "create", staticmethod(lambda _config: fake)) + + await service.create_bucket(None, " new-bucket ") + assert fake.created == ["new-bucket"], "桶名应先 strip 再委托给适配器" + assert fake.closed is True, "无论成功失败都必须关闭适配器" + + +def test_base_adapter_default_create_bucket_is_unsupported() -> None: + from app.modules.task.storage.core.local_adapter import LocalStorageAdapter + + adapter = LocalStorageAdapter(StorageAdapterConfig(protocol=StorageProtocol.LOCAL, port=0, host="/tmp")) + with pytest.raises(NotImplementedError): + adapter._sync_create_bucket("x") + + +def test_s3_create_bucket_passes_location_constraint() -> None: + adapter = S3StorageAdapter(_config(StorageProtocol.S3, region="ap-southeast-1")) + calls: dict[str, Any] = {} + + class _Client: + def create_bucket(self, **kwargs: Any) -> None: + calls.update(kwargs) + + adapter.client = _Client() # type: ignore[assignment] + adapter._sync_create_bucket("new-bucket") + assert calls == { + "Bucket": "new-bucket", + "CreateBucketConfiguration": {"LocationConstraint": "ap-southeast-1"}, + } + + +def test_s3_create_bucket_omits_location_constraint_for_us_east_1() -> None: + adapter = S3StorageAdapter(_config(StorageProtocol.S3, region="us-east-1")) + calls: dict[str, Any] = {} + + class _Client: + def create_bucket(self, **kwargs: Any) -> None: + calls.update(kwargs) + + adapter.client = _Client() # type: ignore[assignment] + adapter._sync_create_bucket("new-bucket") + assert calls == {"Bucket": "new-bucket"} + + +def test_obs_create_bucket_raises_on_error_status() -> None: + adapter = ObsStorageAdapter(_config(StorageProtocol.OBS, endpoint="obs.example.com", region="cn-north-4")) + + class _Resp: + status = 403 + errorMessage = "AccessDenied" + + class _Client: + def createBucket(self, bucketName: str) -> _Resp: # noqa: N802 - 对齐 OBS SDK 命名 + return _Resp() + + adapter.client = _Client() # type: ignore[assignment] + with pytest.raises(CustomException, match="OBS 创建桶失败"): + adapter._sync_create_bucket("new-bucket") + + +def test_obs_create_bucket_succeeds_on_2xx() -> None: + adapter = ObsStorageAdapter(_config(StorageProtocol.OBS, endpoint="obs.example.com", region="cn-north-4")) + + class _Resp: + status = 200 + errorMessage = "" + + class _Client: + def createBucket(self, bucketName: str) -> _Resp: # noqa: N802 - 对齐 OBS SDK 命名 + return _Resp() + + adapter.client = _Client() # type: ignore[assignment] + adapter._sync_create_bucket("new-bucket") # 不抛异常即通过 diff --git a/backend/tests/test_storage_lazy_sdk.py b/backend/tests/test_storage_lazy_sdk.py new file mode 100644 index 000000000..069f7744a --- /dev/null +++ b/backend/tests/test_storage_lazy_sdk.py @@ -0,0 +1,76 @@ +"""存储适配器 SDK 懒加载回归测试（参照 FastHub 的 lazy_sdk 实践）。 + +不变量： +1. 导入适配器工厂（会引入全部 7 个适配器）**不得**拉起任何重型 SDK； +2. 只有构造对应协议的客户端时才导入该 SDK。 + +用子进程验证 `sys.modules`：本进程 import 过的东西无法"卸载"，只能隔离观察。 +""" + +import json +import subprocess +import sys +from pathlib import Path + +_BACKEND_DIR = Path(__file__).resolve().parents[1] +_HEAVY = ("boto3", "botocore", "qcloud_cos", "obs", "paramiko", "alibabacloud_oss_v2") + +_PROBE = r''' +import json, sys +import app.modules.task.storage.core.factory # 触发全部适配器导入 +HEAVY = ("boto3", "botocore", "qcloud_cos", "obs", "paramiko", "alibabacloud_oss_v2") +print("__BEFORE__" + json.dumps([m for m in HEAVY if m in sys.modules])) + +from app.modules.task.storage.core.base import StorageAdapterConfig, StorageProtocol +from app.modules.task.storage.core.s3_adapter import S3StorageAdapter +try: + S3StorageAdapter( + StorageAdapterConfig( + protocol=StorageProtocol.S3, + port=443, + endpoint="https://s3.example.com", + bucket="demo", + username="access-key", + password="secret-key", + ) + ) +except Exception as exc: # 构造期不需要联网，允许因参数不足抛错 + print("__S3ERR__" + type(exc).__name__ + ": " + str(exc)[:200]) + +print("__AFTER__" + json.dumps([m for m in HEAVY if m in sys.modules])) +''' + + +def _probe() -> dict[str, list[str]]: + """在独立进程中导入工厂并构造 S3 适配器，返回前后已载入的重型 SDK。""" + proc = subprocess.run( + [sys.executable, "-c", _PROBE], + cwd=_BACKEND_DIR, + capture_output=True, + text=True, + timeout=180, + ) + assert proc.returncode == 0, f"子进程失败: {proc.stderr[-2000:]}" + data: dict[str, list[str]] = {} + for line in proc.stdout.splitlines(): + if line.startswith("__BEFORE__"): + data["before"] = json.loads(line[len("__BEFORE__") :]) + elif line.startswith("__AFTER__"): + data["after"] = json.loads(line[len("__AFTER__") :]) + assert "before" in data and "after" in data, f"探针输出异常: {proc.stdout[-1000:]}" + return data + + +def test_importing_factory_does_not_load_heavy_sdks() -> None: + """导入工厂不能拉起任何重型 SDK（冷启动不为用不到的协议付代价）。""" + data = _probe() + assert data["before"] == [], f"导入工厂即拉起重型 SDK: {data['before']}" + + +def test_constructing_s3_adapter_loads_only_its_own_sdk() -> None: + """构造 S3 客户端只应载入 boto3/botocore，不得连带载入其它协议 SDK。""" + data = _probe() + after = data["after"] + assert "boto3" in after, f"S3 构造后未导入 boto3: {after}" + for module in ("qcloud_cos", "obs", "alibabacloud_oss_v2", "paramiko"): + assert module not in after, f"S3 构造连带导入了 {module}: {after}" diff --git a/frontend/web/src/api/module_ai/chat.ts b/frontend/web/src/api/module_ai/chat.ts index 66c2ee387..56628bc5b 100644 --- a/frontend/web/src/api/module_ai/chat.ts +++ b/frontend/web/src/api/module_ai/chat.ts @@ -94,8 +94,37 @@ export const AiChatAPI = { method: "post", }); }, + + /** 拉取接口可用模型列表（按填写中的 base_url + api_key 探测，不落库） */ + listAvailableModels(body: { base_url: string; api_key: string }) { + return request>({ + url: `${API_PATH}/model/available`, + method: "post", + data: body, + }); + }, + + /** 测试接口连通性（按填写中的参数探测，不落库） */ + testModelConnection(body: { base_url: string; api_key: string; model_id: string }) { + return request>({ + url: `${API_PATH}/model/test`, + method: "post", + data: body, + }); + }, }; +export interface AiModelOption { + id: string; + owned_by?: string | null; +} + +export interface AiModelTestResult { + ok: boolean; + model_id: string; + latency_ms: number; +} + export interface AiModelConfigInput { name: string; base_url: string; diff --git a/frontend/web/src/api/module_storage/browse.ts b/frontend/web/src/api/module_storage/browse.ts index dc9005c4d..f4e919d4d 100644 --- a/frontend/web/src/api/module_storage/browse.ts +++ b/frontend/web/src/api/module_storage/browse.ts @@ -21,6 +21,15 @@ const StorageAPI = { }); }, + /** 创建存储桶（对象存储；COS 的桶名需带 -appid 后缀） */ + createBucket(body: { source_id?: number | null; bucket: string }) { + return request>({ + url: `${API_PATH}/create-bucket`, + method: "post", + data: body, + }); + }, + /** 上传文件（可传 onProgress 回调获取 0-100 进度百分比） */ uploadFile(formData: FormData, onProgress?: (percent: number) => void) { return request>>({ diff --git a/frontend/web/src/api/module_storage/node.ts b/frontend/web/src/api/module_storage/node.ts index 39a7b01f4..46b3aa97c 100644 --- a/frontend/web/src/api/module_storage/node.ts +++ b/frontend/web/src/api/module_storage/node.ts @@ -80,8 +80,31 @@ const NodeAPI = { method: "get", }); }, + + /** 导出存储源配置（不含凭据） */ + exportConfig() { + return request>({ + url: `${API_PATH}/export`, + method: "get", + }); + }, + + /** 导入存储源配置（不含凭据，按名称去重，已存在则跳过） */ + importConfig(data: StorageConfigExport) { + return request>({ + url: `${API_PATH}/import`, + method: "post", + data, + }); + }, }; +export interface StorageConfigExport { + version: number; + exported_at: string; + nodes: Record[]; +} + export default NodeAPI; /** 存储协议项 */ diff --git a/frontend/web/src/views/module_ai/chat/components/FaAiModelConfigPanel.vue b/frontend/web/src/views/module_ai/chat/components/FaAiModelConfigPanel.vue index 298a26d17..0f92b8ef4 100644 --- a/frontend/web/src/views/module_ai/chat/components/FaAiModelConfigPanel.vue +++ b/frontend/web/src/views/module_ai/chat/components/FaAiModelConfigPanel.vue @@ -222,11 +222,29 @@ /> - +
+
+ + 拉取模型 + 测试连通 +
+
+ + {{ m.id }} + +
+
({ created_time: null, }); +// 模型探测：拉取可用模型 / 连通性测试（按表单中填写、尚未保存的参数） +const probingModels = ref(false); +const testingConnection = ref(false); +const availableModels = ref([]); + +async function fetchAvailableModels() { + if (!form.base_url || !form.api_key) { + ElMessage.warning("请先填写 Base URL 与 API Key"); + return; + } + probingModels.value = true; + try { + const res = await AiChatAPI.listAvailableModels({ + base_url: form.base_url, + api_key: form.api_key, + }); + availableModels.value = res.data.data || []; + if (availableModels.value.length) { + ElMessage.success(`获取到 ${availableModels.value.length} 个可用模型，点击标签即可填入`); + } else { + ElMessage.warning("接口未返回可用模型"); + } + } catch { + // 统一错误提示已由 http 层处理 + } finally { + probingModels.value = false; + } +} + +async function testConnection() { + if (!form.base_url || !form.api_key || !form.model_id) { + ElMessage.warning("请先填写 Base URL、API Key 与模型 ID"); + return; + } + testingConnection.value = true; + try { + const res = await AiChatAPI.testModelConnection({ + base_url: form.base_url, + api_key: form.api_key, + model_id: form.model_id, + }); + ElMessage.success(`连接成功，往返 ${res.data.data?.latency_ms ?? "-"} ms`); + } catch { + // 统一错误提示已由 http 层处理 + } finally { + testingConnection.value = false; + } +} + // 常用 Base URL 预设 - 提升用户输入效率 const baseUrlPresets = [ { label: "OpenAI 官方", url: "https://api.openai.com/v1" }, diff --git a/frontend/web/src/views/module_task/storage/browse/index.vue b/frontend/web/src/views/module_task/storage/browse/index.vue index 3383e03dd..e5e25339b 100644 --- a/frontend/web/src/views/module_task/storage/browse/index.vue +++ b/frontend/web/src/views/module_task/storage/browse/index.vue @@ -210,7 +210,17 @@ v-loading="activeTab.loading" class="min-h-0 flex-1" > -
存储桶（账号下全部，点击进入）
+
+ 存储桶（账号下全部，点击进入） + 新建桶 +
({ value: undefined as string | undefined })); + const name = value?.trim(); + if (!name) return; + bucketCreating.value = true; + try { + await StorageAPI.createBucket({ source_id: tab.nodeId, bucket: name }); + ElMessage.success("创建桶成功"); + await loadTab(tab); + } finally { + bucketCreating.value = false; + } +} + // ─── 文件操作：下载 ──────────────────────────────────────────────── function filenameFromDisposition(header: string | undefined, fallback: string): string { if (!header) return fallback; diff --git a/frontend/web/src/views/module_task/storage/node/index.vue b/frontend/web/src/views/module_task/storage/node/index.vue index 753959963..784ec647d 100644 --- a/frontend/web/src/views/module_task/storage/node/index.vue +++ b/frontend/web/src/views/module_task/storage/node/index.vue @@ -25,15 +25,26 @@ @refresh="refreshData" > @@ -190,6 +201,7 @@ import { useRouter } from "vue-router"; import { ElTag } from "element-plus"; import { useCrudForm } from "@/hooks/core/useCrudForm"; import { resolveStatusColumns, renderTableOperationCell, type TableOperationAction } from "@utils"; +import { saveAs } from "file-saver"; import NodeAPI, { type AdvancedFieldDef, type SourceForm, @@ -477,6 +489,65 @@ async function handleTestConfig() { // ─── 对话框状态 ─── const { dialogVisible } = useCrudDialog(); +// ─── 配置导出 / 导入（不含凭据；导入按名称去重）──────────────────── +const exporting = ref(false); +const importing = ref(false); +const importInputRef = ref(); + +async function handleExportConfig() { + exporting.value = true; + try { + const res = await NodeAPI.exportConfig(); + const payload = res.data.data; + saveAs( + new Blob([JSON.stringify(payload, null, 2)], { type: "application/json" }), + `storage-nodes-${new Date().toISOString().slice(0, 19).replace(/[:T]/g, "-")}.json` + ); + ElMessage.success("已导出（不含凭据，导入后需重填密码）"); + } finally { + exporting.value = false; + } +} + +function triggerImportConfig() { + importInputRef.value?.click(); +} + +async function handleImportFile(event: Event) { + const input = event.target as HTMLInputElement; + const file = input.files?.[0]; + input.value = ""; + if (!file) return; + let payload: { version?: number; nodes?: Record[] }; + try { + payload = JSON.parse(await file.text()); + } catch { + ElMessage.error("文件不是合法的 JSON"); + return; + } + if (!Array.isArray(payload.nodes)) { + ElMessage.error("文件内容缺少 nodes 数组"); + return; + } + importing.value = true; + try { + const res = await NodeAPI.importConfig({ + version: payload.version ?? 1, + exported_at: "", + nodes: payload.nodes, + }); + const result = res.data.data; + ElMessage.success( + `导入完成：新增 ${result?.created ?? 0}，跳过 ${result?.skipped ?? 0}（同名已存在），无效 ${result?.invalid ?? 0}` + ); + refreshData(); + } catch { + // 统一错误提示由 http 层处理 + } finally { + importing.value = false; + } +} + const detailFormData = ref({} as SourceTable); const sourceDetailItems: import("@/components/display/fa-descriptions/index.vue").DescriptionsItem[] = @@ -691,7 +762,10 @@ const sourceDialogFormItems = computed(() => { props: { placeholder: "请选择或输入区域", // ElOption 只认 { label, value }：直接传字符串会展开成下标属性、value 缺省被布尔转为 false，导致选项显示 "false" - options: (regionDict[p]?.options ?? []).map((region) => ({ label: region, value: region })), + options: (regionDict[p]?.options ?? []).map((region) => ({ + label: region, + value: region, + })), filterable: true, allowCreate: true, },