add recovery mode

src/gradio_space_ci/webhook.py

@@ -1,6 +1,7 @@
 import os
 import uuid
 import warnings
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from typing import Any, Dict, List, Literal, Optional, TypedDict, Union
 
@@ -20,6 +21,7 @@ from huggingface_hub import (
     get_repo_discussions,
     get_space_runtime,
     get_space_variables,
+    repo_exists,
     request_space_hardware,
     request_space_storage,
     snapshot_download,
@@ -126,9 +128,38 @@ def configure_space_ci(
     server = WebhooksServer(ui=blocks, webhook_secret=WEBHOOK_SECRET)
     server.add_webhook()(trigger_ci_on_pr)
     configure_webhook_on_hub()
+
+    # Recover missed webhooks
+    recover_after_restart(space_id=SPACE_ID)
     return server
 
 
+###
+# Recovery logic
+###
+# Check if there are any PRs that need to be synced.
+# We might have missed some events while the server was down.
+# => called once at startup (see configure_space_ci)
+
+background_pool = ThreadPoolExecutor(max_workers=1)
+
+
+def recover_after_restart(space_id: str) -> None:
+    for discussion in get_repo_discussions(repo_id=space_id, repo_type="space"):
+        if discussion.is_pull_request:
+            if discussion.status == "open":
+                if not is_pr_synced(space_id=space_id, pr_num=discussion.num):
+                    # Found a PR that is not yet synced
+                    print(f"Recovery. Found an open PR that is not synced: {discussion.url}. Syncing it.")
+                    background_pool.submit(sync_ci_space, space_id=space_id, pr_num=discussion.num)
+            if discussion.status == "merged" or discussion.status == "closed":
+                ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=discussion.num)
+                if repo_exists(repo_id=ci_space_id, repo_type="space"):
+                    # Found a PR for which the CI space has not been deleted
+                    print(f"Recovery. Found a closed PR with an active CI space: {discussion.url}. Deleting it.")
+                    background_pool.submit(delete_ci_space, space_id=space_id, pr_num=discussion.num)
+
+
 ###
 # Define webhook on the Hub logic
 ###
@@ -173,12 +204,7 @@ async def trigger_ci_on_pr(payload: WebhookPayload, task_queue: BackgroundTasks)
     ):
         if not is_pr_synced(space_id=space_id, pr_num=payload.discussion.num):
             # New PR! Sync task scheduled
-            task_queue.add_task(
-                sync_ci_space,
-                space_id=space_id,
-                pr_num=payload.discussion.num,
-                private=payload.repo.private,
-            )
+            task_queue.add_task(sync_ci_space, space_id=space_id, pr_num=payload.discussion.num)
             has_task = True
     elif (
         # Means "a PR has been merged or closed"
@@ -204,12 +230,7 @@ async def trigger_ci_on_pr(payload: WebhookPayload, task_queue: BackgroundTasks)
             if discussion.is_pull_request and discussion.status == "open":
                 if not is_pr_synced(space_id=space_id, pr_num=discussion.num):
                     # Found a PR that is not yet synced
-                    task_queue.add_task(
-                        sync_ci_space,
-                        space_id=space_id,
-                        pr_num=discussion.num,
-                        private=payload.repo.private,
-                    )
+                    task_queue.add_task(sync_ci_space, space_id=space_id, pr_num=discussion.num)
                     has_task = True
 
     if has_task:
@@ -240,7 +261,7 @@ def is_pr_synced(space_id: str, pr_num: int) -> bool:
     return last_synced_sha == last_pr_sha
 
 
-def sync_ci_space(space_id: str, pr_num: int, private: bool) -> None:
+def sync_ci_space(space_id: str, pr_num: int) -> None:
     print(f"New task: sync ephemeral env for {space_id} (PR {pr_num})")
     ci_space_id = _get_ci_space_id(space_id=space_id, pr_num=pr_num)