diff --git a/astrbot/cli/utils/plugin.py b/astrbot/cli/utils/plugin.py index 228a20a547..79d06e266e 100644 --- a/astrbot/cli/utils/plugin.py +++ b/astrbot/cli/utils/plugin.py @@ -63,36 +63,20 @@ def download_repository( Raises: ValueError: If the repository URL is unsupported or invalid. - httpx.HTTPError: If repository metadata or the archive cannot be downloaded. + httpx.HTTPError: If the repository archive cannot be downloaded. """ from astrbot.core.repository import GitHubRepository temp_dir = Path(tempfile.mkdtemp()) try: repository = GitHubRepository.parse(url) - if not repository.branch: - try: - with httpx.Client(follow_redirects=True, trust_env=True) as client: - response = client.get(repository.default_branch_api_url) - response.raise_for_status() - default_branch = str( - response.json().get("default_branch") or "" - ).strip() - except httpx.HTTPError as exc: - default_branch = "" - click.echo( - f"Failed to resolve the default GitHub branch: {exc}. Trying main." - ) - branch = default_branch or "main" - repository = GitHubRepository( - repository.owner, - repository.name, - branch, - ) - + reference = ( + f"branch {repository.branch}" + if repository.branch + else "default reference HEAD" + ) click.echo( - f"Downloading {repository.owner}/{repository.name} " - f"from GitHub branch {repository.branch}" + f"Downloading {repository.owner}/{repository.name} from GitHub {reference}" ) download_url = repository.archive_url if proxy: diff --git a/astrbot/core/repository.py b/astrbot/core/repository.py index a57ebdd130..e682540f93 100644 --- a/astrbot/core/repository.py +++ b/astrbot/core/repository.py @@ -47,7 +47,7 @@ class GitHubRepository: Args: owner: GitHub repository owner. name: GitHub repository name. - branch: Explicit or resolved source branch. + branch: Explicit source branch, or None to use the default HEAD reference. """ owner: str @@ -91,22 +91,15 @@ def parse(cls, url: str) -> "GitHubRepository": raise ValueError("Invalid GitHub repository URL") return cls(owner, name, branch) - @property - def default_branch_api_url(self) -> str: - """Return the GitHub repository metadata API URL.""" - owner = quote(self.owner, safe="") - name = quote(self.name, safe="") - return f"https://api.github.com/repos/{owner}/{name}" - @property def archive_url(self) -> str: - """Return the source ZIP URL for the resolved branch. + """Return the source ZIP URL for a branch or the default HEAD reference. - Raises: - ValueError: If the source branch has not been resolved. + Returns: + Branch archive URL, or the repository HEAD archive when unspecified. """ if not self.branch: - raise ValueError("GitHub source branch has not been resolved") + return self.revision_archive_url("HEAD") owner = quote(self.owner, safe="") name = quote(self.name, safe="") branch = quote(self.branch, safe="/") @@ -130,22 +123,17 @@ def revision_archive_url( return f"https://github.com/{owner}/{name}/archive/{encoded_revision}.zip" def raw_file_url(self, path: str) -> str: - """Return a raw file URL in the resolved branch. + """Return a raw file URL in a branch or the default HEAD reference. Args: path: Repository-relative file path. Returns: - GitHub raw file URL. - - Raises: - ValueError: If the source branch has not been resolved. + GitHub raw file URL, using HEAD when no branch is specified. """ - if not self.branch: - raise ValueError("GitHub source branch has not been resolved") owner = quote(self.owner, safe="") name = quote(self.name, safe="") - branch = quote(self.branch, safe="/") + branch = quote(self.branch or "HEAD", safe="/") encoded_path = quote(path.lstrip("/"), safe="/") return ( f"https://raw.githubusercontent.com/{owner}/{name}/{branch}/{encoded_path}" diff --git a/astrbot/core/star/updater.py b/astrbot/core/star/updater.py index 5a0ae4c517..68fe86878f 100644 --- a/astrbot/core/star/updater.py +++ b/astrbot/core/star/updater.py @@ -9,6 +9,7 @@ from astrbot.core import logger from astrbot.core.repository import ( + GitHubRepository, GitUnavailableError, normalize_repository_url, parse_repository_url, @@ -140,7 +141,7 @@ async def inspect_repository( await self._clone_repository(normalized_url, checkout_path) metadata = self.inspect_plugin_directory(checkout_path)["metadata"] else: - source = await self._resolve_repository_source(normalized_url) + source = GitHubRepository.parse(normalized_url) proxy = proxy.strip().removesuffix("/") async with self._create_httpx_client( timeout=PLUGIN_REPOSITORY_TIMEOUT_SECONDS diff --git a/astrbot/core/zip_updater.py b/astrbot/core/zip_updater.py index 9ba11998ed..2030053df8 100644 --- a/astrbot/core/zip_updater.py +++ b/astrbot/core/zip_updater.py @@ -1,3 +1,4 @@ +import asyncio import inspect import os import re @@ -59,7 +60,6 @@ def __init__( Args: verify: TLS certificate verification configuration for HTTPX. """ - self._rm_on_error = on_error self._httpx_verify = certifi.where() if verify is None else verify def _create_httpx_client(self, timeout: float = 30.0) -> httpx.AsyncClient: @@ -76,72 +76,6 @@ def _truncate_response_body(body: str, max_len: int = 1000) -> str: return body return body[:max_len] + "...[truncated]" - async def _fetch_repository_default_branch( - self, - repository: GitHubRepository, - ) -> str | None: - """Fetch the default branch for a repository. - - Args: - repository: Parsed GitHub repository. - - Returns: - The default branch name, or None if it cannot be resolved. - """ - url = repository.default_branch_api_url - try: - async with self._create_httpx_client(timeout=10.0) as client: - response = await client.get(url) - response.raise_for_status() - repo_info = response.json() - except Exception as exc: - logger.debug( - "Failed to get the default %s branch for %s/%s: %s", - "github", - repository.owner, - repository.name, - exc, - ) - return None - - default_branch = str(repo_info.get("default_branch") or "").strip() - return default_branch or None - - async def _resolve_repository_source( - self, - repo_url: str, - ) -> GitHubRepository: - """Resolve a repository URL to a downloadable source archive. - - Args: - repo_url: Repository URL, optionally with an explicit tree branch. - - Returns: - Resolved provider adapter and repository branch. - - Raises: - ValueError: If the repository URL is unsupported or invalid. - """ - repository = GitHubRepository.parse(repo_url) - if repository.branch: - return repository - - default_branch = await self._fetch_repository_default_branch(repository) - branch = default_branch or "main" - if not default_branch: - logger.info( - "Could not get the default %s branch for %s/%s; trying %s.", - "github", - repository.owner, - repository.name, - branch, - ) - return GitHubRepository( - repository.owner, - repository.name, - branch, - ) - async def _download_file( self, url: str, @@ -201,10 +135,18 @@ async def _emit_progress(payload: dict) -> None: "speed": 0, }, ) - except Exception as e: - logger.error(f"Failed to download file: {url} -> {target_path}: {e}") - if self._rm_on_error and target_path.exists(): - target_path.unlink() + except (asyncio.CancelledError, Exception) as error: + if not isinstance(error, asyncio.CancelledError): + logger.error( + f"Failed to download file: {url} -> {target_path}: {error}" + ) + try: + target_path.unlink(missing_ok=True) + except OSError as cleanup_error: + logger.warning( + "Failed to remove partial download: " + f"{url} -> {target_path}: {cleanup_error}" + ) raise async def _fetch_release_info(self, url: str, latest: bool = True) -> list: @@ -287,15 +229,16 @@ async def _check_update( async def _download_repository( self, target_path: str, repo_url: str, proxy="" ) -> None: - repository = await self._resolve_repository_source(repo_url) + repository = GitHubRepository.parse(repo_url) logger.info(f"Downloading update for {repository.name} ...") logger.info( - "Downloading %s/%s from %s branch %s", + "Downloading %s/%s from github %s", repository.owner, repository.name, - "github", - repository.branch, + f"branch {repository.branch}" + if repository.branch + else "default reference HEAD", ) release_url = repository.archive_url diff --git a/tests/test_repository.py b/tests/test_repository.py index 8b75459d64..b74d85736d 100644 --- a/tests/test_repository.py +++ b/tests/test_repository.py @@ -18,6 +18,22 @@ def test_github_repository_resolves_branch_with_slashes() -> None: assert repository.archive_url == ( "https://github.com/AstrBotDevs/AstrBot/archive/refs/heads/feature/updater.zip" ) + assert repository.raw_file_url("metadata.yaml") == ( + "https://raw.githubusercontent.com/AstrBotDevs/AstrBot/" + "feature/updater/metadata.yaml" + ) + + +def test_github_repository_uses_head_when_branch_is_unspecified() -> None: + repository = GitHubRepository.parse("https://github.com/AstrBotDevs/AstrBot") + + assert repository.branch is None + assert repository.archive_url == ( + "https://github.com/AstrBotDevs/AstrBot/archive/HEAD.zip" + ) + assert repository.raw_file_url("metadata.yaml") == ( + "https://raw.githubusercontent.com/AstrBotDevs/AstrBot/HEAD/metadata.yaml" + ) def test_non_github_http_repository_uses_git_transport() -> None: diff --git a/tests/test_updater_socks.py b/tests/test_updater_socks.py index 15bc1f168c..85f2a3c180 100644 --- a/tests/test_updater_socks.py +++ b/tests/test_updater_socks.py @@ -1,3 +1,4 @@ +import asyncio import ntpath import posixpath import zipfile @@ -47,6 +48,9 @@ async def aiter_bytes(self, chunk_size: int = 8192): class _FakeFailingStreamResponse: + def __init__(self, error: BaseException | None = None): + self._error = error if error is not None else RuntimeError("stream interrupted") + async def __aenter__(self): return self @@ -58,7 +62,7 @@ def raise_for_status(self) -> None: async def aiter_bytes(self, chunk_size: int = 8192): # noqa: ARG002 yield b"partial" - raise RuntimeError("stream interrupted") + raise self._error class _FakeStatusErrorResponse: @@ -162,6 +166,9 @@ async def get(self, url: str): class _FakeFailingStreamAsyncClient: + def __init__(self, error: BaseException | None = None): + self._error = error + async def __aenter__(self): return self @@ -169,7 +176,7 @@ async def __aexit__(self, exc_type, exc, tb) -> None: return None def stream(self, method: str, url: str): # noqa: ARG002 - return _FakeFailingStreamResponse() + return _FakeFailingStreamResponse(self._error) class _FakeZipArchive: @@ -1173,14 +1180,13 @@ async def test_fetch_release_info_uses_httpx_client_with_env_proxy_support( @pytest.mark.asyncio -async def test_download_from_repo_url_uses_httpx_stream_for_zip_download( +async def test_download_from_repo_url_uses_head_without_metadata_lookup( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, fake_async_client_state: _FakeAsyncClientState, ) -> None: import astrbot.core.zip_updater as zip_updater_module - fake_async_client_state.json_payload = {"default_branch": "trunk"} fake_async_client_state.stream_payload = b"zip-data" monkeypatch.setattr( zip_updater_module, @@ -1206,11 +1212,9 @@ async def test_download_from_repo_url_uses_httpx_stream_for_zip_download( ) assert (tmp_path / "AstrBot.zip").read_bytes() == b"zip-data" - assert fake_async_client_state.requested_urls == [ - "https://api.github.com/repos/AstrBotDevs/AstrBot" - ] + assert fake_async_client_state.requested_urls == [] assert fake_async_client_state.stream_urls == [ - "https://github.com/AstrBotDevs/AstrBot/archive/refs/heads/trunk.zip" + "https://github.com/AstrBotDevs/AstrBot/archive/HEAD.zip" ] assert fake_async_client_state.init_kwargs is not None assert fake_async_client_state.init_kwargs["follow_redirects"] is True @@ -1227,20 +1231,10 @@ async def test_download_from_repo_url_uses_explicit_branch_without_default_branc updater = _RepoZipUpdater() calls: list[str] = [] - async def fail_fetch_repository_default_branch( - repository, - ): # noqa: ARG001 - raise AssertionError("explicit branch should not fetch the default branch") - async def fake_download_file(url: str, path: str): calls.append(url) Path(path).write_bytes(b"zip-data") - monkeypatch.setattr( - updater, - "_fetch_repository_default_branch", - fail_fetch_repository_default_branch, - ) monkeypatch.setattr(updater, "_download_file", fake_download_file) await updater._download_repository( @@ -1340,16 +1334,6 @@ async def test_plugin_updater_inspects_github_repository_source( ) -> None: updater = _PluginUpdater() requested_urls: list[str] = [] - source = SimpleNamespace( - raw_file_url=lambda filename: ( - "https://raw.githubusercontent.com/AstrBotDevs/" - f"astrbot-plugin-demo/trunk/{filename}" - ), - ) - - async def fake_resolve_repository_source(repo_url: str): - assert repo_url == "https://github.com/AstrBotDevs/astrbot-plugin-demo" - return source def handle_request(request: httpx.Request) -> httpx.Response: requested_urls.append(str(request.url)) @@ -1367,11 +1351,6 @@ def handle_request(request: httpx.Request) -> httpx.Response: ), ) - monkeypatch.setattr( - updater, - "_resolve_repository_source", - fake_resolve_repository_source, - ) monkeypatch.setattr( updater, "_create_httpx_client", @@ -1388,10 +1367,12 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert result["name"] == "astrbot_plugin_demo" assert result["desc"] == "Demo plugin" - assert requested_urls[-1] == ( + assert requested_urls == [ "https://proxy.example/https://raw.githubusercontent.com/AstrBotDevs/" - "astrbot-plugin-demo/trunk/metadata.yml" - ) + "astrbot-plugin-demo/HEAD/metadata.yaml", + "https://proxy.example/https://raw.githubusercontent.com/AstrBotDevs/" + "astrbot-plugin-demo/HEAD/metadata.yml", + ] @pytest.mark.asyncio @@ -1399,12 +1380,6 @@ async def test_plugin_updater_rejects_large_repository_metadata( monkeypatch: pytest.MonkeyPatch, ) -> None: updater = _PluginUpdater() - source = SimpleNamespace( - raw_file_url=lambda filename: f"https://example.com/{filename}", - ) - - async def fake_resolve_repository_source(repo_url: str): # noqa: ARG001 - return source def handle_request(request: httpx.Request) -> httpx.Response: # noqa: ARG001 return httpx.Response( @@ -1412,11 +1387,6 @@ def handle_request(request: httpx.Request) -> httpx.Response: # noqa: ARG001 headers={"Content-Length": str(1024 * 1024 + 1)}, ) - monkeypatch.setattr( - updater, - "_resolve_repository_source", - fake_resolve_repository_source, - ) monkeypatch.setattr( updater, "_create_httpx_client", @@ -1511,6 +1481,86 @@ async def test_download_file_removes_partial_file_when_stream_fails( assert not target_path.exists() +@pytest.mark.asyncio +async def test_download_file_removes_partial_file_when_cancelled( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + cancellation = asyncio.CancelledError("download cancelled") + monkeypatch.setattr( + _RepoZipUpdater, + "_create_httpx_client", + staticmethod( + lambda timeout=30.0: _FakeFailingStreamAsyncClient( # noqa: ARG005 + cancellation + ) + ), + ) + + target_path = tmp_path / "cancelled.zip" + + with pytest.raises(asyncio.CancelledError) as exc_info: + await _RepoZipUpdater()._download_file( + "https://example.com/archive.zip", + str(target_path), + ) + + assert exc_info.value is cancellation + assert not target_path.exists() + + +@pytest.mark.parametrize( + "download_error", + [ + pytest.param(RuntimeError("stream interrupted"), id="stream-error"), + pytest.param(asyncio.CancelledError("download cancelled"), id="cancellation"), + ], +) +@pytest.mark.asyncio +async def test_download_file_preserves_original_error_when_cleanup_fails( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + download_error: BaseException, +) -> None: + import astrbot.core.zip_updater as zip_updater_module + + target_path = tmp_path / "undeletable.zip" + log_messages: list[str] = [] + original_unlink = Path.unlink + + def fail_target_unlink(path: Path, missing_ok: bool = False) -> None: + if path == target_path: + raise OSError("permission denied") + original_unlink(path, missing_ok=missing_ok) + + monkeypatch.setattr( + _RepoZipUpdater, + "_create_httpx_client", + staticmethod( + lambda timeout=30.0: _FakeFailingStreamAsyncClient( # noqa: ARG005 + download_error + ) + ), + ) + monkeypatch.setattr(Path, "unlink", fail_target_unlink) + monkeypatch.setattr( + zip_updater_module.logger, + "warning", + lambda message: log_messages.append(message), + ) + + with pytest.raises(type(download_error)) as exc_info: + await _RepoZipUpdater()._download_file( + "https://example.com/archive.zip", + str(target_path), + ) + + assert exc_info.value is download_error + assert target_path.exists() + assert any(str(target_path) in message for message in log_messages) + assert any("permission denied" in message for message in log_messages) + + @pytest.mark.asyncio async def test_download_file_logs_url_and_target_path_on_failure( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/test_cli_plugin_utils.py b/tests/unit/test_cli_plugin_utils.py index 8cf8673ade..afd185a71a 100644 --- a/tests/unit/test_cli_plugin_utils.py +++ b/tests/unit/test_cli_plugin_utils.py @@ -1,6 +1,8 @@ +from io import BytesIO from pathlib import Path +from zipfile import ZipFile -from astrbot.cli.utils.plugin import PluginStatus, build_plug_list +from astrbot.cli.utils.plugin import PluginStatus, build_plug_list, download_repository class FakeResponse: @@ -36,6 +38,44 @@ def get(self, url): return FakeResponse() +def test_download_repository_uses_head_without_metadata_lookup( + monkeypatch, tmp_path, capsys +): + archive = BytesIO() + with ZipFile(archive, "w") as zip_file: + zip_file.writestr("plugin-commit/main.py", "VALUE = 1\n") + requested_urls = [] + + class ArchiveResponse: + content = archive.getvalue() + + def raise_for_status(self): + return None + + class ArchiveClient: + def __init__(self, **kwargs): + assert kwargs == {"follow_redirects": True, "trust_env": True} + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def get(self, url): + requested_urls.append(url) + return ArchiveResponse() + + monkeypatch.setattr("astrbot.cli.utils.plugin.httpx.Client", ArchiveClient) + + target_path = tmp_path / "plugin" + download_repository("https://github.com/example/plugin", target_path) + + assert requested_urls == ["https://github.com/example/plugin/archive/HEAD.zip"] + assert (target_path / "main.py").read_text(encoding="utf-8") == "VALUE = 1\n" + assert "default reference HEAD" in capsys.readouterr().out + + def write_metadata(plugin_dir: Path, name: str, version: str) -> None: plugin_dir.mkdir(parents=True, exist_ok=True) plugin_dir.joinpath("metadata.yaml").write_text(