robodock-cli 0.1.0__py3-none-any.whl

sdk/client/auth.py

@@ -0,0 +1,171 @@
+import time
+from typing import Any, Callable
+import json
+
+from fast_api_client import AuthenticatedClient, Client
+from fast_api_client.api.auth import (
+    login_api_p_v1_auth_login_post,
+    logout_api_p_v1_auth_logout_post,
+    refresh_token_api_p_v1_auth_refresh_post,
+    me_api_p_v1_auth_me_get
+)
+from fast_api_client.models.login_request import LoginRequest
+from fast_api_client.models.logout_request import LogoutRequest
+from fast_api_client.models.refresh_request import RefreshRequest
+from config import (
+    ROBODOCK_BASEURL,
+    SESSION_FILE,
+    REFRESH_AHEAD_SECONDS,
+)
+from sdk.exception import (
+    LoginError,
+    LogoutError,
+    NotLoggedInError,
+    TokenRefreshError,
+    UserInfoError,
+)
+    
+
+class AuthAPI:
+    def __init__(self, client):
+        self.client = client
+        self.token = None
+        self.refresh_token = None
+        self.expires_in = None
+        self.expires_at = None
+        self.load_state()
+
+    # ---------- token 管理 ----------
+
+    def _apply_tokens(self, data: dict) -> None:
+        """保存最新的 token pair，并重建带认证的 client。"""
+        self.token = data["access_token"]
+        self.refresh_token = data["refresh_token"]
+        self.expires_in = data["expires_in"]
+        self.expires_at = time.time() + self.expires_in
+        self._refresh_client()
+
+    def _refresh_client(self) -> None:
+        """用当前 token 重建 AuthenticatedClient。"""
+        self.client.client = AuthenticatedClient(
+            base_url=ROBODOCK_BASEURL,
+            token=self.token,
+        )
+
+    def _is_token_expired(self) -> bool:
+        """判断 access token 是否已过期（或即将过期）。"""
+        if self.expires_at is None:
+            return True
+        return time.time() >= self.expires_at - REFRESH_AHEAD_SECONDS
+
+    def _ensure_token(self) -> None:
+        """如果 token 已过期则自动刷新。"""
+        if self.client.client is None:
+            raise NotLoggedInError()
+        if self._is_token_expired():
+            self.refresh()
+
+    # ---------- 后端认证接口 ----------
+
+    def login(self, username: str, password: str):
+        client = Client(base_url=ROBODOCK_BASEURL)
+        req = LoginRequest(login=username, password=password)
+        try:
+            resp = login_api_p_v1_auth_login_post.sync(client=client, body=req)
+        except Exception as e:
+            raise LoginError() from e
+        if resp is None or resp.get("code") != 200:
+            raise LoginError(str(resp))
+        self._apply_tokens(resp["data"])
+
+    def logout(self):
+        if self.client.client is None:
+            raise NotLoggedInError()
+        try:
+            req = LogoutRequest(refresh_token=self.refresh_token)
+            logout_api_p_v1_auth_logout_post.sync(
+                client=self.client.client,
+                body=req,
+            )
+        except Exception as e:
+            raise LogoutError() from e
+        self.client.client = None
+        self.token = None
+        self.refresh_token = None
+        self.expires_in = None
+        self.expires_at = None
+
+    def refresh(self):
+        """用 refresh token 刷新 access token（滚动刷新）。"""
+        if self.refresh_token is None:
+            raise NotLoggedInError("未登录")
+        # 用普通 Client 调刷新接口，避免携带可能已过期的 Authorization 头
+        client = Client(base_url=ROBODOCK_BASEURL)
+        req = RefreshRequest(refresh_token=self.refresh_token)
+        try:
+            resp = refresh_token_api_p_v1_auth_refresh_post.sync(
+                client=client,
+                body=req,
+            )
+        except Exception as e:
+            raise TokenRefreshError() from e
+        if resp is None or resp.get("code") != 200:
+            raise TokenRefreshError(str(resp))
+        self._apply_tokens(resp["data"])
+        
+    def info(self):
+        """获取当前登录用户信息。"""
+        self._ensure_token()
+        try:
+            resp = me_api_p_v1_auth_me_get.sync_detailed(client=self.client.client)
+        except Exception as e:
+            raise UserInfoError() from e
+        if resp is None or resp.status_code != 200:
+            raise UserInfoError(str(resp))
+        resp = json.loads(resp.content.decode("utf-8"))
+        if resp.get("code") != 200:
+            raise UserInfoError(str(resp))
+        return resp["data"]
+
+    # ---------- session文件管理 ----------
+    
+    def load_state(self) -> None:
+        """从持久化的Session文件中恢复登录态。"""
+        try:
+            state = json.loads(SESSION_FILE.read_text(encoding="utf-8"))
+        except (json.JSONDecodeError, OSError):
+            self.client.client = None
+            return 
+        
+        self.token = state.get("token")
+        self.refresh_token = state.get("refresh_token")
+        self.expires_in = state.get("expires_in")
+        self.expires_at = state.get("expires_at")
+        if self.token:
+            self._refresh_client()
+        else:
+            self.client.client = None
+
+    def save_state(self) -> None:
+        """导出登录态，并持久化到本地文件。"""
+        state_info = {
+            "token": self.token,
+            "refresh_token": self.refresh_token,
+            "expires_in": self.expires_in,
+            "expires_at": self.expires_at,
+        }
+        SESSION_FILE.parent.mkdir(parents=True, exist_ok=True)
+        SESSION_FILE.write_text(
+            json.dumps(state_info, ensure_ascii=False, indent=2),
+            encoding="utf-8",
+        )
+        
+    def clear_state(self) -> None:
+        """清除登录态，并删除本地持久化文件。"""
+        self.token = None
+        self.refresh_token = None
+        self.expires_in = None
+        self.expires_at = None
+        self.client.client = None
+        if SESSION_FILE.exists():
+            SESSION_FILE.unlink()

