1 files changed, 11 insertions, 4 deletions
diff --git a/third_party/python/taskcluster_taskgraph/taskgraph/actions/util.py b/third_party/python/taskcluster_taskgraph/taskgraph/actions/util.py
index cf81029da2..41e3b035de 100644
--- a/third_party/python/taskcluster_taskgraph/taskgraph/actions/util.py
+++ b/third_party/python/taskcluster_taskgraph/taskgraph/actions/util.py
@@ -32,8 +32,15 @@ def get_parameters(decision_task_id):
     return get_artifact(decision_task_id, "public/parameters.yml")
 
 
-def fetch_graph_and_labels(parameters, graph_config):
-    decision_task_id = find_decision_task(parameters, graph_config)
+def fetch_graph_and_labels(parameters, graph_config, task_group_id=None):
+    try:
+        # Look up the decision_task id in the index
+        decision_task_id = find_decision_task(parameters, graph_config)
+    except KeyError:
+        if not task_group_id:
+            raise
+        # Not found (e.g. from github-pull-request), fall back to the task group id.
+        decision_task_id = task_group_id
 
     # First grab the graph and labels generated during the initial decision task
     full_task_graph = get_artifact(decision_task_id, "public/full-task-graph.json")
@@ -90,7 +97,7 @@ def fetch_graph_and_labels(parameters, graph_config):
     return (decision_task_id, full_task_graph, label_to_taskid)
 
 
-def create_task_from_def(task_id, task_def, level):
+def create_task_from_def(task_id, task_def, level, trust_domain):
     """Create a new task from a definition rather than from a label
     that is already in the full-task-graph. The task definition will
     have {relative-datestamp': '..'} rendered just like in a decision task.
@@ -98,7 +105,7 @@ def create_task_from_def(task_id, task_def, level):
     It is useful if you want to "edit" the full_task_graph and then hand
     it to this function. No dependencies will be scheduled. You must handle
     this yourself. Seeing how create_tasks handles it might prove helpful."""
-    task_def["schedulerId"] = f"gecko-level-{level}"
+    task_def["schedulerId"] = f"{trust_domain}-level-{level}"
     label = task_def["metadata"]["name"]
     session = get_session()
     create.create_task(session, task_id, label, task_def)