# agent_fastapi.py from __future__ import annotations import asyncio import mimetypes import os import sys import json import re import time import uuid import math import logging import shutil from pathlib import Path from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Tuple, Set from contextlib import asynccontextmanager from starlette.websockets import WebSocketState, WebSocketDisconnect try: import tomllib # Python 3.11+ # type: ignore except ModuleNotFoundError: import tomli as tomllib # Python <= 3.10 import traceback try: from uvicorn.protocols.utils import ClientDisconnected except Exception: ClientDisconnected = None logger = logging.getLogger(__name__) import anyio from fastapi import FastAPI, APIRouter, UploadFile, File, Form, HTTPException, WebSocket, WebSocketDisconnect, Request from fastapi.responses import FileResponse, JSONResponse, Response from fastapi.staticfiles import StaticFiles from langchain_core.messages import SystemMessage, HumanMessage, BaseMessage, AIMessage, ToolMessage # ---- 确保 src 可导入（避免环境差异导致找不到模块）---- ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) SRC_DIR = os.path.join(ROOT_DIR, "src") if SRC_DIR not in sys.path: sys.path.insert(0, SRC_DIR) from open_storyline.agent import build_agent, ClientContext from open_storyline.utils.prompts import get_prompt from open_storyline.utils.media_handler import scan_media_dir from open_storyline.config import load_settings, default_config_path from open_storyline.config import Settings from open_storyline.storage.agent_memory import ArtifactStore from open_storyline.mcp.hooks.node_interceptors import ToolInterceptor from open_storyline.mcp.hooks.chat_middleware import set_mcp_log_sink, reset_mcp_log_sink WEB_DIR = os.path.join(ROOT_DIR, "web") STATIC_DIR = os.path.join(WEB_DIR, "static") INDEX_HTML = os.path.join(WEB_DIR, "index.html") NODE_MAP_HTML = os.path.join(WEB_DIR, "node_map/node_map.html") NODE_MAP_DIR = os.path.join(WEB_DIR, "node_map") SERVER_CACHE_DIR = os.path.join(ROOT_DIR, '.storyline' , ".server_cache") CHUNK_SIZE = 1024 * 1024 # 1MB # 是否根据session_id隔离用户 USE_SESSION_SUBDIR = True CUSTOM_MODEL_KEY = "__custom__" # Load keys DEFAULT_LLM_API_KEY = os.getenv("DEEPSEEK_API_KEY") DEFAULT_LLM_API_URL = os.getenv("DEEPSEEK_API_URL") DEFAULT_LLM_API_NAME = os.getenv("DEEPSEEK_API_NAME", "deepseek-chat") DEFAULT_VLM_API_KEY = os.getenv("GLM_V4_6_API_KEY") DEFAULT_VLM_API_URL = os.getenv("GLM_V4_6_API_URL") DEFAULT_VLM_API_NAME = os.getenv("GLM_V4_6_API_NAME", "qwen3-vl-8b-instruct") print("DEEPSEEK_API_KEY exists:", bool(os.getenv("DEEPSEEK_API_KEY"))) print("QWEN3_VL_8B_API_KEY exists:", bool(os.getenv("QWEN3_VL_8B_API_KEY"))) print("DEEPSEEK_API_URL:", repr(os.getenv("DEEPSEEK_API_URL"))) print("QWEN3_VL_8B_API_URL:", repr(os.getenv("QWEN3_VL_8B_API_URL"))) def debug_traceback_print(cfg: Settings): if cfg.developer.developer_mode: traceback.print_exc() def _s(x: Any) -> str: return str(x or "").strip() def _norm_url(u: Any) -> str: u = _s(u) return u.rstrip("/") if u else "" def _env_fallback_for_model(model_name: str) -> Tuple[str, str]: """ - deepseek* -> DEEPSEEK_API_URL / DEEPSEEK_API_KEY - qwen3* -> QWEN3_VL_8B_API_URL / QWEN3_VL_8B_API_KEY """ m = _s(model_name).lower() if "deepseek" in m: return (_s(os.getenv("DEEPSEEK_API_URL")), _s(os.getenv("DEEPSEEK_API_KEY"))) if m.startswith("qwen3-vl-8b-instruct") or "qwen3-vl-8b-instruct" in m: return (_s(os.getenv("QWEN3_VL_8B_API_URL")), _s(os.getenv("QWEN3_VL_8B_API_KEY"))) return ("", "") def _resolve_default_model_override(cfg: Settings, model_name: str) -> Tuple[Optional[Dict[str, Any]], Optional[str]]: """ 1. get config from [developer.chat_models_config.""] 2. rollback to env """ model_name = _s(model_name) if not model_name: return None, "default model name is empty" model_cfg: Dict[str, Any] = {} try: model_cfg = (cfg.developer.chat_models_config.get(model_name) or {}) if getattr(cfg, "developer", None) else {} except Exception: model_cfg = {} if not isinstance(model_cfg, dict): model_cfg = {} base_url = _norm_url(model_cfg.get("base_url")) api_key = _s(model_cfg.get("api_key")) if not base_url or not api_key: env_url, env_key = _env_fallback_for_model(model_name) if not base_url: base_url = _norm_url(env_url) if not api_key: api_key = _s(env_key) override: Dict[str, Any] = {"model": model_name} if base_url: override["base_url"] = base_url if api_key: override["api_key"] = api_key for k in ("timeout", "temperature", "max_retries", "top_p", "max_tokens"): if k in model_cfg and model_cfg.get(k) not in (None, ""): override[k] = model_cfg.get(k) if not override.get("base_url") or not override.get("api_key"): return None, ( f"cannot find base_url/api_key of default model: {model_name}. " f"please fill in base_url/api_key of [developer.chat_models_config.\"{model_name}\" in config.toml]" f"or set environment variables（DEEPSEEK_API_URL/DEEPSEEK_API_KEY / QWEN3_VL_8B_API_URL/QWEN3_VL_8B_API_KEY）。" ) return override, None def _stable_dict_key(d: Optional[Dict[str, Any]]) -> str: try: return json.dumps(d or {}, sort_keys=True, ensure_ascii=False) except Exception: return str(d or {}) def _parse_service_config(service_cfg: Any) -> Tuple[ Optional[Dict[str, Any]], Optional[Dict[str, Any]], Dict[str, Any], Dict[str, Any], Optional[str]]: """ 返回 (custom_llm, custom_vlm, tts_cfg, pexels, err) - custom_llm/custom_vlm: {"model","base_url","api_key"} 或 None（允许只传 llm 或只传 vlm） - tts_cfg: dict（可能为空） """ if not isinstance(service_cfg, dict): return None, None, {}, {}, None # ---- custom models ---- custom_llm = None custom_vlm = None custom_models = service_cfg.get("custom_models") if custom_models is not None: if not isinstance(custom_models, dict): return None, None, {}, {}, "service_config.custom_models 必须是对象" def _pick(m: Any, label: str) -> Tuple[Optional[Dict[str, str]], Optional[str]]: if m is None: return None, None if not isinstance(m, dict): return None, f"service_config.custom_models.{label} 必须是对象" model = _s(m.get("model")) base_url = _norm_url(m.get("base_url")) api_key = _s(m.get("api_key")) if not (model and base_url and api_key): return None, f"自定义 {label.upper()} 配置不完整：请填写 model/base_url/api_key" if not (base_url.startswith("http://") or base_url.startswith("https://")): return None, f"自定义 {label.upper()} 的 base_url 必须以 http(s) 开头" return {"model": model, "base_url": base_url, "api_key": api_key}, None custom_llm, err1 = _pick(custom_models.get("llm"), "llm") if err1: return None, None, {}, {}, err1 custom_vlm, err2 = _pick(custom_models.get("vlm"), "vlm") if err2: return None, None, {}, {}, err2 # ---- tts ---- tts_cfg: Dict[str, Any] = {} tts = service_cfg.get("tts") if isinstance(tts, dict): provider = (tts.get("provider") or "").strip().lower() if provider: provider_block = tts.get(provider) tts_cfg = {"provider": provider, provider: provider_block} # ---- pexels ---- pexels_cfg: Dict[str, Any] = {} search_media = service_cfg.get("search_media") if isinstance(search_media, dict): # 支持两种格式： # 1) {search_media:{pexels:{mode, api_key}}} # 2) {search_media:{mode, pexel_api_key}} p = search_media.get("pexels") or search_media.get("pexels") if isinstance(p, dict): mode = _s(p.get("mode")).lower() if mode not in ("default", "custom"): mode = "default" api_key = _s(p.get("api_key") or p.get("pexels_api_key") or p.get("pexels_api_key")) pexels_cfg = {"mode": mode, "api_key": api_key} else: mode = _s(search_media.get("mode") or search_media.get("pexels_mode") or search_media.get("pexels_mode")).lower() if mode not in ("default", "custom"): mode = "default" api_key = _s(search_media.get("pexels_api_key") or search_media.get("pexels_api_key")) pexels_cfg = {"mode": mode, "api_key": api_key} return custom_llm, custom_vlm, tts_cfg, pexels_cfg, None def is_developer_mode(cfg: Settings) -> bool: try: return bool(cfg.developer.developer_mode) except Exception: return False def _abs(p: str) -> str: return os.path.abspath(os.path.expanduser(p)) def resolve_media_dir(cfg_media_dir: str, session_id: str) -> str: root = _abs(cfg_media_dir).rstrip("/\\") if not USE_SESSION_SUBDIR: return root project_dir = os.path.dirname(root) leaf = os.path.basename(root) return os.path.join(project_dir, session_id, leaf) def sanitize_filename(name: str) -> str: name = os.path.basename(name or "") name = name.replace("\x00", "") return name or "unnamed" def detect_media_kind(filename: str) -> str: ext = os.path.splitext(filename)[1].lower() if ext in {".png", ".jpg", ".jpeg", ".gif", ".bmp", ".webp"}: return "image" if ext in {".mp4", ".mov", ".avi", ".mkv", ".webm"}: return "video" return "unknown" _MEDIA_RE = re.compile(r"^media_(\d+)", re.IGNORECASE) def make_media_store_filename(seq: int, ext: str) -> str: ext = (ext or "").lower() if ext and not ext.startswith("."): ext = "." + ext return f"{MEDIA_PREFIX}{seq:0{MEDIA_SEQ_WIDTH}d}{ext}" def parse_media_seq(filename: str) -> Optional[int]: m = _MEDIA_RE.match(os.path.basename(filename or "")) if not m: return None try: return int(m.group(1)) except Exception: return None def safe_save_path_no_overwrite(media_dir: str, filename: str) -> str: filename = sanitize_filename(filename) stem, ext = os.path.splitext(filename) path = os.path.join(media_dir, filename) if not os.path.exists(path): return path i = 2 while True: p2 = os.path.join(media_dir, f"{stem} ({i}){ext}") if not os.path.exists(p2): return p2 i += 1 def ensure_thumbs_dir(media_dir: str) -> str: d = os.path.join(media_dir, ".thumbs") os.makedirs(d, exist_ok=True) return d def ensure_uploads_dir(media_dir: str) -> str: d = os.path.join(media_dir, ".uploads") os.makedirs(d, exist_ok=True) return d def guess_media_type(path: str) -> str: mt, _ = mimetypes.guess_type(path) return mt or "application/octet-stream" def _is_under_dir(path: str, root: str) -> bool: try: path = os.path.abspath(path) root = os.path.abspath(root) return os.path.commonpath([path, root]) == root except Exception: return False def video_placeholder_svg_bytes() -> bytes: svg = """ """ return svg.encode("utf-8") def make_image_thumbnail_sync(src_path: str, dst_path: str, max_size: Tuple[int, int] = (320, 320)) -> bool: try: from PIL import Image img = Image.open(src_path).convert("RGB") img.thumbnail(max_size) img.save(dst_path, format="JPEG", quality=85) return True except Exception: return False async def make_video_thumbnail_async( src_video: str, dst_path: str, *, max_size: Tuple[int, int] = (320, 320), seek_sec: float = 0.5, timeout_sec: float = 20.0, ) -> bool: ffmpeg = os.environ.get("FFMPEG_BIN") or shutil.which("ffmpeg") if not ffmpeg: logger.warning("ffmpeg not found (PATH/FFMPEG_BIN). skip video thumbnail. src=%s", src_video) return False src_video = os.path.abspath(src_video) dst_path = os.path.abspath(dst_path) os.makedirs(os.path.dirname(dst_path), exist_ok=True) tmp_path = dst_path + ".tmp.jpg" vf = ( f"scale={max_size[0]}:{max_size[1]}:force_original_aspect_ratio=decrease" f",pad={max_size[0]}:{max_size[1]}:(ow-iw)/2:(oh-ih)/2" ) async def _run(args: list[str]) -> tuple[bool, str]: proc = await asyncio.create_subprocess_exec( *args, stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.PIPE, ) try: _, err = await asyncio.wait_for(proc.communicate(), timeout=timeout_sec) except asyncio.TimeoutError: try: proc.kill() except Exception: pass await proc.wait() return False, f"timeout after {timeout_sec}s" err_text = (err or b"").decode("utf-8", "ignore").strip() return (proc.returncode == 0), err_text # 两种策略：1) -ss 在 -i 前（快，但有些文件/关键帧会失败） # 2) -ss 在 -i 后（慢，但更稳定） common_tail = [ "-an", "-frames:v", "1", "-vf", vf, "-vcodec", "mjpeg", "-q:v", "3", "-f", "image2", tmp_path, ] attempts = [ # fast seek [ffmpeg, "-hide_banner", "-loglevel", "error", "-y", "-ss", f"{seek_sec}", "-i", src_video] + common_tail, # accurate seek [ffmpeg, "-hide_banner", "-loglevel", "error", "-y", "-i", src_video, "-ss", f"{seek_sec}"] + common_tail, # fallback：如果 seek 太靠前导致失败，再试试 1s [ffmpeg, "-hide_banner", "-loglevel", "error", "-y", "-ss", "1.0", "-i", src_video] + common_tail, ] last_err: Optional[str] = None try: for args in attempts: ok, err = await _run(args) if ok and os.path.exists(tmp_path) and os.path.getsize(tmp_path) > 0: os.replace(tmp_path, dst_path) return True last_err = err or last_err # 清理无效临时文件，避免下次误判 try: if os.path.exists(tmp_path): os.remove(tmp_path) except Exception: pass logger.warning("ffmpeg thumbnail failed. src=%s dst=%s err=%s", src_video, dst_path, last_err) return False finally: try: if os.path.exists(tmp_path): os.remove(tmp_path) except Exception: pass def _env_int(name: str, default: int) -> int: try: return int(os.environ.get(name, str(default))) except Exception: return default def _env_float(name: str, default: float) -> float: try: return float(os.environ.get(name, str(default))) except Exception: return float(default) def _rpm_to_rps(rpm: float) -> float: return float(rpm) / 60.0 # 是否信任反向代理头（X-Forwarded-For / X-Real-IP） RATE_LIMIT_TRUST_PROXY_HEADERS = os.environ.get("RATE_LIMIT_TRUST_PROXY_HEADERS", "0") == "1" @dataclass class _RateBucket: tokens: float last_ts: float # monotonic last_seen: float # monotonic (for TTL cleanup) class TokenBucketRateLimiter: """ 内存令牌桶 + 防爆内存： - max_buckets: 限制内部桶表最大条目数（防止海量 IP 导致字典膨胀） - evict_batch: 超过上限后每次驱逐多少条（按插入顺序驱逐最早创建的桶） """ def __init__( self, ttl_sec: int = 900, cleanup_interval_sec: int = 60, *, max_buckets: int = 100000, evict_batch: int = 2000, ): self.ttl_sec = int(ttl_sec) self.cleanup_interval_sec = int(cleanup_interval_sec) self.max_buckets = int(max(1, max_buckets)) self.evict_batch = int(max(1, evict_batch)) self._buckets: Dict[str, _RateBucket] = {} self._lock = asyncio.Lock() self._last_cleanup = time.monotonic() async def allow( self, key: str, *, capacity: float, refill_rate: float, cost: float = 1.0, ) -> Tuple[bool, float, float]: """ 返回: (allowed, retry_after_sec, remaining_tokens) """ now = time.monotonic() capacity = float(max(0.0, capacity)) refill_rate = float(max(0.0, refill_rate)) cost = float(max(0.0, cost)) async with self._lock: b = self._buckets.get(key) if b is None: # 先做一次周期清理 if now - self._last_cleanup > self.cleanup_interval_sec: self._cleanup_locked(now) self._last_cleanup = now # 桶表满了：先清 TTL，再做批量驱逐；仍然满 -> 不再创建新桶，直接拒绝 if len(self._buckets) >= self.max_buckets: self._cleanup_locked(now) if len(self._buckets) >= self.max_buckets: self._evict_locked() if len(self._buckets) >= self.max_buckets: # 不存任何新 key，避免内存继续涨 # retry_after 给一个很短的值即可（客户端会重试） return False, 1.0, 0.0 b = _RateBucket(tokens=capacity, last_ts=now, last_seen=now) self._buckets[key] = b else: b.last_seen = now # refill elapsed = max(0.0, now - b.last_ts) if refill_rate > 0: b.tokens = min(capacity, b.tokens + elapsed * refill_rate) else: b.tokens = min(capacity, b.tokens) b.last_ts = now if b.tokens >= cost: b.tokens -= cost return True, 0.0, float(max(0.0, b.tokens)) # not enough if refill_rate <= 0: retry_after = float(self.ttl_sec) else: need = cost - b.tokens retry_after = need / refill_rate return False, float(retry_after), float(max(0.0, b.tokens)) def _cleanup_locked(self, now: float) -> None: ttl = float(self.ttl_sec) dead = [k for k, b in self._buckets.items() if (now - b.last_seen) > ttl] for k in dead: self._buckets.pop(k, None) def _evict_locked(self) -> None: # 按 dict 插入顺序驱逐最早的一批 bucket（不排序，避免在高压下额外 CPU 开销） n = min(self.evict_batch, len(self._buckets)) for _ in range(n): try: k = next(iter(self._buckets)) except StopIteration: break self._buckets.pop(k, None) def _headers_to_dict(scope_headers: List[Tuple[bytes, bytes]]) -> Dict[str, str]: d: Dict[str, str] = {} for k, v in scope_headers or []: try: dk = k.decode("latin1").lower() dv = v.decode("latin1") except Exception: continue d[dk] = dv return d def _client_ip_from_http_scope(scope: dict, trust_proxy_headers: bool) -> str: headers = _headers_to_dict(scope.get("headers") or []) if trust_proxy_headers: xff = headers.get("x-forwarded-for") if xff: # "client, proxy1, proxy2" -> client return xff.split(",")[0].strip() or "unknown" xri = headers.get("x-real-ip") if xri: return xri.strip() or "unknown" client = scope.get("client") if client and isinstance(client, (list, tuple)) and len(client) >= 1: return str(client[0] or "unknown") return "unknown" def _client_ip_from_ws(ws: WebSocket, trust_proxy_headers: bool) -> str: try: if trust_proxy_headers: xff = ws.headers.get("x-forwarded-for") if xff: return xff.split(",")[0].strip() or "unknown" xri = ws.headers.get("x-real-ip") if xri: return xri.strip() or "unknown" except Exception: pass try: if ws.client: return str(ws.client.host or "unknown") except Exception: pass return "unknown" # 分片上传（绕开网关对单次请求体/单文件的限制） UPLOAD_RESUMABLE_CHUNK_BYTES = _env_int("UPLOAD_RESUMABLE_CHUNK_BYTES", 8 * 1024 * 1024) # 未完成的分片上传状态保留多久（超时自动清理临时文件） RESUMABLE_UPLOAD_TTL_SEC = _env_int("RESUMABLE_UPLOAD_TTL_SEC", 3600) # 1 hour MEDIA_SEQ_WIDTH = 4 # media_0001 MEDIA_PREFIX = "media_" # -------- 注意：在服务器上，所有用户的ip可能是相同的---- # 每个 IP 的总体请求速率（包括 /static、/api、/ 等） HTTP_GLOBAL_RPM = _env_int("RATE_LIMIT_HTTP_GLOBAL_RPM", 3000) HTTP_GLOBAL_BURST = _env_int("RATE_LIMIT_HTTP_GLOBAL_BURST", 600) # 创建 session：防止刷 session 导致内存爆 HTTP_CREATE_SESSION_RPM = _env_int("RATE_LIMIT_CREATE_SESSION_RPM", 3000) HTTP_CREATE_SESSION_BURST = _env_int("RATE_LIMIT_CREATE_SESSION_BURST", 50) # 上传素材：最容易被滥用（大文件 + 频率） HTTP_UPLOAD_MEDIA_RPM = _env_int("RATE_LIMIT_UPLOAD_MEDIA_RPM", 12000) HTTP_UPLOAD_MEDIA_BURST = _env_int("RATE_LIMIT_UPLOAD_MEDIA_BURST", 300) # 上传“成本”换算：content-length 每多少字节算 1 个 token（越大越费 token） UPLOAD_COST_BYTES = _env_int("RATE_LIMIT_UPLOAD_COST_BYTES", 10 * 1024 * 1024) # 默认 10MB = 1 token # 素材个数控制：会话内上线+上传上限 MAX_UPLOAD_FILES_PER_REQUEST = _env_int("MAX_UPLOAD_FILES_PER_REQUEST", 30) # 单次请求最多文件数 MAX_MEDIA_PER_SESSION = _env_int("MAX_MEDIA_PER_SESSION", 30) # 每个 session 总素材上限（pending + 已用） MAX_PENDING_MEDIA_PER_SESSION = _env_int("MAX_PENDING_MEDIA_PER_SESSION", 30) # 每个 session pending 素材上限（UI 友好） HTTP_UPLOAD_MEDIA_COUNT_RPM = _env_int("RATE_LIMIT_UPLOAD_MEDIA_COUNT_RPM", 50000) HTTP_UPLOAD_MEDIA_COUNT_BURST = _env_int("RATE_LIMIT_UPLOAD_MEDIA_COUNT_BURST", 1000) # 下载/缩略图：适中限制（防刷资源） HTTP_MEDIA_GET_RPM = _env_int("RATE_LIMIT_MEDIA_GET_RPM", 2400) HTTP_MEDIA_GET_BURST = _env_int("RATE_LIMIT_MEDIA_GET_BURST", 60) # 清空会话：避免频繁清空扰动 HTTP_CLEAR_RPM = _env_int("RATE_LIMIT_CLEAR_SESSION_RPM", 3000) HTTP_CLEAR_BURST = _env_int("RATE_LIMIT_CLEAR_SESSION_BURST", 50) # 其它 API 默认：比 global 更细一点（可选） HTTP_API_RPM = _env_int("RATE_LIMIT_API_RPM", 2400) HTTP_API_BURST = _env_int("RATE_LIMIT_API_BURST", 120) # WebSocket：连接创建频率 WS_CONNECT_RPM = _env_int("RATE_LIMIT_WS_CONNECT_RPM", 600) WS_CONNECT_BURST = _env_int("RATE_LIMIT_WS_CONNECT_BURST", 50) # WebSocket：chat.send（真正触发 LLM 成本） WS_CHAT_SEND_RPM = _env_int("RATE_LIMIT_WS_CHAT_SEND_RPM", 300) WS_CHAT_SEND_BURST = _env_int("RATE_LIMIT_WS_CHAT_SEND_BURST", 20) # ---- 全局（所有 IP 合并）限流：抵御多 IP 同时访问 ---- HTTP_ALL_RPM = _env_int("RATE_LIMIT_HTTP_ALL_RPM", 1200) # 全站 HTTP 总量：1200/min ~= 20 rps HTTP_ALL_BURST = _env_int("RATE_LIMIT_HTTP_ALL_BURST", 200) CREATE_SESSION_ALL_RPM = _env_int("RATE_LIMIT_CREATE_SESSION_ALL_RPM", 120) CREATE_SESSION_ALL_BURST = _env_int("RATE_LIMIT_CREATE_SESSION_ALL_BURST", 20) UPLOAD_MEDIA_ALL_RPM = _env_int("RATE_LIMIT_UPLOAD_MEDIA_ALL_RPM", 6000) UPLOAD_MEDIA_ALL_BURST = _env_int("RATE_LIMIT_UPLOAD_MEDIA_ALL_BURST", 2000) # “素材个数”限流：默认复用 upload_media 的 rpm/burst UPLOAD_MEDIA_COUNT_ALL_RPM = _env_int("RATE_LIMIT_UPLOAD_MEDIA_COUNT_ALL_RPM", UPLOAD_MEDIA_ALL_RPM) UPLOAD_MEDIA_COUNT_ALL_BURST = _env_int("RATE_LIMIT_UPLOAD_MEDIA_COUNT_ALL_BURST", UPLOAD_MEDIA_ALL_BURST) MEDIA_GET_ALL_RPM = _env_int("RATE_LIMIT_MEDIA_GET_ALL_RPM", 600) MEDIA_GET_ALL_BURST = _env_int("RATE_LIMIT_MEDIA_GET_ALL_BURST", 120) WS_CONNECT_ALL_RPM = _env_int("RATE_LIMIT_WS_CONNECT_ALL_RPM", 60000) WS_CONNECT_ALL_BURST = _env_int("RATE_LIMIT_WS_CONNECT_ALL_BURST", 2000) WS_CHAT_SEND_ALL_RPM = _env_int("RATE_LIMIT_WS_CHAT_SEND_ALL_RPM", 500) WS_CHAT_SEND_ALL_BURST = _env_int("RATE_LIMIT_WS_CHAT_SEND_ALL_BURST", 30) # ---- 全局并发上限：抵御“很多 IP 同时连/同时触发 LLM/同时上传” ---- WS_MAX_CONNECTIONS = _env_int("RATE_LIMIT_WS_MAX_CONNECTIONS", 500) # 同时在线 WS 连接数上限 CHAT_MAX_CONCURRENCY = _env_int("RATE_LIMIT_CHAT_MAX_CONCURRENCY", 80) # 同时跑的 LLM turn 上限 UPLOAD_MAX_CONCURRENCY = _env_int("RATE_LIMIT_UPLOAD_MAX_CONCURRENCY", 100) # 同时处理上传（含缩略图）上限 WS_CONN_SEM = asyncio.Semaphore(WS_MAX_CONNECTIONS) CHAT_TURN_SEM = asyncio.Semaphore(CHAT_MAX_CONCURRENCY) UPLOAD_SEM = asyncio.Semaphore(UPLOAD_MAX_CONCURRENCY) def _global_http_rule_limit(rule_name: str) -> Optional[Tuple[int, int]]: if rule_name == "create_session": return CREATE_SESSION_ALL_BURST, CREATE_SESSION_ALL_RPM if rule_name == "upload_media": return UPLOAD_MEDIA_ALL_BURST, UPLOAD_MEDIA_ALL_RPM if rule_name == "media_get": return MEDIA_GET_ALL_BURST, MEDIA_GET_ALL_RPM return None def _get_content_length(scope: dict) -> Optional[int]: try: headers = _headers_to_dict(scope.get("headers") or []) v = headers.get("content-length") if v is None: return None n = int(v) if n < 0: return None return n except Exception: return None def _match_http_rule(method: str, path: str) -> Tuple[str, int, int, float]: """ 返回 (rule_name, burst, rpm, cost) cost 默认为 1；上传接口会按 content-length 动态计算 cost（在 middleware 内处理）。 """ method = (method or "").upper() path = path or "" # 精确接口优先 if method == "POST" and path == "/api/sessions": return ("create_session", HTTP_CREATE_SESSION_BURST, HTTP_CREATE_SESSION_RPM, 1.0) # 上传素材（含分片接口） if method == "POST" and path.startswith("/api/sessions/"): if path.endswith("/media") or path.endswith("/media/init"): return ("upload_media", HTTP_UPLOAD_MEDIA_BURST, HTTP_UPLOAD_MEDIA_RPM, 1.0) if "/media/" in path and (path.endswith("/chunk") or path.endswith("/complete") or path.endswith("/cancel")): return ("upload_media", HTTP_UPLOAD_MEDIA_BURST, HTTP_UPLOAD_MEDIA_RPM, 1.0) if method == "GET" and path.startswith("/api/sessions/") and (path.endswith("/thumb") or path.endswith("/file")): return ("media_get", HTTP_MEDIA_GET_BURST, HTTP_MEDIA_GET_RPM, 1.0) if method == "POST" and path.startswith("/api/sessions/") and path.endswith("/clear"): return ("clear_session", HTTP_CLEAR_BURST, HTTP_CLEAR_RPM, 1.0) # 其它 API if path.startswith("/api/"): return ("api_general", HTTP_API_BURST, HTTP_API_RPM, 1.0) # 非 /api 的其他请求：只走 global return ("", 0, 0, 1.0) class HttpRateLimitMiddleware: """ ASGI middleware：对 HTTP 请求做限流（WebSocket 不在这里处理）。 """ def __init__(self, app: Any, limiter: TokenBucketRateLimiter, trust_proxy_headers: bool = False): self.app = app self.limiter = limiter self.trust_proxy_headers = bool(trust_proxy_headers) async def __call__(self, scope: dict, receive: Any, send: Any): if scope.get("type") != "http": return await self.app(scope, receive, send) method = scope.get("method", "GET") path = scope.get("path", "/") ip = _client_ip_from_http_scope(scope, self.trust_proxy_headers) # 0) 全局总量桶（所有 IP 合并） ok, retry_after, _ = await self.limiter.allow( key="http:all", capacity=float(HTTP_ALL_BURST), refill_rate=_rpm_to_rps(float(HTTP_ALL_RPM)), cost=1.0, ) if not ok: return await self._reject(send, retry_after) # 1) 单 IP 全局桶（防单点） ok, retry_after, _ = await self.limiter.allow( key=f"http:global:{ip}", capacity=float(HTTP_GLOBAL_BURST), refill_rate=_rpm_to_rps(float(HTTP_GLOBAL_RPM)), cost=1.0, ) if not ok: return await self._reject(send, retry_after) # 2) 规则桶 rule_name, burst, rpm, cost = _match_http_rule(method, path) # 上传接口：按 content-length 动态增加 cost（越大越费 token） if rule_name == "upload_media": cl = _get_content_length(scope) if cl and cl > 0 and UPLOAD_COST_BYTES > 0: cost = max(1.0, float(math.ceil(cl / float(UPLOAD_COST_BYTES)))) if rule_name: # 2.1 规则的“全局桶”（跨 IP） g = _global_http_rule_limit(rule_name) if g: g_burst, g_rpm = g okg, rag, _ = await self.limiter.allow( key=f"http:{rule_name}:all", capacity=float(g_burst), refill_rate=_rpm_to_rps(float(g_rpm)), cost=float(cost), ) if not okg: return await self._reject(send, rag) # 2.2 规则的“单 IP 桶” ok2, retry_after2, _ = await self.limiter.allow( key=f"http:{rule_name}:{ip}", capacity=float(burst), refill_rate=_rpm_to_rps(float(rpm)), cost=float(cost), ) if not ok2: return await self._reject(send, retry_after2) return await self.app(scope, receive, send) async def _reject(self, send: Any, retry_after: float): ra = int(math.ceil(float(retry_after or 0.0))) body = json.dumps( {"detail": "Too Many Requests", "retry_after": ra}, ensure_ascii=False ).encode("utf-8") headers = [ (b"content-type", b"application/json; charset=utf-8"), (b"retry-after", str(ra).encode("ascii")), ] await send({"type": "http.response.start", "status": 429, "headers": headers}) await send({"type": "http.response.body", "body": body, "more_body": False}) RATE_LIMITER = TokenBucketRateLimiter( ttl_sec=_env_int("RATE_LIMIT_TTL_SEC", 900), # 默认 15 分钟：多 IP 攻击时更快释放桶表 cleanup_interval_sec=_env_int("RATE_LIMIT_CLEANUP_INTERVAL_SEC", 60), max_buckets=_env_int("RATE_LIMIT_MAX_BUCKETS", 100000), evict_batch=_env_int("RATE_LIMIT_EVICT_BATCH", 2000), ) @dataclass class MediaMeta: id: str name: str kind: str path: str thumb_path: Optional[str] ts: float @dataclass class ResumableUpload: upload_id: str filename: str # 素材原名（用于UI展示） store_filename: str # 落盘名 media_0001.mp4 size: int chunk_size: int total_chunks: int tmp_path: str kind: str created_ts: float last_ts: float received: Set[int] = field(default_factory=set) closed: bool = False lock: asyncio.Lock = field(default_factory=asyncio.Lock) class MediaStore: """ 专注文件系统层： - 保存上传文件（async chunk） - 生成缩略图（图片：线程；视频：异步子进程） - 删除文件（只删 media_dir 下的文件） """ def __init__(self, media_dir: str): self.media_dir = os.path.abspath(media_dir) os.makedirs(self.media_dir, exist_ok=True) self.thumbs_dir = ensure_thumbs_dir(self.media_dir) async def save_upload(self, uf: UploadFile, *, store_filename: str, display_name: str) -> MediaMeta: media_id = uuid.uuid4().hex[:10] display_name = sanitize_filename(display_name or uf.filename or "unnamed") store_filename = sanitize_filename(store_filename) kind = detect_media_kind(display_name) save_path = os.path.join(self.media_dir, store_filename) os.makedirs(os.path.dirname(save_path), exist_ok=True) if os.path.exists(save_path): raise HTTPException(status_code=409, detail=f"media filename exists: {store_filename}") # async chunk 写盘（不一次性读入内存） async with await anyio.open_file(save_path, "wb") as out: while True: chunk = await uf.read(CHUNK_SIZE) if not chunk: break await out.write(chunk) try: await uf.close() except Exception: pass thumb_path: Optional[str] = None if kind in ("image", "video"): thumb_path = os.path.join(self.thumbs_dir, f"{media_id}.jpg") if kind == "image": ok = await anyio.to_thread.run_sync(make_image_thumbnail_sync, save_path, thumb_path) else: ok = await make_video_thumbnail_async(save_path, thumb_path) if not ok: # 图片缩略图失败 -> 用原图；视频失败 -> 置空（thumb endpoint 返回占位 SVG） thumb_path = save_path if kind == "image" else None return MediaMeta( id=media_id, name=os.path.basename(display_name), kind=kind, path=os.path.abspath(save_path), thumb_path=os.path.abspath(thumb_path) if thumb_path else None, ts=time.time(), ) async def save_from_path( self, src_path: str, *, store_filename: str, display_name: str, ) -> MediaMeta: """ 将分片上传产生的临时文件移动到 media_dir 下的最终文件。 - display_name: UI 展示名（原始文件名） - store_filename: 落盘名（media_0001.mp4），用于记录顺序 """ media_id = uuid.uuid4().hex[:10] display_name = sanitize_filename(display_name or "unnamed") store_filename = sanitize_filename(store_filename or "unnamed") kind = detect_media_kind(display_name) src_path = os.path.abspath(src_path) if not os.path.exists(src_path): raise HTTPException(status_code=400, detail="upload temp file missing") save_path = os.path.abspath(os.path.join(self.media_dir, store_filename)) os.makedirs(os.path.dirname(save_path), exist_ok=True) if os.path.exists(save_path): raise HTTPException(status_code=409, detail=f"media already exists: {store_filename}") # move tmp -> final os.replace(src_path, save_path) thumb_path: Optional[str] = None if kind in ("image", "video"): thumb_path = os.path.join(self.thumbs_dir, f"{media_id}.jpg") if kind == "image": ok = await anyio.to_thread.run_sync(make_image_thumbnail_sync, save_path, thumb_path) else: ok = await make_video_thumbnail_async(save_path, thumb_path) if not ok: thumb_path = save_path if kind == "image" else None return MediaMeta( id=media_id, name=os.path.basename(display_name), # ★ UI 显示原文件名 kind=kind, path=os.path.abspath(save_path), # ★ 磁盘文件名 media_0001.ext thumb_path=os.path.abspath(thumb_path) if thumb_path else None, ts=time.time(), ) async def delete_files(self, meta: MediaMeta) -> None: root = self.media_dir for p in {meta.path, meta.thumb_path}: if not p: continue ap = os.path.abspath(p) if not _is_under_dir(ap, root): continue if os.path.isdir(ap): continue if os.path.exists(ap): try: os.remove(ap) except Exception: pass class ChatSession: """ 一个 session 的全部状态： - agent / lc_messages（LangChain上下文） - history（给前端回放） - load_media / pending_media（staging） - tool trace 索引（支持 tool 事件“就地更新”） """ def __init__(self, session_id: str, cfg: Settings): self.session_id = session_id self.cfg = cfg self.lang = "zh" default_llm = _s(getattr(getattr(cfg, "developer", None), "default_llm", "")) or "deepseek-chat" default_vlm = _s(getattr(getattr(cfg, "developer", None), "default_vlm", "")) or "qwen3-vl-8b-instruct" self.chat_models = [default_llm, CUSTOM_MODEL_KEY] self.chat_model_key = default_llm self.vlm_models = [default_vlm, CUSTOM_MODEL_KEY] self.vlm_model_key = default_vlm self.developer_mode = is_developer_mode(cfg) self.media_dir = resolve_media_dir(cfg.project.media_dir, session_id) self.media_store = MediaStore(self.media_dir) # 分片上传临时目录 + in-flight 状态 self.uploads_dir = ensure_uploads_dir(self.media_dir) self.resumable_uploads: Dict[str, ResumableUpload] = {} # 直传（multipart 多文件）时的“预占位”，避免并发竞争导致超过上限 self._direct_upload_reservations = 0 self.agent: Any = None self.node_manager = None self.client_context = None # 锁分离：避免“流式输出”阻塞上传/删除 pending self.chat_lock = asyncio.Lock() self.media_lock = asyncio.Lock() self.sent_media_total: int = 0 self._attach_stats_msg_idx = 1 self.lc_messages: List[BaseMessage] = [ SystemMessage(content=get_prompt("instruction.system", lang=self.lang)), SystemMessage(content="【User media upload status】{}"), ] self.history: List[Dict[str, Any]] = [] self.load_media: Dict[str, MediaMeta] = {} self.pending_media_ids: List[str] = [] self._tool_history_index: Dict[str, int] = {} # tool_call_id -> history index self.cancel_event = asyncio.Event() # 打断信号 # 服务相关配置 self.custom_llm_config: Optional[Dict[str, Any]] = None self.custom_vlm_config: Optional[Dict[str, Any]] = None self.tts_config: Dict[str, Any] = {} self._agent_build_key: Optional[Tuple[Any, ...]] = None self.pexels_key_mode: str = "default" # "default" | "custom" self.pexels_custom_key: str = "" self._media_seq_inited = False self._media_seq_next = 1 def _ensure_system_prompt(self) -> None: sys = (get_prompt("instruction.system", lang=self.lang) or "").strip() if not sys: return for m in self.lc_messages: if isinstance(m, SystemMessage) and (getattr(m, "content", "") or "").strip() == sys: return self.lc_messages.insert(0, SystemMessage(content=sys)) def _init_media_seq_locked(self) -> None: """ 初始化 self._media_seq_next： - 允许 clear chat 后继续编号，不覆盖旧文件 """ if self._media_seq_inited: return max_seq = 0 # 1) 已落盘文件 try: for fn in os.listdir(self.media_dir): s = parse_media_seq(fn) if s is not None: max_seq = max(max_seq, s) except Exception: pass # 2) 内存里已有 load_media（保险） for meta in (self.load_media or {}).values(): s = parse_media_seq(os.path.basename(meta.path or "")) if s is not None: max_seq = max(max_seq, s) # 3) in-flight resumable（保险） for u in (self.resumable_uploads or {}).values(): s = parse_media_seq(getattr(u, "store_filename", "") or "") if s is not None: max_seq = max(max_seq, s) self._media_seq_next = max_seq + 1 self._media_seq_inited = True def _reserve_store_filenames_locked(self, display_filenames: List[str]) -> List[str]: """ 按传入顺序生成一组 store 文件名（media_0001.ext ...） 注意：这里的“顺序”就是你要固化的上传顺序。 """ self._init_media_seq_locked() out: List[str] = [] seq = int(self._media_seq_next) for disp in display_filenames: disp = sanitize_filename(disp or "unnamed") ext = os.path.splitext(disp)[1].lower() # 不复用旧号；仅在极端情况下跳过已存在文件（防撞） while True: store = make_media_store_filename(seq, ext) if not os.path.exists(os.path.join(self.media_dir, store)): break seq += 1 out.append(store) seq += 1 self._media_seq_next = seq return out def apply_service_config(self, service_cfg: Any) -> Tuple[bool, Optional[str]]: llm, vlm, tts, pexels, err = _parse_service_config(service_cfg) if err: return False, err if llm is not None: self.custom_llm_config = llm if vlm is not None: self.custom_vlm_config = vlm # tts 允许为空；非空才覆盖 if isinstance(tts, dict) and tts: self.tts_config = tts # ---- pexels ---- if isinstance(pexels, dict) and pexels: mode = _s(pexels.get("mode")).lower() if mode == "custom": self.pexels_key_mode = "custom" self.pexels_custom_key = _s(pexels.get("api_key")) else: self.pexels_key_mode = "default" self.pexels_custom_key = "" return True, None async def ensure_agent(self) -> None: # 1) resolve LLM override if self.chat_model_key == CUSTOM_MODEL_KEY: if not isinstance(self.custom_llm_config, dict): raise RuntimeError("please fill in model/base_url/api_key of custom LLM") llm_override = self.custom_llm_config else: llm_override, err = _resolve_default_model_override(self.cfg, self.chat_model_key) if err: raise RuntimeError(err) # 2) resolve VLM override if self.vlm_model_key == CUSTOM_MODEL_KEY: if not isinstance(self.custom_vlm_config, dict): raise RuntimeError("please fill in model/base_url/api_key of custom VLM") vlm_override = self.custom_vlm_config else: vlm_override, err = _resolve_default_model_override(self.cfg, self.vlm_model_key) if err: raise RuntimeError(err) agent_build_key: Tuple[Any, ...] = ( "models", _stable_dict_key(llm_override), _stable_dict_key(vlm_override), ) if self.agent is None or self._agent_build_key != agent_build_key: artifact_store = ArtifactStore(self.cfg.project.outputs_dir, session_id=self.session_id) self.agent, self.node_manager = await build_agent( cfg=self.cfg, session_id=self.session_id, store=artifact_store, tool_interceptors=[ ToolInterceptor.inject_media_content_before, ToolInterceptor.save_media_content_after, ToolInterceptor.inject_tts_config, ToolInterceptor.inject_pexels_api_key, ], llm_override=llm_override, vlm_override=vlm_override, ) self._agent_build_key = agent_build_key if self.client_context is None: self.client_context = ClientContext( cfg=self.cfg, session_id=self.session_id, media_dir=self.media_dir, bgm_dir=self.cfg.project.bgm_dir, outputs_dir=self.cfg.project.outputs_dir, node_manager=self.node_manager, chat_model_key=self.chat_model_key, vlm_model_key=self.vlm_model_key, tts_config=(self.tts_config or None), pexels_api_key=None, lang=self.lang, ) else: self.client_context.chat_model_key = self.chat_model_key self.client_context.vlm_model_key = self.vlm_model_key self.client_context.tts_config = (self.tts_config or None) self.client_context.lang = self.lang # ---- resolve pexels_api_key for runtime context ---- pexels_api_key = "" if (self.pexels_key_mode or "").lower() == "custom": pexels_api_key = _s(self.pexels_custom_key) else: pexels_api_key = _get_default_pexels_api_key(self.cfg) # from config.toml self.client_context.pexels_api_key = (pexels_api_key or None) # ---- DTO / public mapping ---- def public_media(self, meta: MediaMeta) -> Dict[str, Any]: return { "id": meta.id, "name": meta.name, "kind": meta.kind, "thumb_url": f"/api/sessions/{self.session_id}/media/{meta.id}/thumb", "file_url": f"/api/sessions/{self.session_id}/media/{meta.id}/file", } def public_pending_media(self) -> List[Dict[str, Any]]: out: List[Dict[str, Any]] = [] for aid in self.pending_media_ids: meta = self.load_media.get(aid) if meta: out.append(self.public_media(meta)) return out def snapshot(self) -> Dict[str, Any]: return { "session_id": self.session_id, "developer_mode": self.developer_mode, "pending_media": self.public_pending_media(), "history": self.history, "limits": { "max_upload_files_per_request": MAX_UPLOAD_FILES_PER_REQUEST, "max_media_per_session": MAX_MEDIA_PER_SESSION, "max_pending_media_per_session": MAX_PENDING_MEDIA_PER_SESSION, "upload_chunk_bytes": UPLOAD_RESUMABLE_CHUNK_BYTES, }, "stats": { "media_count": len(self.load_media), "pending_count": len(self.pending_media_ids), "inflight_uploads": len(self.resumable_uploads), }, "chat_model_key": self.chat_model_key, "chat_models": self.chat_models, "llm_model_key": self.chat_model_key, "llm_models": self.chat_models, "vlm_model_key": self.vlm_model_key, "vlm_models": self.vlm_models, "lang": self.lang, } # ---- media operations ---- def _cleanup_stale_uploads_locked(self, now: Optional[float] = None) -> None: now = float(now or time.time()) ttl = float(RESUMABLE_UPLOAD_TTL_SEC) dead = [uid for uid, u in self.resumable_uploads.items() if (now - u.last_ts) > ttl] for uid in dead: u = self.resumable_uploads.pop(uid, None) if not u: continue try: if u.tmp_path and os.path.exists(u.tmp_path): os.remove(u.tmp_path) except Exception: pass def _check_media_caps_locked(self, add: int = 0) -> None: add = int(max(0, add)) total = len(self.load_media) + len(self.resumable_uploads) + int(self._direct_upload_reservations) pending = len(self.pending_media_ids) + len(self.resumable_uploads) + int(self._direct_upload_reservations) if MAX_MEDIA_PER_SESSION > 0 and (total + add) > MAX_MEDIA_PER_SESSION: raise HTTPException( status_code=400, detail=f"会话素材总数已达上限：{total}/{MAX_MEDIA_PER_SESSION}", ) if MAX_PENDING_MEDIA_PER_SESSION > 0 and (pending + add) > MAX_PENDING_MEDIA_PER_SESSION: raise HTTPException( status_code=400, detail=f"待发送素材数量已达上限：{pending}/{MAX_PENDING_MEDIA_PER_SESSION}", ) async def add_uploads(self, files: List[UploadFile], store_filenames: List[str]) -> List[MediaMeta]: if len(store_filenames) != len(files): raise HTTPException(status_code=500, detail="store_filenames mismatch") metas: List[MediaMeta] = [] for uf, store_fn in zip(files, store_filenames): display_name = sanitize_filename(uf.filename or "unnamed") metas.append(await self.media_store.save_upload( uf, store_filename=store_fn, display_name=display_name, )) async with self.media_lock: for m in metas: self.load_media[m.id] = m self.pending_media_ids.append(m.id) self.pending_media_ids.sort( key=lambda aid: os.path.basename(self.load_media[aid].path or "") if aid in self.load_media else "" ) return metas async def delete_pending_media(self, media_id: str) -> None: async with self.media_lock: if media_id not in self.pending_media_ids: raise HTTPException(status_code=400, detail="media is not pending (refuse physical delete)") self.pending_media_ids = [x for x in self.pending_media_ids if x != media_id] meta = self.load_media.pop(media_id, None) if meta: await self.media_store.delete_files(meta) async def take_pending_media_for_message(self, attachment_ids: Optional[List[str]]) -> List[MediaMeta]: async with self.media_lock: if attachment_ids: pick = [aid for aid in attachment_ids if aid in self.pending_media_ids] else: pick = list(self.pending_media_ids) pick_set = set(pick) self.pending_media_ids = [aid for aid in self.pending_media_ids if aid not in pick_set] metas = [self.load_media[aid] for aid in pick if aid in self.load_media] return metas # ---- tool trace handling ---- def _ensure_tool_record(self, tcid: str, server: str, name: str, args: Any) -> Dict[str, Any]: idx = self._tool_history_index.get(tcid) if idx is None: rec = { "id": f"tool_{tcid}", "role": "tool", "tool_call_id": tcid, "server": server, "name": name, "args": args, "state": "running", "progress": 0.0, "message": "", "summary": None, "ts": time.time(), } self.history.append(rec) self._tool_history_index[tcid] = len(self.history) - 1 return rec return self.history[idx] def apply_tool_event(self, raw: Dict[str, Any]) -> Optional[Dict[str, Any]]: et = raw.get("type") tcid = raw.get("tool_call_id") if et not in ("tool_start", "tool_progress", "tool_end") or not tcid: return None server = raw.get("server") or "" name = raw.get("name") or "" args = raw.get("args") or {} rec = self._ensure_tool_record(tcid, server, name, args) if et == "tool_start": rec.update({ "server": server, "name": name, "args": args, "state": "running", "progress": 0.0, "message": "Starting...", "summary": None, }) elif et == "tool_progress": progress = float(raw.get("progress", 0.0)) total = raw.get("total") if total and float(total) > 0: p = progress / float(total) else: p = progress / 100.0 if progress > 1 else progress p = max(0.0, min(1.0, p)) rec.update({ "state": "running", "progress": p, "message": raw.get("message") or "", }) elif et == "tool_end": is_error = bool(raw.get("is_error")) summary = raw.get("summary") try: json.dumps(summary, ensure_ascii=False) except Exception: summary = str(summary) if summary is not None else None rec.update({ "state": "error" if is_error else "complete", "progress": 1.0, "summary": summary, "message": raw.get("message") or rec.get("message") or "", }) return rec class SessionStore: def __init__(self, cfg: Settings): self.cfg = cfg self._lock = asyncio.Lock() self._sessions: Dict[str, ChatSession] = {} async def create(self) -> ChatSession: sid = uuid.uuid4().hex sess = ChatSession(sid, self.cfg) async with self._lock: self._sessions[sid] = sess return sess async def get(self, sid: str) -> Optional[ChatSession]: async with self._lock: return self._sessions.get(sid) async def get_or_404(self, sid: str) -> ChatSession: sess = await self.get(sid) if not sess: raise HTTPException(status_code=404, detail="session not found") return sess @asynccontextmanager async def lifespan(app: FastAPI): cfg = load_settings(default_config_path()) app.state.cfg = cfg app.state.developer_mode = is_developer_mode(cfg) app.state.sessions = SessionStore(cfg) yield app = FastAPI(title="OpenStoryline Web", version="1.0.0", lifespan=lifespan) app.add_middleware( HttpRateLimitMiddleware, limiter=RATE_LIMITER, trust_proxy_headers=RATE_LIMIT_TRUST_PROXY_HEADERS, ) if os.path.isdir(STATIC_DIR): app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static") if os.path.isdir(NODE_MAP_DIR): app.mount("/node_map", StaticFiles(directory=NODE_MAP_DIR), name="node_map") api = APIRouter(prefix="/api") def _rate_limit_reject_json(retry_after: float) -> JSONResponse: ra = int(math.ceil(float(retry_after or 0.0))) return JSONResponse( {"detail": "Too Many Requests", "retry_after": ra}, status_code=429, headers={"Retry-After": str(ra)}, ) async def _enforce_upload_media_count_limit(request: Request, cost: float) -> Optional[JSONResponse]: ip = _client_ip_from_http_scope(request.scope, RATE_LIMIT_TRUST_PROXY_HEADERS) cost = float(max(0.0, cost)) ok, ra, _ = await RATE_LIMITER.allow( key="http:upload_media_count:all", capacity=float(UPLOAD_MEDIA_COUNT_ALL_BURST), refill_rate=_rpm_to_rps(float(UPLOAD_MEDIA_COUNT_ALL_RPM)), cost=cost, ) if not ok: return _rate_limit_reject_json(ra) ok2, ra2, _ = await RATE_LIMITER.allow( key=f"http:upload_media_count:{ip}", capacity=float(HTTP_UPLOAD_MEDIA_COUNT_BURST), refill_rate=_rpm_to_rps(float(HTTP_UPLOAD_MEDIA_COUNT_RPM)), cost=cost, ) if not ok2: return _rate_limit_reject_json(ra2) return None _TTS_UI_SECRET_KEYS = { "api_key", "access_token", "authorization", "token", "password", "secret", "x-api-key", "apikey", "access_key", "accesskey", } def _is_secret_field_name(k: str) -> bool: if str(k or "").strip().lower() in _TTS_UI_SECRET_KEYS: return True return False def _read_config_toml(path: str) -> dict: if tomllib is None: return {} try: p = Path(path) with p.open("rb") as f: return tomllib.load(f) or {} except Exception: return {} def _get_default_pexels_api_key(cfg: Settings) -> str: # 1) try Settings.search_media.pexels_api_key try: search_media = getattr(cfg, "search_media", None) pexels_api_key = _s(getattr(search_media, "pexels_api_key", None) if search_media else None) if pexels_api_key: return pexels_api_key else: return "" except Exception: return "" def _normalize_field_item(item) -> dict | None: """ item 支持： - "uid" - { key="uid", label="UID", required=true, secret=false, placeholder="..." } """ if isinstance(item, str): key = item.strip() if not key: return None return { "key": key, "secret": _is_secret_field_name(key), } return None def _build_provider_schema(provider: str, label: str | None, fields: list[dict]) -> dict: seen = set() out = [] for f in fields: k = str(f.get("key") or "").strip() if not k or k in seen: continue seen.add(k) out.append({ "key": k, "label": f.get("label") or k, "placeholder": f.get("placeholder") or f.get("label") or k, "required": bool(f.get("required", False)), "secret": bool(f.get("secret", False)), }) return {"provider": provider, "label": label or provider, "fields": out} def _build_tts_ui_schema_from_config(config_path: str) -> dict: """ 返回： { "providers": [ {"provider":"bytedance","label":"字节跳动","fields":[{"key":"uid",...}, ...]}, ... ] } """ cfg = _read_config_toml(config_path) tts = cfg.get("generate_voiceover", {}) providers_out: list[dict] = [] # 格式：[tts.providers.] providers = tts.get("providers") if isinstance(providers, dict): for provider, provider_cfg in providers.items(): fields: list[dict] = [] label = str(provider_cfg.get("label") or provider_cfg.get("name") or provider) for key in provider_cfg.keys(): f = _normalize_field_item(str(key)) if f: fields.append(f) providers_out.append(_build_provider_schema(provider, label, fields)) return {"providers": providers_out} @app.get("/") async def index(): if not os.path.exists(INDEX_HTML): return Response("index.html not found. Put it under ./web/index.html", media_type="text/plain", status_code=404) return FileResponse(INDEX_HTML, media_type="text/html") @app.get("/node-map") async def node_map(): if not os.path.exists(NODE_MAP_HTML): return Response( "node_map.html not found. Put it under ./web/node_map/node_map.html", media_type="text/plain", status_code=404, ) return FileResponse(NODE_MAP_HTML, media_type="text/html") @api.get("/meta/tts") async def get_tts_ui_schema(): schema = _build_tts_ui_schema_from_config(default_config_path()) return JSONResponse(schema) # ------------------------- # Sessions (REST) # ------------------------- @api.post("/sessions") async def create_session(): store: SessionStore = app.state.sessions sess = await store.create() return JSONResponse(sess.snapshot()) @api.get("/sessions/{session_id}") async def get_session(session_id: str): store: SessionStore = app.state.sessions sess = await store.get_or_404(session_id) return JSONResponse(sess.snapshot()) @api.post("/sessions/{session_id}/clear") async def clear_session_chat(session_id: str): store: SessionStore = app.state.sessions sess = await store.get_or_404(session_id) async with sess.chat_lock: sess.sent_media_total = 0 sess._attach_stats_msg_idx = 1 sess.lc_messages = [ SystemMessage(content=get_prompt("instruction.system", lang=sess.lang)), SystemMessage(content="【User media upload status】{}"), ] sess._attach_stats_msg_idx = 1 sess.history = [] sess._tool_history_index = {} return JSONResponse({"ok": True}) @api.post("/sessions/{session_id}/cancel") async def cancel_session_turn(session_id: str): """ 打断当前正在进行的 LLM turn（流式回复/工具调用）。 - 不清空 history / lc_messages - 仅设置 cancel_event，由 WS 侧在流式循环中感知并安全收尾 """ store: SessionStore = app.state.sessions sess = await store.get_or_404(session_id) sess.cancel_event.set() return JSONResponse({"ok": True}) # ------------------------- # media (REST, session-scoped) # ------------------------- @api.post("/sessions/{session_id}/media") async def upload_media(session_id: str, request: Request, files: List[UploadFile] = File(...)): if not isinstance(files, list) or not files: raise HTTPException(status_code=400, detail="no files") if MAX_UPLOAD_FILES_PER_REQUEST > 0 and len(files) > MAX_UPLOAD_FILES_PER_REQUEST: raise HTTPException(status_code=400, detail=f"单次上传最多 {MAX_UPLOAD_FILES_PER_REQUEST} 个文件") # 按素材个数限流（cost = 文件数） rej = await _enforce_upload_media_count_limit(request, cost=float(len(files))) if rej: return rej if UPLOAD_SEM.locked(): raise HTTPException(status_code=429, detail="上传并发过高，请稍后重试") await UPLOAD_SEM.acquire() n = len(files) try: store: SessionStore = app.state.sessions sess = await store.get_or_404(session_id) # session cap 检查 + 预占位（避免并发竞争） async with sess.media_lock: sess._cleanup_stale_uploads_locked() sess._check_media_caps_locked(add=n) sess._direct_upload_reservations += n display_names = [sanitize_filename(uf.filename or "unnamed") for uf in files] store_filenames = sess._reserve_store_filenames_locked(display_names) try: metas = await sess.add_uploads(files, store_filenames=store_filenames) finally: async with sess.media_lock: sess._direct_upload_reservations = max(0, sess._direct_upload_reservations - n) return JSONResponse({ "media": [sess.public_media(m) for m in metas], "pending_media": sess.public_pending_media(), }) finally: try: UPLOAD_SEM.release() except Exception: pass @api.post("/sessions/{session_id}/media/init") async def init_resumable_media_upload(session_id: str, request: Request): try: data = await request.json() if not isinstance(data, dict): data = {} except Exception: data = {} filename = sanitize_filename((data.get("filename") or data.get("name") or "unnamed")) size = int(data.get("size") or 0) if size <= 0: raise HTTPException(status_code=400, detail="invalid size") # 按素材个数限流：init 视为“新增 1 个素材” rej = await _enforce_upload_media_count_limit(request, cost=1.0) if rej: return rej store: SessionStore = app.state.sessions sess = await store.get_or_404(session_id) async with sess.media_lock: sess._cleanup_stale_uploads_locked() sess._check_media_caps_locked(add=1) store_filename = sess._reserve_store_filenames_locked([filename])[0] upload_id = uuid.uuid4().hex chunk_size = int(max(1, UPLOAD_RESUMABLE_CHUNK_BYTES)) total_chunks = int(math.ceil(size / float(chunk_size))) tmp_path = os.path.join(sess.uploads_dir, f"{upload_id}.part") os.makedirs(os.path.dirname(tmp_path), exist_ok=True) try: with open(tmp_path, "wb"): pass except Exception as e: raise HTTPException(status_code=500, detail=f"cannot create temp file: {e}") u = ResumableUpload( upload_id=upload_id, filename=filename, store_filename=store_filename, size=size, chunk_size=chunk_size, total_chunks=total_chunks, tmp_path=os.path.abspath(tmp_path), kind=detect_media_kind(filename), created_ts=time.time(), last_ts=time.time(), ) sess.resumable_uploads[upload_id] = u return JSONResponse({ "upload_id": upload_id, "chunk_size": chunk_size, "total_chunks": total_chunks, "filename": filename, }) @api.post("/sessions/{session_id}/media/{upload_id}/chunk") async def upload_resumable_media_chunk( session_id: str, upload_id: str, index: int = Form(...), chunk: UploadFile = File(...), ): if UPLOAD_SEM.locked(): raise HTTPException(status_code=429, detail="上传并发过高，请稍后重试") await UPLOAD_SEM.acquire() try: store: SessionStore = app.state.sessions sess = await store.get_or_404(session_id) async with sess.media_lock: sess._cleanup_stale_uploads_locked() u = sess.resumable_uploads.get(upload_id) if not u: raise HTTPException(status_code=404, detail="upload_id not found or expired") idx = int(index) if idx < 0 or idx >= u.total_chunks: raise HTTPException(status_code=400, detail="invalid chunk index") # 期望长度（最后一片可能小于 chunk_size） expected_len = u.size - idx * u.chunk_size if expected_len <= 0: raise HTTPException(status_code=400, detail="invalid chunk index") expected_len = min(u.chunk_size, expected_len) written = 0 async with u.lock: if u.closed: raise HTTPException(status_code=400, detail="upload already closed") async with await anyio.open_file(u.tmp_path, "r+b") as out: await out.seek(idx * u.chunk_size) while True: buf = await chunk.read(CHUNK_SIZE) if not buf: break written += len(buf) if written > expected_len: raise HTTPException(status_code=400, detail="chunk too large") await out.write(buf) try: await chunk.close() except Exception: pass if written != expected_len: raise HTTPException(status_code=400, detail=f"chunk size mismatch: {written} != {expected_len}") u.received.add(idx) u.last_ts = time.time() return JSONResponse({ "ok": True, "received_chunks": len(u.received), "total_chunks": u.total_chunks, }) finally: try: UPLOAD_SEM.release() except Exception: pass @api.post("/sessions/{session_id}/media/{upload_id}/complete") async def complete_resumable_media_upload(session_id: str, upload_id: str): if UPLOAD_SEM.locked(): raise HTTPException(status_code=429, detail="上传并发过高，请稍后重试") await UPLOAD_SEM.acquire() try: store: SessionStore = app.state.sessions sess = await store.get_or_404(session_id) async with sess.media_lock: sess._cleanup_stale_uploads_locked() u = sess.resumable_uploads.get(upload_id) if not u: raise HTTPException(status_code=404, detail="upload_id not found or expired") # 锁住此 upload，防止 chunk 并发写 async with u.lock: u.closed = True if len(u.received) != u.total_chunks: missing = u.total_chunks - len(u.received) raise HTTPException(status_code=400, detail=f"chunks missing: {missing}") # 从索引移除（释放会话额度） async with sess.media_lock: u2 = sess.resumable_uploads.pop(upload_id, None) if not u2: raise HTTPException(status_code=404, detail="upload_id not found") meta = await sess.media_store.save_from_path( u2.tmp_path, store_filename=u2.store_filename, display_name=u2.filename, ) async with sess.media_lock: sess.load_media[meta.id] = meta sess.pending_media_ids.append(meta.id) return JSONResponse({ "media": sess.public_media(meta), "pending_media": sess.public_pending_media(), }) finally: try: UPLOAD_SEM.release() except Exception: pass @api.post("/sessions/{session_id}/media/{upload_id}/cancel") async def cancel_resumable_media_upload(session_id: str, upload_id: str): store: SessionStore = app.state.sessions sess = await store.get_or_404(session_id) async with sess.media_lock: u = sess.resumable_uploads.pop(upload_id, None) if not u: return JSONResponse({"ok": True}) async with u.lock: u.closed = True try: if u.tmp_path and os.path.exists(u.tmp_path): os.remove(u.tmp_path) except Exception: pass return JSONResponse({"ok": True}) @api.get("/sessions/{session_id}/media/pending") async def get_pending_media(session_id: str): store: SessionStore = app.state.sessions sess = await store.get_or_404(session_id) return JSONResponse({"pending_media": sess.public_pending_media()}) @api.delete("/sessions/{session_id}/media/pending/{media_id}") async def delete_pending_media(session_id: str, media_id: str): store: SessionStore = app.state.sessions sess = await store.get_or_404(session_id) await sess.delete_pending_media(media_id) return JSONResponse({"ok": True, "pending_media": sess.public_pending_media()}) @api.get("/sessions/{session_id}/media/{media_id}/thumb") async def get_media_thumb(session_id: str, media_id: str): store: SessionStore = app.state.sessions sess = await store.get_or_404(session_id) meta = sess.load_media.get(media_id) if not meta: raise HTTPException(status_code=404, detail="media not found") # thumb 存在优先 if meta.thumb_path and os.path.exists(meta.thumb_path): return FileResponse(meta.thumb_path, media_type="image/jpeg") # video 无 thumb => placeholder if meta.kind == "video": return Response(content=video_placeholder_svg_bytes(), media_type="image/svg+xml") # image thumb 失败 => 用原图 if meta.path and os.path.exists(meta.path): return FileResponse(meta.path, media_type=guess_media_type(meta.path)) raise HTTPException(status_code=404, detail="thumb not available") @api.get("/sessions/{session_id}/media/{media_id}/file") async def get_media_file(session_id: str, media_id: str): store: SessionStore = app.state.sessions sess = await store.get_or_404(session_id) meta = sess.load_media.get(media_id) if not meta: raise HTTPException(status_code=404, detail="media not found") if not meta.path or (not os.path.exists(meta.path)): raise HTTPException(status_code=404, detail="file not found") # 安全：只允许 media_dir 下 if not _is_under_dir(meta.path, sess.media_store.media_dir): raise HTTPException(status_code=403, detail="forbidden") return FileResponse( meta.path, media_type=guess_media_type(meta.path), filename=meta.name, ) @api.get("/sessions/{session_id}/preview") async def preview_local_file(session_id: str, path: str): """ 把 summary.preview_urls 里的“服务器本地路径”安全地转成可访问 URL。 只允许访问：media_dir / outputs_dir / outputs_dir / bgm_dir / .server_cache 这些根目录下的文件。 """ store: SessionStore = app.state.sessions sess = await store.get_or_404(session_id) p = (path or "").strip() if not p: raise HTTPException(status_code=400, detail="empty path") if "\x00" in p: raise HTTPException(status_code=400, detail="bad path") # 兼容 file:// 前缀（如果未来有） if p.startswith("file://"): p = p[len("file://"):] # 相对路径：默认相对 ROOT_DIR if os.path.isabs(p): ap = os.path.abspath(p) else: ap = os.path.abspath(os.path.join(ROOT_DIR, p)) allowed_roots = [ os.path.abspath(sess.media_dir), os.path.abspath(app.state.cfg.project.outputs_dir), os.path.abspath(app.state.cfg.project.outputs_dir), os.path.abspath(app.state.cfg.project.bgm_dir), os.path.abspath(SERVER_CACHE_DIR), ] if not any(_is_under_dir(ap, r) for r in allowed_roots): raise HTTPException(status_code=403, detail="forbidden") if (not os.path.exists(ap)) or os.path.isdir(ap): raise HTTPException(status_code=404, detail="file not found") # 对 cache 文件强缓存 headers = {"Cache-Control": "public, max-age=31536000, immutable"} if _is_under_dir(ap, SERVER_CACHE_DIR) else None return FileResponse( ap, media_type=guess_media_type(ap), filename=os.path.basename(ap), headers=headers, ) app.include_router(api) # ------------------------- # WebSocket: session-scoped chat stream # ------------------------- def extract_text_delta(msg_chunk: Any) -> str: # 兼容 content_blocks (qwen3 常见) blocks = getattr(msg_chunk, "content_blocks", None) or [] if blocks: out = "" for b in blocks: if isinstance(b, dict) and b.get("type") == "text": out += b.get("text", "") return out c = getattr(msg_chunk, "content", "") return c if isinstance(c, str) else "" async def ws_send(ws: WebSocket, type_: str, data: Any = None): if getattr(ws, "client_state", None) != WebSocketState.CONNECTED: return False try: await ws.send_json({"type": type_, "data": data}) return True except WebSocketDisconnect: return False except RuntimeError: return False except Exception as e: if ClientDisconnected is not None and isinstance(e, ClientDisconnected): return False logger.exception("ws_send failed: type=%s err=%r", type_, e) return False @asynccontextmanager async def mcp_sink_context(sink_func): token = set_mcp_log_sink(sink_func) try: yield finally: reset_mcp_log_sink(token) @app.websocket("/ws/sessions/{session_id}/chat") async def ws_chat(ws: WebSocket, session_id: str): client_ip = _client_ip_from_ws(ws, RATE_LIMIT_TRUST_PROXY_HEADERS) ok, retry_after, _ = await RATE_LIMITER.allow( key=f"ws:connect:{client_ip}", capacity=float(WS_CONNECT_BURST), refill_rate=_rpm_to_rps(float(WS_CONNECT_RPM)), cost=1.0, ) if not ok: try: await ws.close(code=1013, reason=f"rate_limited, retry after {int(math.ceil(retry_after))}s") except Exception: debug_traceback_print(app.state.cfg) pass return if WS_CONN_SEM.locked(): try: await ws.close(code=1013, reason="Server busy (websocket connections limit)") except Exception: debug_traceback_print(app.state.cfg) pass return await WS_CONN_SEM.acquire() try: await ws.accept() store: SessionStore = app.state.sessions sess = await store.get(session_id) if not sess: await ws.close(code=4404, reason="session not found") return sess = await store.get_or_404(session_id) await ws_send(ws, "session.snapshot", sess.snapshot()) try: while True: req = await ws.receive_json() if not isinstance(req, dict): continue t = req.get("type") if t == "ping": await ws_send(ws, "pong", {"ts": time.time()}) continue if t == "session.set_lang": data = (req.get("data") or {}) lang = (data.get("lang") or "").strip().lower() if lang not in ("zh", "en"): lang = "zh" sess.lang = lang if sess.client_context: sess.client_context.lang = lang await ws_send(ws, "session.lang", {"lang": lang}) continue if t == "chat.clear": async with sess.chat_lock: sess.sent_media_total = 0 sess._attach_stats_msg_idx = 1 sess.lc_messages = [ SystemMessage(content=get_prompt("instruction.system", lang=sess.lang)), SystemMessage(content="【User media upload status】{}"), ] sess._attach_stats_msg_idx = 1 sess.history = [] sess._tool_history_index = {} await ws_send(ws, "chat.cleared", {"ok": True}) continue if t != "chat.send": await ws_send(ws, "error", {"message": f"unknown type: {t}"}) continue # ---- WebSocket message rate limit: only limit expensive "chat.send" ---- if sess.chat_lock.locked(): await ws_send(ws, "error", {"message": "上一条消息尚未完成，请稍后再发送"}) continue ok, retry_after, _ = await RATE_LIMITER.allow( key="ws:chat_send:all", capacity=float(WS_CHAT_SEND_ALL_BURST), refill_rate=_rpm_to_rps(float(WS_CHAT_SEND_ALL_RPM)), cost=1.0, ) if not ok: await ws_send(ws, "error", { "message": f"触发全局限流：请 {int(math.ceil(retry_after))} 秒后再试", "retry_after": int(math.ceil(retry_after)), }) continue ok, retry_after, _ = await RATE_LIMITER.allow( key=f"ws:chat_send:{client_ip}", capacity=float(WS_CHAT_SEND_BURST), refill_rate=_rpm_to_rps(float(WS_CHAT_SEND_RPM)), cost=1.0, ) if not ok: await ws_send(ws, "error", { "message": f"触发限流：请 {int(math.ceil(retry_after))} 秒后再试", "retry_after": int(math.ceil(retry_after)), }) continue if CHAT_TURN_SEM.locked(): await ws_send(ws, "error", {"message": "服务器繁忙（模型并发已满），请稍后再试"}) continue await CHAT_TURN_SEM.acquire() try: # 再次确认（期间有 await，锁状态可能变化） if sess.chat_lock.locked(): await ws_send(ws, "error", {"message": "上一条消息尚未完成，请稍后再发送"}) continue data = (req.get("data", {}) or {}) prompt = data.get("text", "") prompt = (prompt or "").strip() if not prompt: continue requested_llm = data.get("llm_model") requested_vlm = data.get("vlm_model") attachment_ids = data.get("attachment_ids") if not isinstance(attachment_ids, list): attachment_ids = None async with sess.chat_lock: # 新 turn 开始：清掉上一次残留的 cancel 信号 sess.cancel_event.clear() # 0.0) 应用 service_config（自定义模型 / TTS） ok_cfg, err_cfg = sess.apply_service_config(data.get("service_config")) if not ok_cfg: await ws_send(ws, "error", {"message": err_cfg or "service_config invalid"}) continue # 0) 如果前端传了 model，则更新会话当前对话模型 if isinstance(requested_llm, str): m = requested_llm.strip() if m: sess.chat_model_key = m if sess.client_context: sess.client_context.chat_model_key = m if isinstance(requested_vlm, str): m2 = requested_vlm.strip() if m2: sess.vlm_model_key = m2 if sess.client_context: sess.client_context.vlm_model_key = m2 requested_lang = data.get("lang") if isinstance(requested_lang, str): lang = requested_lang.strip().lower() if lang in ("zh", "en"): sess.lang = lang # 0.1) 可能需要重建 agent（比如切换到 __custom__ 或者自定义配置变化） try: await sess.ensure_agent() except Exception as e: await ws_send(ws, "error", {"message": f"{type(e).__name__}: {e}"}) continue sess._ensure_system_prompt() if sess.client_context: sess.client_context.lang = sess.lang # 1) 从 pending 里拿本次要发送的附件 attachments = await sess.take_pending_media_for_message(attachment_ids) attachments_public = [sess.public_media(m) for m in attachments] # 统计本轮和累计发送了几个素材 turn_attached_count = len(attachments) sess.sent_media_total = int(getattr(sess, "sent_media_total", 0)) + turn_attached_count stats = { "Number of media carried in this message sent by the user": turn_attached_count, "Total number of media sent by the user in all conversations": sess.sent_media_total, "Total number of media in user's media library": scan_media_dir(resolve_media_dir(app.state.cfg.project.media_dir, session_id=session_id)), } idx = int(getattr(sess, "_attach_stats_msg_idx", 1)) if len(sess.lc_messages) <= idx: while len(sess.lc_messages) <= idx: sess.lc_messages.append(SystemMessage(content="")) sess.lc_messages[idx] = SystemMessage( content="【User media upload status】The following fields are used to determine the nature of the media provided by the user: \n" + json.dumps(stats, ensure_ascii=False) ) # 2.1 写入 history + lc context user_msg = { "id": uuid.uuid4().hex[:12], "role": "user", "content": prompt, "attachments": attachments_public, "ts": time.time(), } sess.history.append(user_msg) sess.lc_messages.append(HumanMessage(content=prompt)) # if app.state.cfg.developer.developer_mode: # print("[LLM_CTX]", session_id, sess.lc_messages) # 2.2 ack：让前端更新 pending + 插入 user 消息（前端也可本地先插入） await ws_send(ws, "chat.user", { "text": prompt, "attachments": attachments_public, "pending_media": sess.public_pending_media(), "llm_model_key": sess.chat_model_key, "vlm_model_key": sess.vlm_model_key, }) # 2.3 建立“单通道事件队列”，确保 ws.send_json 不会并发冲突 loop = asyncio.get_running_loop() out_q: asyncio.Queue[Tuple[str, Any]] = asyncio.Queue() def sink(ev: Any): # MCP interceptor 可能 emit 非 dict；这里只收 dict if isinstance(ev, dict): loop.call_soon_threadsafe(out_q.put_nowait, ("mcp", ev)) new_messages: List[BaseMessage] = [] async def pump_agent(): nonlocal new_messages try: stream = sess.agent.astream( {"messages": sess.lc_messages}, context=sess.client_context, stream_mode=["messages", "updates"], ) async for mode, chunk in stream: if mode == "messages": msg_chunk, meta = chunk if meta.get("langgraph_node") == "model": delta = extract_text_delta(msg_chunk) if delta: await out_q.put(("assistant.delta", delta)) elif mode == "updates": if isinstance(chunk, dict): for _step, data in chunk.items(): msgs = (data or {}).get("messages") or [] new_messages.extend(msgs) await out_q.put(("agent.done", None)) except asyncio.CancelledError: # 被用户打断 / 连接关闭导致的取消，不属于“真正异常” # 不要发 agent.error；给主循环一个 cancelled 信号即可 try: out_q.put_nowait(("agent.cancelled", None)) except Exception: debug_traceback_print(app.state.cfg) pass raise # 让任务保持 cancelled 状态，finally 里 await 时会抛 CancelledError except Exception as e: # 关键：异常也要让主循环“可结束”，否则 UI 卡死 await out_q.put(("agent.error", f"{type(e).__name__}: {e}")) async def safe_send(type_: str, data: Any = None) -> bool: try: await ws_send(ws, type_, data) return True except WebSocketDisconnect: return False except RuntimeError as e: # starlette: Cannot call "send" once a close message has been sent. if 'Cannot call "send" once a close message has been sent.' in str(e): return False raise except Exception as e: # uvicorn: ClientDisconnected（不同版本类路径不稳定，用类名兜底） if e.__class__.__name__ == "ClientDisconnected": return False raise # turn 开始（前端可禁用发送按钮/显示占位） if not await ws_send(ws, "assistant.start", {}): return # 当前 assistant 分段缓冲：用于在 tool_start 到来前“封口” seg_text = "" seg_ts: Optional[float] = None async def flush_segment(send_flush_event: bool): """ - send_flush_event=True：告诉前端立刻结束当前 assistant 气泡（不结束整个 turn） - 若 seg_text 有内容：写入 history（用于刷新/回放） """ nonlocal seg_text, seg_ts if send_flush_event: if not await ws_send(ws, "assistant.flush", {}): return text = (seg_text or "").strip() if text: sess.history.append({ "id": uuid.uuid4().hex[:12], "role": "assistant", "content": text, "ts": seg_ts or time.time(), }) seg_text = "" seg_ts = None pump_task: Optional[asyncio.Task] = None # helper: 从 AIMessage 提取 tool_call_id（兼容不同 provider 的结构） def _tool_call_ids_from_ai_message(m: BaseMessage) -> set[str]: ids: set[str] = set() tc = getattr(m, "tool_calls", None) or [] for c in tc: _id = None if isinstance(c, dict): _id = c.get("id") or c.get("tool_call_id") else: _id = getattr(c, "id", None) or getattr(c, "tool_call_id", None) if _id: ids.add(str(_id)) ak = getattr(m, "additional_kwargs", None) or {} tc2 = ak.get("tool_calls") or [] for c in tc2: if isinstance(c, dict): _id = c.get("id") or c.get("tool_call_id") if _id: ids.add(str(_id)) return ids # helper: new_messages 里有哪些 tool_call_id def _tool_call_ids_in_msgs(msgs: List[BaseMessage]) -> set[str]: ids: set[str] = set() for m in msgs: if isinstance(m, AIMessage): ids |= _tool_call_ids_from_ai_message(m) return ids # helper: new_messages 里哪些 tool_call_id 已经有 ToolMessage 结果了 def _tool_result_ids_in_msgs(msgs: List[BaseMessage]) -> set[str]: ids: set[str] = set() for m in msgs: if isinstance(m, ToolMessage): tcid = getattr(m, "tool_call_id", None) if tcid: ids.add(str(tcid)) return ids # helper: 把“已存在的 ToolMessage”强制替换成 cancelled（避免工具其实返回了但用户打断没看到，导致上下文和 UI 不一致） def _force_cancelled_tool_results(msgs: List[BaseMessage], cancel_ids: set[str]) -> List[BaseMessage]: if not cancel_ids: return msgs cancelled_content = json.dumps({"cancelled": True}, ensure_ascii=False) out: List[BaseMessage] = [] for m in msgs: if isinstance(m, ToolMessage): tcid = getattr(m, "tool_call_id", None) if tcid and str(tcid) in cancel_ids: out.append(ToolMessage(content=cancelled_content, tool_call_id=str(tcid))) continue out.append(m) return out def _inject_cancelled_tool_messages(msgs: List[BaseMessage], tool_call_ids: List[str]) -> List[BaseMessage]: if not tool_call_ids: return msgs out = list(msgs) existing = set() for m in out: if isinstance(m, ToolMessage): tcid = getattr(m, "tool_call_id", None) if tcid: existing.add(str(tcid)) cancelled_content = json.dumps({"cancelled": True}, ensure_ascii=False) for tcid in tool_call_ids: tcid = str(tcid) if tcid in existing: continue insert_at = None for i in range(len(out) - 1, -1, -1): m = out[i] if isinstance(m, AIMessage) and (tcid in _tool_call_ids_from_ai_message(m)): insert_at = i + 1 break if insert_at is None: continue out.insert(insert_at, ToolMessage(content=cancelled_content, tool_call_id=tcid)) existing.add(tcid) return out def _sanitize_new_messages_on_cancel( new_messages: List[BaseMessage], *, interrupted_text: str, cancelled_tool_ids_from_ui: List[str], ) -> List[BaseMessage]: """ 返回：应该写回 sess.lc_messages 的消息序列（只包含“用户可见/认可”的那部分） - 工具：对未返回的 tool_call 补 ToolMessage({"cancelled": true}) - 回复：用 interrupted_text 替换末尾 final AIMessage，避免把完整回复泄漏进上下文 """ msgs = list(new_messages or []) interrupted_text = (interrupted_text or "").strip() # 1) 工具：找出“AI 发起了 tool_call 但没有 ToolMessage 结果”的那些 id ai_tool_ids = _tool_call_ids_in_msgs(msgs) tool_result_ids = _tool_result_ids_in_msgs(msgs) pending_tool_ids = ai_tool_ids - tool_result_ids # UI 认为被取消的 tool（running -> cancelled） ui_cancel_ids = {str(x) for x in (cancelled_tool_ids_from_ui or [])} # 统一要取消的集合： # - UI 侧 running 的（用户按下打断时看见的） # - 以及 messages 里缺结果的（防止漏标） cancel_ids = set(ui_cancel_ids) | set(pending_tool_ids) # 2) 如果 new_messages 里已经有 ToolMessage(真实结果) 但用户打断了， # 为了“UI/上下文一致”，强制替换成 cancelled msgs = _force_cancelled_tool_results(msgs, cancel_ids) # 3) 注入缺失的 ToolMessage(cancelled) msgs = _inject_cancelled_tool_messages(msgs, list(cancel_ids)) # 4) 处理 assistant 最终文本（避免把完整 answer 写回） # - 如果 interrupted_text 非空：用它替换最后一个“非 tool_call 的 AIMessage” # - 如果 interrupted_text 为空：只在“末尾存在一个 non-toolcall AIMessage（且它后面没有 tool_call）”时移除它 def _is_toolcall_ai(m: BaseMessage) -> bool: return isinstance(m, AIMessage) and bool(_tool_call_ids_from_ai_message(m)) def _is_text_ai(m: BaseMessage) -> bool: if not isinstance(m, AIMessage): return False if _tool_call_ids_from_ai_message(m): return False c = getattr(m, "content", None) return isinstance(c, str) and bool(c.strip()) # 找最后一个“文本 AIMessage（非 tool_call）” last_text_ai_idx = None for i in range(len(msgs) - 1, -1, -1): if _is_text_ai(msgs[i]): last_text_ai_idx = i break if interrupted_text: if last_text_ai_idx is None: msgs.append(AIMessage(content=interrupted_text)) else: # 用用户看见的部分替换，且丢弃后面所有消息（防止泄漏） msgs = msgs[:last_text_ai_idx] + [AIMessage(content=interrupted_text)] return msgs # interrupted_text 为空：用户没看见任何本段 token # 只移除“末尾的 final answer AIMessage”，避免把 unseen answer 写进上下文； # 但如果该 AIMessage 后面还有 tool_call（说明它是 pre-tool 文本），就不要删 if last_text_ai_idx is not None: has_toolcall_after = any(_is_toolcall_ai(m) for m in msgs[last_text_ai_idx + 1 :]) if not has_toolcall_after: msgs = msgs[:last_text_ai_idx] return msgs pump_task: Optional[asyncio.Task] = None cancel_wait_task: Optional[asyncio.Task] = None was_interrupted = False # 本 turn 是否已经走了“打断收尾” try: async with mcp_sink_context(sink): pump_task = asyncio.create_task(pump_agent()) cancel_wait_task = asyncio.create_task(sess.cancel_event.wait()) while True: # 同时等：queue 出事件 或 cancel_event get_task = asyncio.create_task(out_q.get()) done, _ = await asyncio.wait( {get_task, cancel_wait_task}, return_when=asyncio.FIRST_COMPLETED, ) # 优先处理队列事件（避免 done/flush 已经在队列里时被 cancel 抢占） if get_task in done: kind, payload = get_task.result() else: # cancel_event 触发：不再等 queue try: get_task.cancel() await get_task except asyncio.CancelledError: debug_traceback_print(app.state.cfg) pass except Exception: debug_traceback_print(app.state.cfg) pass kind, payload = ("agent.cancelled", None) # ------------------------ # 1) 处理打断 # ------------------------ if kind == "agent.cancelled": # 防止重复触发（cancel_event + pump_agent cancelled 都可能来一次） if was_interrupted: break was_interrupted = True # 1.1 cancel agent 流（停止继续产出 token/工具） if pump_task and (not pump_task.done()): pump_task.cancel() # 1.2 将所有 running 的工具卡片标记为 error cancelled_tool_recs: List[Dict[str, Any]] = [] for tcid, idx in list(sess._tool_history_index.items()): rec = sess.history[idx] if rec.get("role") == "tool" and rec.get("state") == "running": rec.update({ "state": "error", "progress": 1.0, "message": "Cancelled by user", "summary": {"cancelled": True}, }) cancelled_tool_recs.append(rec) # 推送 tool.end，确保前端停止 spinner for rec in cancelled_tool_recs: await ws_send(ws, "tool.end", { "tool_call_id": rec["tool_call_id"], "server": rec["server"], "name": rec["name"], "is_error": True, "summary": rec.get("summary"), }) # 1.3 把已输出的 seg_text 写入 history（UI 看到的内容） interrupted_text = (seg_text or "").strip() if interrupted_text: sess.history.append({ "id": uuid.uuid4().hex[:12], "role": "assistant", "content": interrupted_text, "ts": seg_ts or time.time(), }) # 1.4 上下文：只写回“用户真实看到/认可”的消息序列 cancelled_tool_ids = [rec["tool_call_id"] for rec in cancelled_tool_recs] commit_msgs = _sanitize_new_messages_on_cancel( new_messages, interrupted_text=interrupted_text, cancelled_tool_ids_from_ui=cancelled_tool_ids, ) if commit_msgs: sess.lc_messages.extend(commit_msgs) elif interrupted_text: # 极端情况：updates 没来得及给任何消息，但用户已看到 token sess.lc_messages.append(AIMessage(content=interrupted_text)) # ★打断：只发 assistant.end，带 interrupted=true await ws_send(ws, "assistant.end", {"text": interrupted_text, "interrupted": True}) sess.cancel_event.clear() break # ------------------------ # 2) 事件处理 # ------------------------ if kind == "assistant.delta": delta = payload or "" if delta: if seg_ts is None: seg_ts = time.time() seg_text += delta if not await ws_send(ws, "assistant.delta", {"delta": delta}): raise WebSocketDisconnect() continue if kind == "mcp": raw = payload if raw.get("type") == "tool_start": await flush_segment(send_flush_event=True) rec = sess.apply_tool_event(raw) if rec: if raw["type"] == "tool_start": await ws_send(ws, "tool.start", { "tool_call_id": rec["tool_call_id"], "server": rec["server"], "name": rec["name"], "args": rec["args"], }) elif raw["type"] == "tool_progress": await ws_send(ws, "tool.progress", { "tool_call_id": rec["tool_call_id"], "server": rec["server"], "name": rec["name"], "progress": rec["progress"], "message": rec["message"], }) elif raw["type"] == "tool_end": await ws_send(ws, "tool.end", { "tool_call_id": rec["tool_call_id"], "server": rec["server"], "name": rec["name"], "is_error": rec["state"] == "error", "summary": rec["summary"], }) continue if kind == "agent.done": final_text = (seg_text or "").strip() if final_text: sess.history.append({ "id": uuid.uuid4().hex[:12], "role": "assistant", "content": final_text, "ts": seg_ts or time.time(), }) if new_messages: sess.lc_messages.extend(new_messages) if not await ws_send(ws, "assistant.end", {"text": final_text}): return break if kind == "agent.error": err_text = str(payload or "unknown error") partial = (seg_text or "").strip() # 把已输出部分落盘/落上下文（避免丢上下文） if partial: sess.history.append({ "id": uuid.uuid4().hex[:12], "role": "assistant", "content": partial, "ts": seg_ts or time.time(), }) sess.lc_messages.append(AIMessage(content=partial)) if new_messages: sess.lc_messages.extend(new_messages) # ★ 真异常：只发 error（并带 partial_text 让前端结束当前气泡） await ws_send(ws, "error", {"message": err_text, "partial_text": partial}) break except WebSocketDisconnect: return except asyncio.CancelledError: # 连接关闭/任务取消：不当作 error return except Exception as e: # 如果已经走了打断收尾，别再发 error（避免“打断=报错”） if was_interrupted: return await ws_send(ws, "error", {"message": f"{type(e).__name__}: {e}", "partial_text": (seg_text or "").strip()}) return finally: # 结束 cancel_wait_task if cancel_wait_task and (not cancel_wait_task.done()): cancel_wait_task.cancel() # pump_task 取消/收尾：避免 await 卡死，加一个短超时保护 if pump_task and (not pump_task.done()): pump_task.cancel() if pump_task: try: await asyncio.wait_for(pump_task, timeout=2.0) except asyncio.TimeoutError: debug_traceback_print(app.state.cfg) pass except asyncio.CancelledError: debug_traceback_print(app.state.cfg) pass except Exception: debug_traceback_print(app.state.cfg) pass finally: try: CHAT_TURN_SEM.release() except Exception: debug_traceback_print(app.state.cfg) pass except WebSocketDisconnect: return finally: try: WS_CONN_SEM.release() except: pass