sdk/client/client.py

@@ -0,0 +1,9 @@
+from sdk.client.auth import AuthAPI
+from sdk.client.task_management import TaskManagementAPI
+from fast_api_client.client import AuthenticatedClient
+
+class RobodockClient:
+    def __init__(self):
+        self.client : AuthenticatedClient = None
+        self.auth = AuthAPI(self)
+        self.task_management = TaskManagementAPI(self)

sdk/client/task_management.py

@@ -0,0 +1,29 @@
+from fast_api_client.api.task_management import get_checkpoint_download_url_api_p_v1_run_checkpoint_download_post
+from fast_api_client.models.checkpoint_download_request import CheckpointDownloadRequest
+from sdk.exception import (
+    GetCheckpointDownloadUrlError
+)
+class TaskManagementAPI:
+    def __init__(self, client):
+        self.client = client
+
+    def get_checkpoint_download_url(self, run_uuid:str, checkpoint_path):
+        
+        req = CheckpointDownloadRequest(
+            run_uuid=run_uuid,
+            checkpoint_path=checkpoint_path
+        )
+        
+        self.client.auth._ensure_token()
+        try:
+            resp = get_checkpoint_download_url_api_p_v1_run_checkpoint_download_post.sync(
+                client=self.client.client,
+                body=req,
+            )
+        except Exception as e:
+            raise GetCheckpointDownloadUrlError() from e
+
+        if resp is None or resp.get("code") != 200:
+            raise GetCheckpointDownloadUrlError(str(resp))
+
+        return resp.get("data")

sdk/exception.py

