From b03b9a0c16d0473d48e10d50b7694b15da423b08 Mon Sep 17 00:00:00 2001
From: abhishekmadan30 <amadan@mozilla.com>
Date: Mon, 6 Oct 2025 17:00:37 -0400
Subject: [PATCH] fix: add pagination for list incomplete tasks

---
 src/taskgraph/util/taskcluster.py | 27 +++++++++++++--------------
 test/test_util_taskcluster.py     | 28 +++++++++++++++++++++++++---
 2 files changed, 38 insertions(+), 17 deletions(-)

diff --git a/src/taskgraph/util/taskcluster.py b/src/taskgraph/util/taskcluster.py
index ab9f89290..9878a464f 100644
--- a/src/taskgraph/util/taskcluster.py
+++ b/src/taskgraph/util/taskcluster.py
@@ -411,20 +411,19 @@ def send_email(address, subject, content, link):
 def list_task_group_incomplete_tasks(task_group_id):
     """Generate the incomplete tasks in a task group"""
     queue = get_taskcluster_client("queue")
-    response = queue.listTaskGroup(task_group_id)
-
-    if not response or "tasks" not in response:
-        return
-
-    tasks = response.get("tasks", [])
-    for task in tasks:
-        if (status := task.get("status")) is not None:  # type: ignore
-            if (task_id := status.get("taskId")) and status.get("state") in [
-                "running",
-                "pending",
-                "unscheduled",
-            ]:
-                yield task_id
+
+    incomplete_tasks = []
+
+    def pagination_handler(response):
+        incomplete_tasks.extend(
+            task["status"]["taskId"]
+            for task in response["tasks"]
+            if task["status"]["state"] in ("running", "pending", "unscheduled")
+        )
+
+    queue.listTaskGroup(task_group_id, paginationHandler=pagination_handler)
+
+    return incomplete_tasks
 
 
 @functools.lru_cache(maxsize=None)
diff --git a/test/test_util_taskcluster.py b/test/test_util_taskcluster.py
index 45791e69d..59b840c68 100644
--- a/test/test_util_taskcluster.py
+++ b/test/test_util_taskcluster.py
@@ -401,15 +401,37 @@ def test_list_task_group_incomplete_tasks(responses, root_url):
         json={
             "tasks": [
                 {"status": {"taskId": "1", "state": "pending"}},
-                {"status": {"taskId": "2", "state": "unscheduled"}},
+                {"status": {"taskId": "2", "state": "completed"}},
                 {"status": {"taskId": "3", "state": "running"}},
-                {"status": {"taskId": "4", "state": "completed"}},
+            ],
+            "continuationToken": "page2",
+        },
+    )
+
+    responses.get(
+        f"{root_url}/api/queue/v1/task-group/{tgid}/list",
+        json={
+            "tasks": [
+                {"status": {"taskId": "4", "state": "unscheduled"}},
+                {"status": {"taskId": "5", "state": "failed"}},
+                {"status": {"taskId": "6", "state": "pending"}},
+            ],
+            "continuationToken": "page3",
+        },
+    )
+
+    responses.get(
+        f"{root_url}/api/queue/v1/task-group/{tgid}/list",
+        json={
+            "tasks": [
+                {"status": {"taskId": "7", "state": "running"}},
+                {"status": {"taskId": "8", "state": "completed"}},
             ]
         },
     )
 
     result = list(tc.list_task_group_incomplete_tasks(tgid))
-    assert result == ["1", "2", "3"]
+    assert result == ["1", "3", "4", "6", "7"]
 
 
 def test_get_ancestors(responses, root_url):