"""QQ Bot chunked upload flow. The QQ v2 API caps inline base64 uploads (``file_data`` / ``url``) at ~10 MB. Files between 10 MB and ~100 MB use the three-step chunked flow:: 1. POST /v2/{users|groups}/{id}/upload_prepare → upload_id, block_size, and pre-signed COS part URLs. 2. Per part: PUT bytes to the COS URL, then POST .../upload_part_finish. 3. POST /v2/{users|groups}/{id}/files with {"upload_id": ...} → ``file_info`` token used in a RichMedia message. Error codes (QQ Bot v2 spec): ``40093001`` — ``upload_part_finish`` retryable until the server's ``retry_timeout`` (or a local cap) elapses; ``40093002`` — daily upload quota exceeded, surfaced as :class:`UploadDailyLimitExceededError`. Other API/I/O failures raise ``RuntimeError``. Ported from WideLee's qqbot-agent-sdk v1.2.2 (``media_loader.py::ChunkedUploader``). Authorship preserved via Co-authored-by. """ from __future__ import annotations import asyncio import hashlib import logging from dataclasses import dataclass from pathlib import Path from typing import Any, Awaitable, Callable, Dict, List from gateway.platforms.qqbot.constants import FILE_UPLOAD_TIMEOUT logger = logging.getLogger(__name__) _BIZ_CODE_DAILY_LIMIT = 40093002 # upload_prepare: daily cumulative limit _BIZ_CODE_PART_RETRYABLE = 40093001 # upload_part_finish: transient _DEFAULT_CONCURRENT_PARTS = 1 _MAX_CONCURRENT_PARTS = 10 _PART_UPLOAD_TIMEOUT = 300.0 # 5 minutes per COS PUT _PART_UPLOAD_MAX_RETRIES = 2 _PART_FINISH_RETRY_INTERVAL = 1.0 _PART_FINISH_DEFAULT_TIMEOUT = 120.0 _PART_FINISH_MAX_TIMEOUT = 600.0 _COMPLETE_UPLOAD_MAX_RETRIES = 2 _COMPLETE_UPLOAD_BASE_DELAY = 2.0 # First 10,002,432 bytes used for the ``md5_10m`` hash (per QQ API spec). _MD5_10M_SIZE = 10_002_432 # ── Exceptions ─────────────────────────────────────────────────────── class UploadDailyLimitExceededError(Exception): """Raised when ``upload_prepare`` returns biz_code 40093002 (daily quota hit).""" def __init__(self, file_name: str, file_size: int, message: str = "") -> None: self.file_name = file_name self.file_size = file_size super().__init__(message or f"Daily upload limit exceeded for {file_name!r}") @property def file_size_human(self) -> str: return format_size(self.file_size) class UploadFileTooLargeError(Exception): """Raised when a file exceeds the platform per-file size limit.""" def __init__(self, file_name: str, file_size: int, limit_bytes: int = 0, message: str = "") -> None: self.file_name = file_name self.file_size = file_size self.limit_bytes = limit_bytes limit_str = f" ({format_size(limit_bytes)})" if limit_bytes else "" super().__init__( message or f"File {file_name!r} ({format_size(file_size)}) exceeds platform limit{limit_str}" ) @property def file_size_human(self) -> str: return format_size(self.file_size) @property def limit_human(self) -> str: return format_size(self.limit_bytes) if self.limit_bytes else "unknown" # ── Prepare-response shape ─────────────────────────────────────────── @dataclass class _PreparePart: index: int presigned_url: str block_size: int = 0 @dataclass class _PrepareResult: upload_id: str block_size: int parts: List[_PreparePart] concurrency: int = _DEFAULT_CONCURRENT_PARTS retry_timeout: float = 0.0 def _parse_prepare_response(raw: Dict[str, Any]) -> _PrepareResult: """Parse upload_prepare response (either bare or wrapped in ``data``).""" src = raw.get("data") if isinstance(raw.get("data"), dict) else raw upload_id = str(src.get("upload_id", "")) if not upload_id: raise ValueError(f"upload_prepare response missing upload_id: {str(raw)[:200]}") block_size = int(src.get("block_size", 0)) raw_parts = src.get("parts") or src.get("part_list") or [] if not isinstance(raw_parts, list) or not raw_parts: raise ValueError(f"upload_prepare response missing parts: {str(raw)[:200]}") parts = [ _PreparePart( index=int(p.get("part_index") or p.get("index") or 0), presigned_url=str(p.get("presigned_url") or p.get("url") or ""), block_size=int(p.get("block_size", 0)), ) for p in raw_parts if isinstance(p, dict) ] return _PrepareResult( upload_id=upload_id, block_size=block_size, parts=parts, concurrency=int(src.get("concurrency", _DEFAULT_CONCURRENT_PARTS)) or _DEFAULT_CONCURRENT_PARTS, retry_timeout=float(src.get("retry_timeout", 0.0) or 0.0), ) def _api_path(chat_type: str, target_id: str, endpoint: str) -> str: base = "/v2/users" if chat_type == "c2c" else "/v2/groups" return f"{base}/{target_id}/{endpoint}" # ── Chunked upload driver ──────────────────────────────────────────── class ChunkedUploader: """Run the prepare → PUT parts → complete sequence. :param api_request: Bound ``_api_request(method, path, body=..., timeout=...)`` coroutine from the adapter (passed in rather than imported to avoid circular imports). Must raise ``RuntimeError`` with the biz_code in the message on API errors. :param http_put: Coroutine ``(url, data, headers) -> httpx-like response`` for COS part uploads. :param log_tag: Log prefix. """ def __init__( self, api_request: Callable[..., Awaitable[Dict[str, Any]]], http_put: Callable[..., Awaitable[Any]], log_tag: str = "QQBot", ) -> None: self._api_request = api_request self._http_put = http_put self._log_tag = log_tag async def upload( self, chat_type: str, target_id: str, file_path: str, file_type: int, file_name: str, ) -> Dict[str, Any]: """Run the full chunked upload and return the ``complete_upload`` response. :param chat_type: ``'c2c'`` or ``'group'``. :param file_type: ``MEDIA_TYPE_*`` constant. :returns: Raw ``complete_upload`` response dict (contains ``file_info``). :raises UploadDailyLimitExceededError: On biz_code 40093002. :raises UploadFileTooLargeError: When the file exceeds the platform limit. :raises RuntimeError: On other API or I/O failures. """ if chat_type not in {"c2c", "group"}: raise ValueError(f"ChunkedUploader: unsupported chat_type {chat_type!r}") file_size = Path(file_path).stat().st_size logger.info( "[%s] Chunked upload start: file=%s size=%s type=%d", self._log_tag, file_name, format_size(file_size), file_type, ) # Hashing is blocking I/O → executor. hashes = await asyncio.get_running_loop().run_in_executor( None, _compute_file_hashes, file_path, file_size ) prepare = await self._prepare(chat_type, target_id, file_type, file_name, file_size, hashes) max_concurrent = min(prepare.concurrency, _MAX_CONCURRENT_PARTS) retry_timeout = min( prepare.retry_timeout if prepare.retry_timeout > 0 else _PART_FINISH_DEFAULT_TIMEOUT, _PART_FINISH_MAX_TIMEOUT, ) logger.info( "[%s] Prepared: upload_id=%s block_size=%s parts=%d concurrency=%d", self._log_tag, prepare.upload_id, format_size(prepare.block_size), len(prepare.parts), max_concurrent, ) total_parts = len(prepare.parts) completed = [0] # shared counter for progress logging sem = asyncio.Semaphore(max(max_concurrent, 1)) async def _run(part: _PreparePart) -> None: async with sem: await self._upload_one_part( chat_type, target_id, file_path, file_size, prepare.upload_id, prepare.block_size, part, retry_timeout, total_parts, completed, ) await asyncio.gather(*(_run(p) for p in prepare.parts)) logger.info("[%s] All %d parts uploaded, completing…", self._log_tag, total_parts) return await self._complete(chat_type, target_id, prepare.upload_id) async def _prepare( self, chat_type: str, target_id: str, file_type: int, file_name: str, file_size: int, hashes: Dict[str, str], ) -> _PrepareResult: body = { "file_type": file_type, "file_name": file_name, "file_size": file_size, "md5": hashes["md5"], "sha1": hashes["sha1"], "md5_10m": hashes["md5_10m"], } try: raw = await self._api_request( "POST", _api_path(chat_type, target_id, "upload_prepare"), body=body, timeout=FILE_UPLOAD_TIMEOUT, ) except RuntimeError as exc: err_msg = str(exc) if f"{_BIZ_CODE_DAILY_LIMIT}" in err_msg: raise UploadDailyLimitExceededError(file_name, file_size, err_msg) from exc raise return _parse_prepare_response(raw) async def _upload_one_part( self, chat_type: str, target_id: str, file_path: str, file_size: int, upload_id: str, rsp_block_size: int, part: _PreparePart, retry_timeout: float, total_parts: int, completed: List[int], ) -> None: """PUT one part to COS, then call ``upload_part_finish``.""" part_index = part.index # Per-part block_size wins; fall back to the response-level value. actual_block_size = part.block_size if part.block_size > 0 else rsp_block_size offset = (part_index - 1) * rsp_block_size length = min(actual_block_size, file_size - offset) data = await asyncio.get_running_loop().run_in_executor( None, _read_file_chunk, file_path, offset, length ) md5_hex = hashlib.md5(data).hexdigest() logger.debug( "[%s] Part %d/%d: uploading %s (offset=%d md5=%s)", self._log_tag, part_index, total_parts, format_size(length), offset, md5_hex, ) await self._put_to_presigned_url(part.presigned_url, data, part_index, total_parts) await self._part_finish_with_retry( chat_type, target_id, upload_id, part_index, length, md5_hex, retry_timeout, ) completed[0] += 1 logger.debug( "[%s] Part %d/%d done (%d/%d total)", self._log_tag, part_index, total_parts, completed[0], total_parts, ) async def _with_retries( self, attempt_fn: Callable[[], Awaitable[Any]], *, max_retries: int, base_delay: float, label: str, failure_label: str, ) -> Any: """Run *attempt_fn* up to ``max_retries + 1`` times with exponential backoff.""" last_exc: Exception | None = None for attempt in range(max_retries + 1): try: return await attempt_fn() except Exception as exc: last_exc = exc if attempt < max_retries: delay = base_delay * (2 ** attempt) logger.warning( "[%s] %s attempt %d failed, retry in %.1fs: %s", self._log_tag, label, attempt + 1, delay, exc, ) await asyncio.sleep(delay) raise RuntimeError(f"{failure_label} failed after {max_retries + 1} attempts: {last_exc}") async def _put_to_presigned_url(self, url: str, data: bytes, part_index: int, total_parts: int) -> None: """PUT part data to a pre-signed COS URL with retry.""" async def _attempt() -> None: resp = await asyncio.wait_for( self._http_put(url, data=data, headers={"Content-Length": str(len(data))}), timeout=_PART_UPLOAD_TIMEOUT, ) status = getattr(resp, "status_code", 0) if 200 <= status < 300: logger.debug("[%s] PUT part %d/%d: %d OK", self._log_tag, part_index, total_parts, status) return body_preview = "" try: body_preview = getattr(resp, "text", "")[:200] except Exception: # pragma: no cover — defensive pass raise RuntimeError(f"COS PUT returned {status}: {body_preview}") await self._with_retries( _attempt, max_retries=_PART_UPLOAD_MAX_RETRIES, base_delay=1.0, label=f"PUT part {part_index}/{total_parts}", failure_label=f"Part {part_index}/{total_parts} upload", ) async def _part_finish_with_retry( self, chat_type: str, target_id: str, upload_id: str, part_index: int, block_size: int, md5: str, retry_timeout: float, ) -> None: """Call ``upload_part_finish``, retrying on biz_code 40093001 until *retry_timeout*.""" path = _api_path(chat_type, target_id, "upload_part_finish") body = {"upload_id": upload_id, "part_index": part_index, "block_size": block_size, "md5": md5} loop = asyncio.get_running_loop() start = loop.time() attempt = 0 while True: try: await self._api_request("POST", path, body=body, timeout=FILE_UPLOAD_TIMEOUT) return except RuntimeError as exc: if f"{_BIZ_CODE_PART_RETRYABLE}" not in str(exc): raise elapsed = loop.time() - start if elapsed >= retry_timeout: raise RuntimeError( f"upload_part_finish persistent retry timed out " f"after {retry_timeout:.0f}s ({attempt} retries): {exc}" ) from exc attempt += 1 logger.debug( "[%s] part_finish retryable error, attempt %d, elapsed=%.1fs: %s", self._log_tag, attempt, elapsed, exc, ) await asyncio.sleep(_PART_FINISH_RETRY_INTERVAL) async def _complete(self, chat_type: str, target_id: str, upload_id: str) -> Dict[str, Any]: """Call ``complete_upload`` with retry. Reuses the ``/files`` endpoint (same as the simple URL-based upload) but signals the chunked-completion path by sending only ``upload_id``. """ path = _api_path(chat_type, target_id, "files") body = {"upload_id": upload_id} return await self._with_retries( lambda: self._api_request("POST", path, body=body, timeout=FILE_UPLOAD_TIMEOUT), max_retries=_COMPLETE_UPLOAD_MAX_RETRIES, base_delay=_COMPLETE_UPLOAD_BASE_DELAY, label="complete_upload", failure_label="complete_upload", ) # ── Helpers (module-level for testability) ─────────────────────────── def format_size(size_bytes: int) -> str: """Return a human-readable file size string (e.g. ``'12.3 MB'``).""" size = float(size_bytes) for unit in ("B", "KB", "MB", "GB"): if size < 1024.0: return f"{size:.1f} {unit}" size /= 1024.0 return f"{size:.1f} TB" def _read_file_chunk(file_path: str, offset: int, length: int) -> bytes: """Read *length* bytes at *offset*; raises IOError on a short read (truncated file).""" with open(file_path, "rb") as fh: fh.seek(offset) data = fh.read(length) if len(data) != length: raise IOError( f"Short read from {file_path}: expected {length} bytes at " f"offset {offset}, got {len(data)} (file may be truncated)" ) return data def _compute_file_hashes(file_path: str, file_size: int) -> Dict[str, str]: """Compute md5, sha1, and md5_10m in a single pass.""" md5 = hashlib.md5() sha1 = hashlib.sha1() md5_10m = hashlib.md5() need_10m = file_size > _MD5_10M_SIZE bytes_read = 0 with open(file_path, "rb") as fh: while chunk := fh.read(65536): md5.update(chunk) sha1.update(chunk) if need_10m: remaining = _MD5_10M_SIZE - bytes_read if remaining > 0: md5_10m.update(chunk[:remaining]) bytes_read += len(chunk) full_md5 = md5.hexdigest() return { "md5": full_md5, "sha1": sha1.hexdigest(), # For small files the "10m" hash is just the full md5. "md5_10m": md5_10m.hexdigest() if need_10m else full_md5, }