@@ -0,0 +1,95 @@
+"""Robodock SDK 统一异常定义。
+
+所有业务异常都继承自 :class:`RobodockError`，上层代码只需捕获它即可统一处理。
+"""
+
+
+class RobodockError(Exception):
+    """Robodock SDK 基础异常，所有业务异常都继承自此。"""
+    def __init__(self, message: str = ""):
+        default_msg = "Robodock SDK异常"
+        if message:
+            full_msg = f"{default_msg}: {message}"
+        else:
+            full_msg = default_msg
+        super().__init__(full_msg)
+
+class NotLoggedInError(RobodockError):
+    """未登录异常。"""
+    def __init__(self, message: str = ""):
+        default_msg = "用户未登录，请先登录"
+        if message:
+            full_msg = f"{default_msg}: {message}"
+        else:
+            full_msg = default_msg
+        super().__init__(full_msg)
+
+
+class LoginError(RobodockError):
+    """登录失败异常。"""
+    def __init__(self, message: str = ""):
+        default_msg = "登录失败"
+        if message:
+            full_msg = f"{default_msg}: {message}"
+        else:
+            full_msg = default_msg
+        super().__init__(full_msg)
+
+class LogoutError(RobodockError):
+    """退出登录失败异常。"""
+    def __init__(self, message: str = ""):
+        default_msg = "退出登录失败"
+        if message:
+            full_msg = f"{default_msg}: {message}"
+        else:
+            full_msg = default_msg
+        super().__init__(full_msg)
+    
+
+class TokenRefreshError(RobodockError):
+    """token 刷新失败异常。"""
+    def __init__(self, message: str = ""):
+        default_msg = "token刷新失败"
+        if message:
+            full_msg = f"{default_msg}: {message}"
+        else:
+            full_msg = default_msg
+        super().__init__(full_msg)
+class UserInfoError(RobodockError):
+    """获取用户信息失败异常。"""
+    def __init__(self, message: str = ""):
+        default_msg = "获取用户信息失败"
+        if message:
+            full_msg = f"{default_msg}: {message}"
+        else:
+            full_msg = default_msg
+        super().__init__(full_msg)
+
+class GetCheckpointDownloadUrlError(RobodockError):
+    """获取模型下载链接失败"""
+    def __init__(self, message: str = ""):
+        default_msg = "获取模型下载链接失败"
+        if message:
+            full_msg = f"{default_msg}: {message}"
+        else:
+            full_msg = default_msg
+        super().__init__(full_msg)
+
+class APIRequestError(RobodockError):
+    """API 请求失败异常，携带 HTTP 状态码。"""
+
+    def __init__(self, status_code: int, message: str | None = None):
+        self.status_code = status_code
+        if message is None:
+            message = f"请求失败: HTTP {status_code}"
+        super().__init__(message)
+
+class DownloadError(RobodockError):
+    """基于 presigned URL 下载文件失败。"""
+    def __init__(self, message: str = ""):
+            default_msg = "基于 presigned URL 下载文件失败"
+            if message:
+                full_msg = f"{default_msg}: {message}"
+            else:
+                full_msg = default_msg
+            super().__init__(full_msg)

sdk/utils/__init__.py

@@ -0,0 +1 @@
+"""Robodock SDK 工具模块。"""

sdk/utils/downloader.py

