From 65253385f11f9705bcdb3093711946080fa482e7 Mon Sep 17 00:00:00 2001
From: Robin Huang <robin.j.huang@gmail.com>
Date: Thu, 7 Nov 2024 13:42:22 -0800
Subject: [PATCH] Handle case when Github releases rate limit is reached.
 (#204)

* Handle case when Github releases rate limit is reached.

* Attach GITHUB_TOKEN to releases request if available.

* Fix test.

* Fix test.

* Ruff fix.
---
 comfy_cli/command/install.py    | 37 ++++++++++++++++++++++++++++++---
 tests/comfy_cli/test_install.py |  2 +-
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/comfy_cli/command/install.py b/comfy_cli/command/install.py
index b80cba2..1ddf16c 100644
--- a/comfy_cli/command/install.py
+++ b/comfy_cli/command/install.py
@@ -188,13 +188,17 @@ def execute(
         clone_comfyui(url=url, repo_dir=repo_dir)
 
     if version != "nightly":
-        checkout_stable_comfyui(version=version, repo_dir=repo_dir)
+        try:
+            checkout_stable_comfyui(version=version, repo_dir=repo_dir)
+        except GitHubRateLimitError as e:
+            rprint(f"[bold red]Error checking out ComfyUI version: {e}[/bold red]")
+            sys.exit(1)
 
     elif not check_comfy_repo(repo_dir)[0]:
         rprint(
             f"[bold red]'{repo_dir}' already exists. But it is an invalid ComfyUI repository. Remove it and retry.[/bold red]"
         )
-        exit(-1)
+        sys.exit(-1)
 
     # checkout specified commit
     if commit is not None:
@@ -281,12 +285,39 @@ def validate_version(version: str) -> Optional[str]:
         ) from exc
 
 
+class GitHubRateLimitError(Exception):
+    """Raised when GitHub API rate limit is exceeded"""
+
+
 def fetch_github_releases(repo_owner: str, repo_name: str) -> List[Dict[str, str]]:
     """
     Fetch the list of releases from the GitHub API.
+    Handles rate limiting by logging the wait time.
     """
     url = f"https://api.github.com/repos/{repo_owner}/{repo_name}/releases"
-    response = requests.get(url)
+
+    headers = {}
+    if github_token := os.getenv("GITHUB_TOKEN"):
+        headers["Authorization"] = f"Bearer {github_token}"
+
+    response = requests.get(url, headers=headers, timeout=5)
+
+    # Handle rate limiting
+    if response.status_code in (403, 429):
+        # Check rate limit headers
+        remaining = int(response.headers.get("x-ratelimit-remaining", 0))
+        if remaining == 0:
+            reset_time = int(response.headers.get("x-ratelimit-reset", 0))
+            message = f"Primary rate limit from Github exceeded! Please retry after: {reset_time})"
+            raise GitHubRateLimitError(message)
+
+        if "retry-after" in response.headers:
+            wait_seconds = int(response.headers["retry-after"])
+            message = f"Rate limit from Github exceeded! Please wait {wait_seconds} seconds before retrying."
+            rprint(f"[yellow]{message}[/yellow]")
+            raise GitHubRateLimitError(message)
+
+    response.raise_for_status()
     return response.json()
 
 
diff --git a/tests/comfy_cli/test_install.py b/tests/comfy_cli/test_install.py
index 1e97b96..05d5290 100644
--- a/tests/comfy_cli/test_install.py
+++ b/tests/comfy_cli/test_install.py
@@ -53,7 +53,7 @@ def test_fetch_releases_success(mock_get):
     assert len(releases) == 2
     assert releases[0]["tag_name"] == "v1.0.0"
     assert releases[1]["tag_name"] == "v1.1.0"
-    mock_get.assert_called_once_with("https://api.github.com/repos/owner/repo/releases")
+    mock_get.assert_called_once_with("https://api.github.com/repos/owner/repo/releases", headers={}, timeout=5)
 
 
 @patch("requests.get")