@@ -0,0 +1,267 @@
+"""基于 presigned URL 的文件下载工具，支持断点续传。
+
+典型用法（配合 SDK 获取 presigned 下载链接）::
+
+    from pathlib import Path
+
+    from sdk.utils.downloader import download_file
+
+    resp = client.task_management.get_checkpoint_download_url(
+        run_uuid="xxx", checkpoint_path="model.ckpt",
+    )
+    data = resp["data"]
+    url = data["url"] if isinstance(data, dict) else data
+
+    download_file(
+        url,
+        Path("model.ckpt"),
+        on_progress=lambda done, total: print(f"{done}/{total}"),
+    )
+
+断点续传原理
+------------
+1. 本地文件已存在时，以「已有文件大小」作为已下载字节数，通过
+   ``Range: bytes=<start>-`` 请求剩余部分；
+2. 服务端返回 ``206 Partial Content`` 表示支持 Range，从偏移量继续追加写入；
+3. 服务端返回 ``200`` 表示忽略了 Range（或不支持），则从头重新下载；
+4. 服务端返回 ``416 Range Not Satisfiable`` 表示本地文件已完整，直接完成。
+
+另外针对 presigned URL 的特点做了两点增强：
+
+- 下载中断（网络抖动、超时、5xx）后自动重试，并从本地已写入的字节数继续；
+- presigned URL 过期（401/403）时，可通过 ``get_presigned_url`` 回调重新获取
+  链接后继续下载。
+"""
+
+from __future__ import annotations
+
+import random
+import time
+from pathlib import Path
+from typing import Callable, Optional, Union
+
+import httpx
+
+from sdk.exception import RobodockError, DownloadError
+
+__all__ = ["download_file", "DownloadError"]
+
+class _RestartNeeded(Exception):
+    """服务端忽略 Range 或返回的偏移与本地不一致，需要从头重新下载。"""
+
+
+class _HttpDownloadError(Exception):
+    """HTTP 状态码非预期。"""
+
+    def __init__(self, status_code: int, message: str, *, refresh_url: bool = False):
+        super().__init__(message)
+        self.status_code = status_code
+        self.refresh_url = refresh_url
+
+
+def download_file(
+    presigned_url: str,
+    local_path: Union[str, Path],
+    *,
+    chunk_size: int = 1024 * 1024,
+    timeout: Union[float, httpx.Timeout, None] = None,
+    overwrite: bool = False,
+    max_retries: int = 5,
+    retry_backoff: float = 1.0,
+    verify: bool = True,
+    get_presigned_url: Optional[Callable[[], str]] = None,
+    on_progress: Optional[Callable[[int, Optional[int]], None]] = None,
+) -> Path:
+    """根据 presigned URL 下载文件到本地，支持断点续传。
+
+    Args:
+        presigned_url: 后端返回的 presigned 下载地址。
+        local_path: 本地保存路径。文件已存在时会从已有大小处断点续传。
+        chunk_size: 每次从网络读取并写入磁盘的字节数，默认 1 MiB。
+        timeout: 单次 HTTP 请求的超时时间。
+        overwrite: 为 True 时忽略本地已有文件，强制从头下载。
+        max_retries: 遇到网络错误/5xx/链接失效时的最大重试次数（不含首次）。
+        retry_backoff: 重试退避基数（秒），实际等待
+            ``backoff * 2 ** 已重试次数 + 随机抖动``。
+        verify: 是否校验 TLS 证书。
+        get_presigned_url: 可选的链接刷新回调。presigned URL 过期（HTTP 401/403）
+            且仍可重试时会被调用，用返回值替换下载链接后继续。
+            未提供时链接失效会直接抛出 :class:`DownloadError`。
+        on_progress: 进度回调，参数为 ``(downloaded_bytes, total_bytes)``；
+            total 无法从响应头解析时为 ``None``。
+
+    Returns:
+        下载完成的本地文件路径。
+
+    Raises:
+        DownloadError: 下载失败（含重试耗尽、链接失效且无法刷新等）。
+    """
+    local_path = Path(local_path)
+    url = presigned_url
+    attempts = 0
+
+    while True:
+        attempts += 1
+        try:
+            _download_once(
+                url=url,
+                local_path=local_path,
+                chunk_size=chunk_size,
+                timeout=timeout,
+                overwrite=overwrite,
+                verify=verify,
+                on_progress=on_progress,
+            )
+            return local_path
+        except _RestartNeeded:
+            # 服务端不支持 Range 或偏移不一致：强制从头下载一次
+            overwrite = True
+            continue
+        except (httpx.TransportError, _HttpDownloadError) as exc:
+            if isinstance(exc, _HttpDownloadError) and exc.status_code == 404:
+                raise DownloadError("下载失败: 文件不存在 (HTTP 404)") from exc
+            if isinstance(exc, _HttpDownloadError) and exc.status_code in (401, 403):
+                if get_presigned_url is None:
+                    raise DownloadError(
+                        f"下载失败: presigned URL 已失效 (HTTP {exc.status_code})，"
+                        "且未提供 get_presigned_url 回调，无法刷新链接"
+                    ) from exc
+                if attempts > max_retries:
+                    raise DownloadError(
+                        f"下载失败: presigned URL 已失效 (HTTP {exc.status_code})"
+                    ) from exc
+                url = get_presigned_url()
+            elif attempts > max_retries:
+                raise DownloadError(
+                    f"下载失败（已重试 {max_retries} 次）: {exc}"
+                ) from exc
+            time.sleep(_backoff(attempts, retry_backoff))
+
+
+def _download_once(
+    *,
+    url: str,
+    local_path: Path,
+    chunk_size: int,
+    timeout: Union[float, httpx.Timeout, None],
+    overwrite: bool,
+    verify: bool,
+    on_progress: Optional[Callable[[int, Optional[int]], None]],
+) -> None:
+    """完成一次（含可能的一次从头重来）流式下载。"""
+    if overwrite:
+        start = 0
+    else:
+        start = local_path.stat().st_size if local_path.exists() else 0
+
+    while True:
+        try:
+            _stream_to_file(
+                url=url,
+                local_path=local_path,
+                start=start,
+                chunk_size=chunk_size,
+                timeout=timeout,
+                verify=verify,
+                on_progress=on_progress,
+            )
+            return
+        except _RestartNeeded:
+            start = 0
+
+
+def _stream_to_file(
+    *,
+    url: str,
+    local_path: Path,
+    start: int,
+    chunk_size: int,
+    timeout: Union[float, httpx.Timeout, None],
+    verify: bool,
+    on_progress: Optional[Callable[[int, Optional[int]], None]],
+) -> None:
+    """向服务端发起一次 Range 请求，并把响应体写入本地文件。"""
+    local_path.parent.mkdir(parents=True, exist_ok=True)
+
+    headers: dict[str, str] = {}
+    if start > 0:
+        headers["Range"] = f"bytes={start}-"
+
+    with httpx.Client(follow_redirects=True, verify=verify) as client:
+        with client.stream("GET", url, headers=headers, timeout=timeout) as resp:
+            status = resp.status_code
+
+            # 416: 请求的 Range 起始偏移已超出文件末尾，说明本地文件已完整
+            if status == 416:
+                total = _parse_total(resp.headers)
+                if on_progress is not None and total is not None and start >= total:
+                    on_progress(total, total)
+                return
+
+            if status == 200:
+                # 服务端忽略了 Range（不支持断点续传）：从头重新下载
+                if start > 0:
+                    raise _RestartNeeded()
+                mode = "wb"
+            elif status == 206:
+                # 校验服务端返回的起始偏移与本地已下载字节数一致，否则从头下载
+                range_start = _parse_range_start(resp.headers.get("content-range"))
+                if range_start is not None and range_start != start:
+                    raise _RestartNeeded()
+                mode = "wb" if start == 0 else "ab"
+            elif status >= 500:
+                raise _HttpDownloadError(status, f"服务端错误 (HTTP {status})")
+            elif status in (401, 403):
+                raise _HttpDownloadError(
+                    status, f"presigned URL 已失效 (HTTP {status})", refresh_url=True
+                )
+            else:
+                raise _HttpDownloadError(status, f"下载失败 (HTTP {status})")
+
+            total = _parse_total(resp.headers)
+            downloaded = start
+            with open(local_path, mode) as fp:
+                for chunk in resp.iter_bytes(chunk_size=chunk_size):
+                    fp.write(chunk)
+                    downloaded += len(chunk)
+                    if on_progress is not None:
+                        on_progress(downloaded, total)
+
+
+def _parse_total(headers: httpx.Headers) -> Optional[int]:
+    """从响应头解析文件总大小。
+
+    优先取 ``Content-Range`` 中的总大小（206 场景），否则取 ``Content-Length``。
+    """
+    content_range = headers.get("content-range")
+    if content_range:
+        # 形如 "bytes 0-1023/20480" 或 "bytes */20480"
+        try:
+            total = content_range.split("/", 1)[1]
+            return int(total) if total != "*" else None
+        except (IndexError, ValueError):
+            return None
+
+    content_length = headers.get("content-length")
+    if content_length:
+        try:
+            return int(content_length)
+        except ValueError:
+            return None
+    return None
+
+
+def _parse_range_start(content_range: Optional[str]) -> Optional[int]:
+    """从 ``Content-Range: bytes <start>-<end>/<total>`` 中解析起始偏移。"""
+    if not content_range:
+        return None
+    try:
+        # 先按空格切出 "start-end/total"，再按 "-" 取 start
+        return int(content_range.split(" ", 1)[1].split("-", 1)[0])
+    except (IndexError, ValueError):
+        return None
+
+
+def _backoff(attempt: int, base: float) -> float:
+    """指数退避 + 随机抖动。"""
+    return base * (2 ** (attempt - 1)) + random.uniform(0, base)