From 315db1f27b49dc92d6f16e513944da63f9f697f0 Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Tue, 9 Sep 2025 13:07:12 -0400
Subject: [PATCH 01/10] chore(docker): remove unnecessary deploy warning

We don't really bother with REGISTRY and VERSION files anymore, so the
scary warning when we don't have a versioned tag is more misleading than
helpful.
---
 src/taskgraph/docker.py | 19 -------------------
 test/test_docker.py     |  2 --
 2 files changed, 21 deletions(-)

diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py
index 9301e458..223a19c4 100644
--- a/src/taskgraph/docker.py
+++ b/src/taskgraph/docker.py
@@ -29,22 +29,6 @@
     get_task_definition,
 )
 
-DEPLOY_WARNING = """
-*****************************************************************
-WARNING: Image is not suitable for deploying/pushing.
-
-To automatically tag the image the following files are required:
-- {image_dir}/REGISTRY
-- {image_dir}/VERSION
-
-The REGISTRY file contains the Docker registry hosting the image.
-A default REGISTRY file may also be defined in the parent docker
-directory.
-
-The VERSION file contains the version of the image.
-*****************************************************************
-"""
-
 
 def get_image_digest(image_name: str) -> str:
     """Get the digest of a docker image by its name.
@@ -212,9 +196,6 @@ def build_image(
         msg += f" and tagged with {tag}"
     print(msg)
 
-    if not tag or tag.endswith(":latest"):
-        print(DEPLOY_WARNING.format(image_dir=os.path.relpath(image_dir), image=name))
-
 
 def load_image(
     url: str, imageName: Optional[str] = None, imageTag: Optional[str] = None
diff --git a/test/test_docker.py b/test/test_docker.py
index 4a7b8ca8..894067e1 100644
--- a/test/test_docker.py
+++ b/test/test_docker.py
@@ -49,7 +49,6 @@ def test_build_image(capsys, mock_docker_build):
 
     out, _ = capsys.readouterr()
     assert f"Successfully built {image} and tagged with {tag}" in out
-    assert "Image is not suitable for deploying/pushing" not in out
 
 
 def test_build_image_no_tag(capsys, mock_docker_build):
@@ -74,7 +73,6 @@ def test_build_image_no_tag(capsys, mock_docker_build):
 
     out, _ = capsys.readouterr()
     assert f"Successfully built {image}" in out
-    assert "Image is not suitable for deploying/pushing" in out
 
 
 def test_build_image_error(capsys, mock_docker_build):

From c2ea2cbeffdee3bb8047b93eab0f0bb75de13fb2 Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Tue, 9 Sep 2025 13:07:12 -0400
Subject: [PATCH 02/10] chore(docker): use logging rather than print

---
 src/taskgraph/docker.py | 36 +++++++++++++++++++-----------------
 test/test_docker.py     | 32 +++++++++++++++++++-------------
 2 files changed, 38 insertions(+), 30 deletions(-)

diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py
index 223a19c4..b2244534 100644
--- a/src/taskgraph/docker.py
+++ b/src/taskgraph/docker.py
@@ -2,11 +2,10 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 
-
+import logging
 import os
 import shlex
 import subprocess
-import sys
 import tarfile
 import tempfile
 from io import BytesIO
@@ -29,6 +28,8 @@
     get_task_definition,
 )
 
+logger = logging.getLogger(__name__)
+
 
 def get_image_digest(image_name: str) -> str:
     """Get the digest of a docker image by its name.
@@ -81,14 +82,12 @@ def load_image_by_name(image_name: str, tag: Optional[str] = None) -> Optional[s
     task_id = IndexSearch().should_replace_task(task, {}, None, indexes)
 
     if task_id in (True, False):
-        print(
+        logger.error(
             "Could not find artifacts for a docker image "
-            "named `{image_name}`. Local commits and other changes "
+            f"named `{image_name}`. Local commits and other changes "
             "in your checkout may cause this error. Try "
-            "updating to a fresh checkout of {project} "
-            "to download image.".format(
-                image_name=image_name, project=params["project"]
-            )
+            f"updating to a fresh checkout of {params['project']} "
+            "to download image."
         )
         return None
 
@@ -111,12 +110,12 @@ def load_image_by_task_id(task_id: str, tag: Optional[str] = None) -> str:
     """
     artifact_url = get_artifact_url(task_id, "public/image.tar.zst")
     result = load_image(artifact_url, tag)
-    print("Found docker image: {}:{}".format(result["image"], result["tag"]))
+    logger.info(f"Found docker image: {result['image']}:{result['tag']}")
     if tag:
-        print(f"Re-tagged as: {tag}")
+        logger.info(f"Re-tagged as: {tag}")
     else:
-        tag = "{}:{}".format(result["image"], result["tag"])
-    print(f"Try: docker run -ti --rm {tag} bash")
+        tag = f"{result['image']}:{result['tag']}"
+    logger.info(f"Try: docker run -ti --rm {tag} bash")
     return tag
 
 
@@ -175,6 +174,7 @@ def build_image(
         ValueError: If name is not provided.
         Exception: If the image directory does not exist.
     """
+    logger.info(f"Building {name} image")
     if not name:
         raise ValueError("must provide a Docker image name")
 
@@ -194,7 +194,7 @@ def build_image(
     msg = f"Successfully built {name}"
     if tag:
         msg += f" and tagged with {tag}"
-    print(msg)
+    logger.info(msg)
 
 
 def load_image(
@@ -244,7 +244,7 @@ def load_image(
     def download_and_modify_image() -> Generator[bytes, None, None]:
         # This function downloads and edits the downloaded tar file on the fly.
         # It emits chunked buffers of the edited tar file, as a generator.
-        print(f"Downloading from {url}")
+        logger.info(f"Downloading from {url}")
         # get_session() gets us a requests.Session set to retry several times.
         req = get_session().get(url, stream=True)
         req.raise_for_status()
@@ -404,21 +404,23 @@ def load_task(
     user = user or "worker"
     task_def = get_task_definition(task_id)
 
+    logger.info(f"Loading '{task_def['metadata']['name']}' from {task_id}")
+
     if "payload" not in task_def or not (image := task_def["payload"].get("image")):
-        print("Tasks without a `payload.image` are not supported!")
+        logger.error("Tasks without a `payload.image` are not supported!")
 
         return 1
 
     command = task_def["payload"].get("command")  # type: ignore
     if not command or not command[0].endswith("run-task"):
-        print("Only tasks using `run-task` are supported!")
+        logger.error("Only tasks using `run-task` are supported!")
         return 1
 
     try:
         image = custom_image or image
         image_tag = _resolve_image(image, graph_config)
     except Exception as e:
-        print(e, file=sys.stderr)
+        logger.exception(e)
         return 1
 
     # Remove the payload section of the task's command. This way run-task will
diff --git a/test/test_docker.py b/test/test_docker.py
index 894067e1..f201f0eb 100644
--- a/test/test_docker.py
+++ b/test/test_docker.py
@@ -1,3 +1,4 @@
+import logging
 import re
 import tempfile
 
@@ -26,7 +27,8 @@ def side_effect(topsrcdir, context_dir, out_file, image_name=None, args=None):
     return (m_stream, m_run)
 
 
-def test_build_image(capsys, mock_docker_build):
+def test_build_image(caplog, mock_docker_build):
+    caplog.set_level(logging.DEBUG)
     m_stream, m_run = mock_docker_build
     image = "hello-world-tag"
     tag = f"test/{image}:1.0"
@@ -47,11 +49,11 @@ def test_build_image(capsys, mock_docker_build):
         check=True,
     )
 
-    out, _ = capsys.readouterr()
-    assert f"Successfully built {image} and tagged with {tag}" in out
+    assert f"Successfully built {image} and tagged with {tag}" in caplog.text
 
 
-def test_build_image_no_tag(capsys, mock_docker_build):
+def test_build_image_no_tag(caplog, mock_docker_build):
+    caplog.set_level(logging.DEBUG)
     m_stream, m_run = mock_docker_build
     image = "hello-world"
 
@@ -71,11 +73,11 @@ def test_build_image_no_tag(capsys, mock_docker_build):
         check=True,
     )
 
-    out, _ = capsys.readouterr()
-    assert f"Successfully built {image}" in out
+    assert f"Successfully built {image}" in caplog.text
 
 
-def test_build_image_error(capsys, mock_docker_build):
+def test_build_image_error(caplog, mock_docker_build):
+    caplog.set_level(logging.DEBUG)
     m_stream, m_run = mock_docker_build
 
     def mock_run(*popenargs, check=False, **kwargs):
@@ -103,8 +105,7 @@ def mock_run(*popenargs, check=False, **kwargs):
         check=True,
     )
 
-    out, _ = capsys.readouterr()
-    assert f"Successfully built {image}" not in out
+    assert f"Successfully built {image}" not in caplog.text
 
 
 @pytest.fixture
@@ -147,7 +148,7 @@ def inner(task, remove=False, custom_image=None):
 
 
 def test_load_task_invalid_task(run_load_task):
-    task = {}
+    task = {"metadata": {"name": "abc"}}
     assert run_load_task(task)[0] == 1
 
     task["payload"] = {}
@@ -167,6 +168,7 @@ def test_load_task_invalid_task(run_load_task):
 def test_load_task(run_load_task):
     image_task_id = "def"
     task = {
+        "metadata": {"name": "test-task"},
         "payload": {
             "command": [
                 "/usr/bin/run-task",
@@ -241,6 +243,7 @@ def test_load_task_env_init_and_remove(mocker, run_load_task):
 
     image_task_id = "def"
     task = {
+        "metadata": {"name": "test-task-env"},
         "payload": {
             "command": [
                 "/usr/bin/run-task",
@@ -307,6 +310,7 @@ def test_load_task_with_different_image_types(
     task_id = "abc"
     image_task_id = "xyz"
     task = {
+        "metadata": {"name": "test-task-image-types"},
         "payload": {
             "command": [
                 "/usr/bin/run-task",
@@ -328,8 +332,10 @@ def test_load_task_with_different_image_types(
     mocks["load_image_by_task_id"].assert_called_once_with(image_task_id)
 
 
-def test_load_task_with_unsupported_image_type(capsys, run_load_task):
+def test_load_task_with_unsupported_image_type(caplog, run_load_task):
+    caplog.set_level(logging.DEBUG)
     task = {
+        "metadata": {"name": "test-task-unsupported"},
         "payload": {
             "command": [
                 "/usr/bin/run-task",
@@ -344,13 +350,13 @@ def test_load_task_with_unsupported_image_type(capsys, run_load_task):
     ret, _ = run_load_task(task)
     assert ret == 1
 
-    _, err = capsys.readouterr()
-    assert "Tasks with unsupported-type images are not supported!" in err
+    assert "Tasks with unsupported-type images are not supported!" in caplog.text
 
 
 @pytest.fixture
 def task():
     return {
+        "metadata": {"name": "test-task-fixture"},
         "payload": {
             "command": [
                 "/usr/bin/run-task",

From c3bf9585e70c145fb518c318c37a1bf5359e525e Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Tue, 9 Sep 2025 13:07:12 -0400
Subject: [PATCH 03/10] feat: support any graph_attr and forward **kwargs in
 load_tasks_from_kind(s)

Graph_attr is added so we can generated the morphed graph.

Kwargs are being forwarded so we can use `write_artifacts=True` with
`load_tasks_from_kind`.
---
 src/taskgraph/generator.py | 14 +++++++++-----
 test/test_generator.py     | 20 ++++++++++++++++++++
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/src/taskgraph/generator.py b/src/taskgraph/generator.py
index f1a42730..89876168 100644
--- a/src/taskgraph/generator.py
+++ b/src/taskgraph/generator.py
@@ -550,28 +550,32 @@ def verify(self, name, *args, **kwargs):
             return name, args[0]
 
 
-def load_tasks_for_kinds(parameters, kinds, root_dir=None):
+def load_tasks_for_kinds(
+    parameters, kinds, root_dir=None, graph_attr=None, **tgg_kwargs
+):
     """
     Get all the tasks of the given kinds.
 
     This function is designed to be called from outside of taskgraph.
     """
+    graph_attr = graph_attr or "full_task_set"
+
     # make parameters read-write
     parameters = dict(parameters)
     parameters["target-kinds"] = kinds
     parameters = parameters_loader(spec=None, strict=False, overrides=parameters)
-    tgg = TaskGraphGenerator(root_dir=root_dir, parameters=parameters)
+    tgg = TaskGraphGenerator(root_dir=root_dir, parameters=parameters, **tgg_kwargs)
     return {
         task.task["metadata"]["name"]: task
-        for task in tgg.full_task_set
+        for task in getattr(tgg, graph_attr)
         if task.kind in kinds
     }
 
 
-def load_tasks_for_kind(parameters, kind, root_dir=None):
+def load_tasks_for_kind(parameters, kind, root_dir=None, **tgg_kwargs):
     """
     Get all the tasks of a given kind.
 
     This function is designed to be called from outside of taskgraph.
     """
-    return load_tasks_for_kinds(parameters, [kind], root_dir)
+    return load_tasks_for_kinds(parameters, [kind], root_dir, **tgg_kwargs)
diff --git a/test/test_generator.py b/test/test_generator.py
index 39a587c9..ca72ac21 100644
--- a/test/test_generator.py
+++ b/test/test_generator.py
@@ -215,6 +215,26 @@ def test_load_tasks_for_kind(monkeypatch):
         and tasks["_example-kind-t-1"].label == "_example-kind-t-1"
     )
 
+    # Test that **kwargs are forwarded to TaskGraphGenerator
+    tasks_with_kwargs = load_tasks_for_kind(
+        {"_kinds": [("_example-kind", []), ("docker-image", [])]},
+        "_example-kind",
+        "/root/taskcluster",
+        write_artifacts=True,  # This should be forwarded to TaskGraphGenerator
+    )
+    assert isinstance(tasks_with_kwargs, dict)
+    assert "_example-kind-t-1" in tasks_with_kwargs
+
+    # Test graph_attr parameter
+    tasks_with_graph_attr = load_tasks_for_kinds(
+        {"_kinds": [("_example-kind", []), ("docker-image", [])]},
+        ["_example-kind"],
+        "/root/taskcluster",
+        graph_attr="full_task_set",
+    )
+    assert isinstance(tasks_with_graph_attr, dict)
+    assert "_example-kind-t-1" in tasks_with_graph_attr
+
 
 def test_loader_backwards_compat_interface(graph_config):
     """Ensure loaders can be called even if they don't support a

From 301258d4f39b771b3bc3ba18656f1f89ba3c4994 Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Tue, 9 Sep 2025 13:07:12 -0400
Subject: [PATCH 04/10] feat(load-task): support piping JSON task definition
 via stdin

This allows you to quickly iterate on a task definition without needing
to actually run it in Taskcluster first.
---
 src/taskgraph/docker.py | 18 +++++---
 src/taskgraph/main.py   | 16 +++++++-
 test/test_docker.py     | 52 +++++++++++++++++++----
 test/test_main.py       | 91 +++++++++++++++++++++++++++++++++++++++++
 4 files changed, 162 insertions(+), 15 deletions(-)

diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py
index b2244534..1f8895a3 100644
--- a/src/taskgraph/docker.py
+++ b/src/taskgraph/docker.py
@@ -11,7 +11,7 @@
 from io import BytesIO
 from pathlib import Path
 from textwrap import dedent
-from typing import Dict, Generator, List, Mapping, Optional, Union
+from typing import Any, Dict, Generator, List, Mapping, Optional, Union
 
 try:
     import zstandard as zstd
@@ -374,7 +374,7 @@ def _resolve_image(image: Union[str, Dict[str, str]], graph_config: GraphConfig)
 
 def load_task(
     graph_config: GraphConfig,
-    task_id: str,
+    task: Union[str, Dict[str, Any]],
     remove: bool = True,
     user: Optional[str] = None,
     custom_image: Optional[str] = None,
@@ -388,7 +388,7 @@ def load_task(
 
     Args:
         graph_config: The graph configuration object.
-        task_id: The ID of the task to load.
+        task: The ID of the task, or task definition to load.
         remove: Whether to remove the container after exit (default True).
         user: The user to switch to in the container (default 'worker').
         custom_image: A custom image to use instead of the task's image.
@@ -402,9 +402,17 @@ def load_task(
         in the interactive shell.
     """
     user = user or "worker"
-    task_def = get_task_definition(task_id)
+    if isinstance(task, str):
+        task_id = task
+        task_def = get_task_definition(task)
+        source = f"task {task_id}"
+    else:
+        # We're running a locally generated definition and don't have a proper id.
+        task_id = "fake"
+        task_def = task
+        source = "provided definition"
 
-    logger.info(f"Loading '{task_def['metadata']['name']}' from {task_id}")
+    logger.info(f"Loading '{task_def['metadata']['name']}' from {source}")
 
     if "payload" not in task_def or not (image := task_def["payload"].get("image")):
         logger.error("Tasks without a `payload.image` are not supported!")
diff --git a/src/taskgraph/main.py b/src/taskgraph/main.py
index 10ec361c..d1a0149e 100644
--- a/src/taskgraph/main.py
+++ b/src/taskgraph/main.py
@@ -674,7 +674,11 @@ def image_digest(args):
     "The task's payload.command will be replaced with 'bash'. You need to have "
     "docker installed and running for this to work.",
 )
-@argument("task_id", help="The task id to load into a docker container.")
+@argument(
+    "task",
+    help="The task id or definition to load into a docker container. Can use "
+    "'-' to read from stdin.",
+)
 @argument(
     "--keep",
     dest="remove",
@@ -699,14 +703,22 @@ def image_digest(args):
 def load_task(args):
     from taskgraph.config import load_graph_config  # noqa: PLC0415
     from taskgraph.docker import load_task  # noqa: PLC0415
+    from taskgraph.util import json  # noqa: PLC0415
 
     validate_docker()
 
+    if args["task"] == "-":
+        data = sys.stdin.read()
+        try:
+            args["task"] = json.loads(data)
+        except ValueError:
+            args["task"] = data  # assume it is a taskId
+
     root = args["root"]
     graph_config = load_graph_config(root)
     return load_task(
         graph_config,
-        args["task_id"],
+        args["task"],
         remove=args["remove"],
         user=args["user"],
         custom_image=args["image"],
diff --git a/test/test_docker.py b/test/test_docker.py
index f201f0eb..81ffdc29 100644
--- a/test/test_docker.py
+++ b/test/test_docker.py
@@ -110,9 +110,7 @@ def mock_run(*popenargs, check=False, **kwargs):
 
 @pytest.fixture
 def run_load_task(mocker):
-    task_id = "abc"
-
-    def inner(task, remove=False, custom_image=None):
+    def inner(task, remove=False, custom_image=None, pass_task_def=False):
         proc = mocker.MagicMock()
         proc.returncode = 0
 
@@ -128,9 +126,6 @@ def inner(task, remove=False, custom_image=None):
             "build_image": mocker.patch.object(
                 docker, "build_image", return_value=None
             ),
-            "get_task_definition": mocker.patch.object(
-                docker, "get_task_definition", return_value=task
-            ),
             "load_image_by_task_id": mocker.patch.object(
                 docker, "load_image_by_task_id", return_value="image/tag"
             ),
@@ -139,8 +134,19 @@ def inner(task, remove=False, custom_image=None):
             ),
         }
 
+        # If testing with task ID, mock get_task_definition
+        if not pass_task_def:
+            task_id = "abc"
+            mocks["get_task_definition"] = mocker.patch.object(
+                docker, "get_task_definition", return_value=task
+            )
+            input_arg = task_id
+        else:
+            # Testing with task definition directly
+            input_arg = task
+
         ret = docker.load_task(
-            graph_config, task_id, remove=remove, custom_image=custom_image
+            graph_config, input_arg, remove=remove, custom_image=custom_image
         )
         return ret, mocks
 
@@ -183,7 +189,8 @@ def test_load_task(run_load_task):
     ret, mocks = run_load_task(task)
     assert ret == 0
 
-    mocks["get_task_definition"].assert_called_once_with("abc")
+    if "get_task_definition" in mocks:
+        mocks["get_task_definition"].assert_called_once_with("abc")
     mocks["load_image_by_task_id"].assert_called_once_with(image_task_id)
 
     expected = [
@@ -353,6 +360,35 @@ def test_load_task_with_unsupported_image_type(caplog, run_load_task):
     assert "Tasks with unsupported-type images are not supported!" in caplog.text
 
 
+def test_load_task_with_task_definition(run_load_task, caplog):
+    # Test passing a task definition directly instead of a task ID
+    caplog.set_level(logging.INFO)
+    image_task_id = "def"
+    task = {
+        "metadata": {"name": "test-task-direct"},
+        "payload": {
+            "command": [
+                "/usr/bin/run-task",
+                "--repo-checkout=/builds/worker/vcs/repo",
+                "--task-cwd=/builds/worker/vcs/repo",
+                "--",
+                "echo foo",
+            ],
+            "image": {"taskId": image_task_id, "type": "task-image"},
+        },
+    }
+
+    ret, mocks = run_load_task(task, pass_task_def=True)
+    assert ret == 0
+
+    # Should not call get_task_definition when passing a definition directly
+    assert "get_task_definition" not in mocks
+    mocks["load_image_by_task_id"].assert_called_once_with(image_task_id)
+
+    # Check logging output shows it's from provided definition
+    assert "Loading 'test-task-direct' from provided definition" in caplog.text
+
+
 @pytest.fixture
 def task():
     return {
diff --git a/test/test_main.py b/test/test_main.py
index f01a77bd..7327dd04 100644
--- a/test/test_main.py
+++ b/test/test_main.py
@@ -405,3 +405,94 @@ def fake_actions_load(_graph_config):
         action = actions[0]
         for field_name in ("name", "title", "description"):
             assert action[field_name] == expected_fields[field_name]
+
+
+@pytest.fixture
+def run_load_task(mocker, monkeypatch):
+    def inner(args, stdin_data=None):
+        # Mock the docker module functions
+        m_validate_docker = mocker.patch("taskgraph.main.validate_docker")
+        m_load_graph_config = mocker.patch("taskgraph.config.load_graph_config")
+        m_docker_load_task = mocker.patch("taskgraph.docker.load_task")
+
+        # Mock graph config
+        graph_config = mocker.MagicMock()
+        m_load_graph_config.return_value = graph_config
+
+        # Mock stdin if data provided
+        if stdin_data is not None:
+            mock_stdin = mocker.MagicMock()
+            mock_stdin.read.return_value = stdin_data
+            monkeypatch.setattr(sys, "stdin", mock_stdin)
+
+        # Set default return value for docker.load_task
+        m_docker_load_task.return_value = 0
+
+        mocks = {
+            "validate_docker": m_validate_docker,
+            "load_graph_config": m_load_graph_config,
+            "docker_load_task": m_docker_load_task,
+            "graph_config": graph_config,
+        }
+
+        # Run the command
+        result = taskgraph_main(args)
+
+        return result, mocks
+
+    return inner
+
+
+def test_load_task_command(run_load_task):
+    # Test normal task ID
+    result, mocks = run_load_task(["load-task", "task-id-123"])
+
+    assert result == 0
+    mocks["validate_docker"].assert_called_once()
+    mocks["load_graph_config"].assert_called_once_with("taskcluster")
+    mocks["docker_load_task"].assert_called_once_with(
+        mocks["graph_config"],
+        "task-id-123",
+        remove=True,
+        user=None,
+        custom_image=None,
+    )
+
+
+def test_load_task_command_with_stdin(run_load_task):
+    # Test with JSON task definition from stdin
+    task_def = {
+        "metadata": {"name": "test-task"},
+        "payload": {
+            "image": {"type": "task-image", "taskId": "image-task-id"},
+            "command": ["echo", "hello"],
+        },
+    }
+    stdin_data = json.dumps(task_def)
+
+    result, mocks = run_load_task(["load-task", "-"], stdin_data=stdin_data)
+
+    assert result == 0
+    mocks["docker_load_task"].assert_called_once_with(
+        mocks["graph_config"],
+        task_def,
+        remove=True,
+        user=None,
+        custom_image=None,
+    )
+
+
+def test_load_task_command_with_task_id(run_load_task):
+    # Test with task ID from stdin (non-JSON)
+    stdin_data = "task-id-from-stdin"
+
+    result, mocks = run_load_task(["load-task", "-"], stdin_data=stdin_data)
+
+    assert result == 0
+    mocks["docker_load_task"].assert_called_once_with(
+        mocks["graph_config"],
+        "task-id-from-stdin",
+        remove=True,
+        user=None,
+        custom_image=None,
+    )

From 87484b54043c6e2a7ee4e161cc36fdb7b2046147 Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Tue, 9 Sep 2025 13:07:12 -0400
Subject: [PATCH 05/10] feat(load-task): implement interactive=False to execute
 tasks without pausing

By default, we load the task with `run-task` then pause execution so the
user can mess around before executing the task. This introduces a way to
support non run-task based tasks by setting interactive=False. This
simply runs the task right away without pausing execution.
---
 src/taskgraph/docker.py | 96 +++++++++++++++++++++++------------------
 test/test_docker.py     | 73 +++++++++++++++++++++++++++++--
 2 files changed, 123 insertions(+), 46 deletions(-)

diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py
index 1f8895a3..2d243df8 100644
--- a/src/taskgraph/docker.py
+++ b/src/taskgraph/docker.py
@@ -6,6 +6,7 @@
 import os
 import shlex
 import subprocess
+import sys
 import tarfile
 import tempfile
 from io import BytesIO
@@ -378,6 +379,7 @@ def load_task(
     remove: bool = True,
     user: Optional[str] = None,
     custom_image: Optional[str] = None,
+    interactive: Optional[bool] = True,
 ) -> int:
     """Load and run a task interactively in a Docker container.
 
@@ -392,14 +394,12 @@ def load_task(
         remove: Whether to remove the container after exit (default True).
         user: The user to switch to in the container (default 'worker').
         custom_image: A custom image to use instead of the task's image.
+        interactive: If True, execution of the task will be paused and user
+          will be dropped into a shell. They can run `exec-task` to resume
+          it (default: True).
 
     Returns:
         int: The exit code from the Docker container.
-
-    Note:
-        Only supports tasks that use 'run-task' and have a payload.image.
-        The task's actual command is made available via an 'exec-task' function
-        in the interactive shell.
     """
     user = user or "worker"
     if isinstance(task, str):
@@ -419,9 +419,9 @@ def load_task(
 
         return 1
 
-    command = task_def["payload"].get("command")  # type: ignore
-    if not command or not command[0].endswith("run-task"):
-        logger.error("Only tasks using `run-task` are supported!")
+    task_command = task_def["payload"].get("command")  # type: ignore
+    if interactive and (not task_command or not task_command[0].endswith("run-task")):
+        logger.error("Only tasks using `run-task` are supported with interactive!")
         return 1
 
     try:
@@ -431,33 +431,40 @@ def load_task(
         logger.exception(e)
         return 1
 
-    # Remove the payload section of the task's command. This way run-task will
-    # set up the task (clone repos, download fetches, etc) but won't actually
-    # start the core of the task. Instead we'll drop the user into an interactive
-    # shell and provide the ability to resume the task command.
-    task_command = None
-    if index := _index(command, "--"):
-        task_command = shlex.join(command[index + 1 :])
-        # I attempted to run the interactive bash shell here, but for some
-        # reason when executed through `run-task`, the interactive shell
-        # doesn't work well. There's no shell prompt on newlines and tab
-        # completion doesn't work. That's why it is executed outside of
-        # `run-task` below, and why we need to parse `--task-cwd`.
-        command[index + 1 :] = [
-            "echo",
-            "Task setup complete!\nRun `exec-task` to execute the task's command.",
-        ]
-
-    # Parse `--task-cwd` so we know where to execute the task's command later.
-    if index := _index(command, "--task-cwd"):
-        task_cwd = command[index + 1]
-    else:
-        for arg in command:
-            if arg.startswith("--task-cwd="):
-                task_cwd = arg.split("=", 1)[1]
-                break
+    exec_command = task_cwd = None
+    if interactive:
+        # Remove the payload section of the task's command. This way run-task will
+        # set up the task (clone repos, download fetches, etc) but won't actually
+        # start the core of the task. Instead we'll drop the user into an interactive
+        # shell and provide the ability to resume the task command.
+        if index := _index(task_command, "--"):
+            exec_command = shlex.join(task_command[index + 1 :])
+            # I attempted to run the interactive bash shell here, but for some
+            # reason when executed through `run-task`, the interactive shell
+            # doesn't work well. There's no shell prompt on newlines and tab
+            # completion doesn't work. That's why it is executed outside of
+            # `run-task` below, and why we need to parse `--task-cwd`.
+            task_command[index + 1 :] = [
+                "echo",
+                "Task setup complete!\nRun `exec-task` to execute the task's command.",
+            ]
+
+        # Parse `--task-cwd` so we know where to execute the task's command later.
+        if index := _index(task_command, "--task-cwd"):
+            task_cwd = task_command[index + 1]
         else:
-            task_cwd = "$TASK_WORKDIR"
+            for arg in task_command:
+                if arg.startswith("--task-cwd="):
+                    task_cwd = arg.split("=", 1)[1]
+                    break
+            else:
+                task_cwd = "$TASK_WORKDIR"
+
+        task_command = [
+            "bash",
+            "-c",
+            f"{shlex.join(task_command)} && cd $TASK_WORKDIR && su -p {user}",
+        ]
 
     # Set some env vars the worker would normally set.
     env = {
@@ -476,17 +483,21 @@ def load_task(
 
     envfile = None
     initfile = None
+    isatty = os.isatty(sys.stdin.fileno())
     try:
         command = [
             "docker",
             "run",
-            "-it",
-            image_tag,
-            "bash",
-            "-c",
-            f"{shlex.join(command)} && cd $TASK_WORKDIR && su -p {user}",
+            "-i",
         ]
 
+        if isatty:
+            command.append("-t")
+
+        command.append(image_tag)
+
+        if task_command:
+            command.extend(task_command)
         if remove:
             command.insert(2, "--rm")
 
@@ -497,16 +508,16 @@ def load_task(
 
             command.insert(2, f"--env-file={envfile.name}")
 
-        if task_command:
+        if exec_command:
             initfile = tempfile.NamedTemporaryFile("w+", delete=False)
             os.fchmod(initfile.fileno(), 0o644)
             initfile.write(
                 dedent(
                     f"""
             function exec-task() {{
-                echo Starting task: {shlex.quote(task_command)}
+                echo Starting task: {shlex.quote(exec_command)}
                 pushd {task_cwd}
-                {task_command}
+                {exec_command}
                 popd
             }}
             """
@@ -516,6 +527,7 @@ def load_task(
 
             command[2:2] = ["-v", f"{initfile.name}:/builds/worker/.bashrc"]
 
+        logger.info(f"Running: {' '.join(command)}")
         proc = subprocess.run(command)
     finally:
         if envfile:
diff --git a/test/test_docker.py b/test/test_docker.py
index 81ffdc29..0f5a5823 100644
--- a/test/test_docker.py
+++ b/test/test_docker.py
@@ -110,7 +110,9 @@ def mock_run(*popenargs, check=False, **kwargs):
 
 @pytest.fixture
 def run_load_task(mocker):
-    def inner(task, remove=False, custom_image=None, pass_task_def=False):
+    def inner(
+        task, remove=False, custom_image=None, pass_task_def=False, interactive=True
+    ):
         proc = mocker.MagicMock()
         proc.returncode = 0
 
@@ -134,6 +136,12 @@ def inner(task, remove=False, custom_image=None, pass_task_def=False):
             ),
         }
 
+        # Mock sys.stdin.fileno() to avoid issues in pytest
+        mock_stdin = mocker.MagicMock()
+        mock_stdin.fileno.return_value = 0
+        mocker.patch.object(docker.sys, "stdin", mock_stdin)
+        mocker.patch.object(docker.os, "isatty", return_value=True)
+
         # If testing with task ID, mock get_task_definition
         if not pass_task_def:
             task_id = "abc"
@@ -146,7 +154,11 @@ def inner(task, remove=False, custom_image=None, pass_task_def=False):
             input_arg = task
 
         ret = docker.load_task(
-            graph_config, input_arg, remove=remove, custom_image=custom_image
+            graph_config,
+            input_arg,
+            remove=remove,
+            custom_image=custom_image,
+            interactive=interactive,
         )
         return ret, mocks
 
@@ -199,7 +211,8 @@ def test_load_task(run_load_task):
         "-v",
         re.compile(f"{tempfile.gettempdir()}/tmp.*:/builds/worker/.bashrc"),
         re.compile(f"--env-file={tempfile.gettempdir()}/tmp.*"),
-        "-it",
+        "-i",
+        "-t",
         "image/tag",
         "bash",
         "-c",
@@ -389,6 +402,43 @@ def test_load_task_with_task_definition(run_load_task, caplog):
     assert "Loading 'test-task-direct' from provided definition" in caplog.text
 
 
+def test_load_task_with_interactive_false(run_load_task):
+    # Test non-interactive mode that doesn't require run-task
+    # Task that doesn't use run-task (would fail in interactive mode)
+    task = {
+        "metadata": {"name": "test-task-non-interactive"},
+        "payload": {
+            "command": ["echo", "hello world"],
+            "image": {"taskId": "def", "type": "task-image"},
+        },
+    }
+
+    # Test with interactive=False - should succeed
+    ret, mocks = run_load_task(task, pass_task_def=True, interactive=False)
+    assert ret == 0
+
+    # Verify subprocess was called
+    mocks["subprocess_run"].assert_called_once()
+    command = mocks["subprocess_run"].call_args[0][0]
+
+    # Should run the task command directly
+    # Find and remove --env-file arg as it contains a tempdir
+    for i, arg in enumerate(command):
+        if arg.startswith("--env-file="):
+            del command[i]
+            break
+
+    assert command == [
+        "docker",
+        "run",
+        "-i",
+        "-t",
+        "image/tag",
+        "echo",
+        "hello world",
+    ]
+
+
 @pytest.fixture
 def task():
     return {
@@ -413,7 +463,22 @@ def test_load_task_with_custom_image_in_tree(run_load_task, task):
 
     mocks["build_image"].assert_called_once()
     args = mocks["subprocess_run"].call_args[0][0]
-    tag = args[args.index("-it") + 1]
+    # Find the image tag - it should be after all docker options and before the command
+    # Structure: ['docker', 'run', ...options..., 'image:tag', ...command...]
+    image_index = None
+    for i, arg in enumerate(args):
+        if (
+            not arg.startswith("-")
+            and not arg.startswith("/")
+            and arg != "docker"
+            and arg != "run"
+            and ":" in arg
+            and not arg.startswith("/tmp")
+        ):
+            image_index = i
+            break
+    assert image_index is not None, f"Could not find image tag in {args}"
+    tag = args[image_index]
     assert tag == f"taskcluster/{image}:latest"
 
 

From 21dad94493946612c470e8070160649b9859ac5d Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Wed, 24 Sep 2025 10:00:08 -0400
Subject: [PATCH 06/10] feat(load-task)!: add a -i/--interactive flag

BREAKING CHANGE: taskgraph load-task non-interactive by default.

This changes interactive to `False` by default. The reasoning is that
only some tasks (aka run-task based one), support this feature. Rather
than having a complex soup of behaviours.. e.g it is True by default
sometimes, but False for certain tasks.. I think it makes much better UX
to just make it something opt-in for all cases.
---
 src/taskgraph/docker.py |  4 ++--
 src/taskgraph/main.py   | 11 +++++++++++
 test/test_main.py       | 18 +++++++++++++++++-
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py
index 2d243df8..d0b7729d 100644
--- a/src/taskgraph/docker.py
+++ b/src/taskgraph/docker.py
@@ -379,7 +379,7 @@ def load_task(
     remove: bool = True,
     user: Optional[str] = None,
     custom_image: Optional[str] = None,
-    interactive: Optional[bool] = True,
+    interactive: Optional[bool] = False,
 ) -> int:
     """Load and run a task interactively in a Docker container.
 
@@ -396,7 +396,7 @@ def load_task(
         custom_image: A custom image to use instead of the task's image.
         interactive: If True, execution of the task will be paused and user
           will be dropped into a shell. They can run `exec-task` to resume
-          it (default: True).
+          it (default: False).
 
     Returns:
         int: The exit code from the Docker container.
diff --git a/src/taskgraph/main.py b/src/taskgraph/main.py
index d1a0149e..2f925d11 100644
--- a/src/taskgraph/main.py
+++ b/src/taskgraph/main.py
@@ -679,6 +679,16 @@ def image_digest(args):
     help="The task id or definition to load into a docker container. Can use "
     "'-' to read from stdin.",
 )
+@argument(
+    "-i",
+    "--interactive",
+    action="store_true",
+    default=False,
+    help="Setup the task but pause execution before executing its command. "
+    "Repositories will be cloned, environment variables will be set and an"
+    "executable script named `exec-task` will be provided to resume task "
+    "execution. Only supported for `run-task` based tasks.",
+)
 @argument(
     "--keep",
     dest="remove",
@@ -719,6 +729,7 @@ def load_task(args):
     return load_task(
         graph_config,
         args["task"],
+        interactive=args["interactive"],
         remove=args["remove"],
         user=args["user"],
         custom_image=args["image"],
diff --git a/test/test_main.py b/test/test_main.py
index 7327dd04..3e175fa0 100644
--- a/test/test_main.py
+++ b/test/test_main.py
@@ -444,7 +444,7 @@ def inner(args, stdin_data=None):
 
 
 def test_load_task_command(run_load_task):
-    # Test normal task ID
+    # Test normal task ID (default non-interactive)
     result, mocks = run_load_task(["load-task", "task-id-123"])
 
     assert result == 0
@@ -453,6 +453,20 @@ def test_load_task_command(run_load_task):
     mocks["docker_load_task"].assert_called_once_with(
         mocks["graph_config"],
         "task-id-123",
+        interactive=False,
+        remove=True,
+        user=None,
+        custom_image=None,
+    )
+
+    # Test with interactive flag
+    result, mocks = run_load_task(["load-task", "-i", "task-id-456"])
+
+    assert result == 0
+    mocks["docker_load_task"].assert_called_once_with(
+        mocks["graph_config"],
+        "task-id-456",
+        interactive=True,
         remove=True,
         user=None,
         custom_image=None,
@@ -476,6 +490,7 @@ def test_load_task_command_with_stdin(run_load_task):
     mocks["docker_load_task"].assert_called_once_with(
         mocks["graph_config"],
         task_def,
+        interactive=False,
         remove=True,
         user=None,
         custom_image=None,
@@ -492,6 +507,7 @@ def test_load_task_command_with_task_id(run_load_task):
     mocks["docker_load_task"].assert_called_once_with(
         mocks["graph_config"],
         "task-id-from-stdin",
+        interactive=False,
         remove=True,
         user=None,
         custom_image=None,

From b9c4b3ae473447f6ddb80411e8f07fe4bfa52d1b Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Tue, 9 Sep 2025 13:07:12 -0400
Subject: [PATCH 07/10] feat(load-task): support specifying volumes

---
 src/taskgraph/docker.py |  6 ++++++
 test/test_docker.py     | 16 ++++++++++++++--
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py
index d0b7729d..a390ef57 100644
--- a/src/taskgraph/docker.py
+++ b/src/taskgraph/docker.py
@@ -380,6 +380,7 @@ def load_task(
     user: Optional[str] = None,
     custom_image: Optional[str] = None,
     interactive: Optional[bool] = False,
+    volumes: Optional[Dict[str, str]] = None,
 ) -> int:
     """Load and run a task interactively in a Docker container.
 
@@ -498,6 +499,11 @@ def load_task(
 
         if task_command:
             command.extend(task_command)
+
+        if volumes:
+            for k, v in volumes.items():
+                command[2:2] = ["-v", f"{k}:{v}"]
+
         if remove:
             command.insert(2, "--rm")
 
diff --git a/test/test_docker.py b/test/test_docker.py
index 0f5a5823..282c33ab 100644
--- a/test/test_docker.py
+++ b/test/test_docker.py
@@ -111,7 +111,12 @@ def mock_run(*popenargs, check=False, **kwargs):
 @pytest.fixture
 def run_load_task(mocker):
     def inner(
-        task, remove=False, custom_image=None, pass_task_def=False, interactive=True
+        task,
+        remove=False,
+        custom_image=None,
+        pass_task_def=False,
+        interactive=True,
+        volumes=None,
     ):
         proc = mocker.MagicMock()
         proc.returncode = 0
@@ -159,6 +164,7 @@ def inner(
             remove=remove,
             custom_image=custom_image,
             interactive=interactive,
+            volumes=volumes,
         )
         return ret, mocks
 
@@ -198,7 +204,9 @@ def test_load_task(run_load_task):
             "image": {"taskId": image_task_id, "type": "task-image"},
         },
     }
-    ret, mocks = run_load_task(task)
+    # Test with custom volumes
+    volumes = {"/host/path": "/container/path", "/another/host": "/another/container"}
+    ret, mocks = run_load_task(task, volumes=volumes)
     assert ret == 0
 
     if "get_task_definition" in mocks:
@@ -211,6 +219,10 @@ def test_load_task(run_load_task):
         "-v",
         re.compile(f"{tempfile.gettempdir()}/tmp.*:/builds/worker/.bashrc"),
         re.compile(f"--env-file={tempfile.gettempdir()}/tmp.*"),
+        "-v",
+        "/another/host:/another/container",
+        "-v",
+        "/host/path:/container/path",
         "-i",
         "-t",
         "image/tag",

From db885fc97d453ad4c0a1a2162a1499bfdf1b719f Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Tue, 9 Sep 2025 13:07:12 -0400
Subject: [PATCH 08/10] feat(build-image)!: use image_builder to build images
 locally

BREAKING CHANGE: The `--tag` flag to `taskgraph build-image` is no
longer supported.
---
 src/taskgraph/docker.py | 152 ++++++++++++++++----------
 src/taskgraph/main.py   |  20 ++--
 test/test_docker.py     | 233 ++++++++++++++++++++++++++--------------
 3 files changed, 252 insertions(+), 153 deletions(-)

diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py
index a390ef57..d3163cd6 100644
--- a/src/taskgraph/docker.py
+++ b/src/taskgraph/docker.py
@@ -4,7 +4,9 @@
 
 import logging
 import os
+import re
 import shlex
+import shutil
 import subprocess
 import sys
 import tarfile
@@ -12,7 +14,11 @@
 from io import BytesIO
 from pathlib import Path
 from textwrap import dedent
-from typing import Any, Dict, Generator, List, Mapping, Optional, Union
+from typing import Any, Dict, Generator, List, Optional, Union
+
+from requests import HTTPError
+
+from taskgraph.generator import load_tasks_for_kind
 
 try:
     import zstandard as zstd
@@ -20,6 +26,7 @@
     zstd = e
 
 from taskgraph.config import GraphConfig
+from taskgraph.transforms.docker_image import IMAGE_BUILDER_IMAGE
 from taskgraph.util import docker, json
 from taskgraph.util.taskcluster import (
     find_task_id,
@@ -27,6 +34,7 @@
     get_root_url,
     get_session,
     get_task_definition,
+    status_task,
 )
 
 logger = logging.getLogger(__name__)
@@ -120,56 +128,27 @@ def load_image_by_task_id(task_id: str, tag: Optional[str] = None) -> str:
     return tag
 
 
-def build_context(
-    name: str,
-    outputFile: str,
-    graph_config: GraphConfig,
-    args: Optional[Mapping[str, str]] = None,
-) -> None:
-    """Build a context.tar for image with specified name.
-
-    Creates a Docker build context tar file for the specified image,
-    which can be used to build the Docker image.
-
-    Args:
-        name: The name of the Docker image to build context for.
-        outputFile: Path to the output tar file to create.
-        graph_config: The graph configuration object.
-        args: Optional mapping of arguments to pass to context creation.
-
-    Raises:
-        ValueError: If name or outputFile is not provided.
-        Exception: If the image directory does not exist.
-    """
-    if not name:
-        raise ValueError("must provide a Docker image name")
-    if not outputFile:
-        raise ValueError("must provide a outputFile")
-
-    image_dir = docker.image_path(name, graph_config)
-    if not os.path.isdir(image_dir):
-        raise Exception(f"image directory does not exist: {image_dir}")
-
-    docker.create_context_tar(".", image_dir, outputFile, args)
-
-
 def build_image(
-    name: str,
-    tag: Optional[str],
     graph_config: GraphConfig,
-    args: Optional[Mapping[str, str]] = None,
-) -> None:
+    name: str,
+    context_file: Optional[str] = None,
+    save_image: Optional[str] = None,
+) -> str:
     """Build a Docker image of specified name.
 
-    Builds a Docker image from the specified image directory and optionally
-    tags it. Output from image building process will be printed to stdout.
+    Builds a Docker image from the specified image directory.
 
     Args:
-        name: The name of the Docker image to build.
-        tag: Optional tag for the built image. If not provided, uses
-            the default tag from docker_image().
         graph_config: The graph configuration.
-        args: Optional mapping of arguments to pass to the build process.
+        name: The name of the Docker image to build.
+        context_file: Path to save the docker context to. If specified,
+            only the context is generated and the image isn't built.
+        save_image: If specified, the resulting `image.tar` will be saved to
+            the specified path. Otherwise, the image is loaded into docker.
+
+    Returns:
+        str: The tag of the loaded image, or absolute path to the image
+            if save_image is specified.
 
     Raises:
         ValueError: If name is not provided.
@@ -183,19 +162,78 @@ def build_image(
     if not os.path.isdir(image_dir):
         raise Exception(f"image directory does not exist: {image_dir}")
 
-    tag = tag or docker.docker_image(name, by_tag=True)
+    label = f"docker-image-{name}"
+    image_tasks = load_tasks_for_kind(
+        {"do_not_optimize": [label]},
+        "docker-image",
+        graph_attr="morphed_task_graph",
+        write_artifacts=True,
+    )
 
-    buf = BytesIO()
-    docker.stream_context_tar(".", image_dir, buf, args)
-    cmdargs = ["docker", "image", "build", "--no-cache", "-"]
-    if tag:
-        cmdargs.insert(-1, f"-t={tag}")
-    subprocess.run(cmdargs, input=buf.getvalue(), check=True)
+    image_context = Path(f"docker-contexts/{name}.tar.gz").resolve()
+    if context_file:
+        shutil.move(image_context, context_file)
+        return ""
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_dir = Path(temp_dir)
+        output_dir = temp_dir / "artifacts"
+        output_dir.mkdir()
+        volumes = {
+            # TODO write artifacts to tmpdir
+            str(output_dir): "/workspace/out",
+            str(image_context): "/workspace/context.tar.gz",
+        }
+
+        assert label in image_tasks
+        task = image_tasks[label]
+        task_def = task.task
+
+        # If the image we're building has a parent image, it may need to be
+        # rebuilt as well if it's cached_task hash changed.
+        if parent_id := task_def["payload"].get("env", {}).get("PARENT_TASK_ID"):
+            try:
+                status_task(parent_id)
+            except HTTPError as e:
+                if e.response.status_code != 404:
+                    raise
+
+                # Parent id doesn't exist, needs to be re-built as well.
+                parent = task.dependencies["parent"][len("docker-image-") :]
+                parent_tar = temp_dir / "parent.tar"
+                build_image(graph_config, parent, save_image=str(parent_tar))
+                volumes[str(parent_tar)] = "/workspace/parent.tar"
+
+        task_def["payload"]["env"]["CHOWN_OUTPUT"] = f"{os.getuid()}:{os.getgid()}"
+        load_task(
+            graph_config,
+            task_def,
+            custom_image=IMAGE_BUILDER_IMAGE,
+            interactive=False,
+            volumes=volumes,
+        )
+        logger.info(f"Successfully built {name} image")
 
-    msg = f"Successfully built {name}"
-    if tag:
-        msg += f" and tagged with {tag}"
-    logger.info(msg)
+        image_tar = output_dir / "image.tar"
+        if save_image:
+            result = Path(save_image).resolve()
+            shutil.copy(image_tar, result)
+        else:
+            proc = subprocess.run(
+                ["docker", "load", "-i", str(image_tar)],
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+            logger.info(proc.stdout)
+
+            m = re.match(r"^Loaded image: (\S+)$", proc.stdout)
+            if m:
+                result = m.group(1)
+            else:
+                result = f"{name}:latest"
+
+    return str(result)
 
 
 def load_image(
@@ -356,9 +394,7 @@ def _resolve_image(image: Union[str, Dict[str, str]], graph_config: GraphConfig)
         # if so build it.
         image_dir = docker.image_path(image, graph_config)
         if Path(image_dir).is_dir():
-            tag = f"taskcluster/{image}:latest"
-            build_image(image, tag, graph_config, os.environ)
-            return tag
+            return build_image(graph_config, image)
 
         # Check if we're referencing a task or index.
         if image.startswith("task-id="):
diff --git a/src/taskgraph/main.py b/src/taskgraph/main.py
index 2f925d11..6113ac2f 100644
--- a/src/taskgraph/main.py
+++ b/src/taskgraph/main.py
@@ -571,9 +571,6 @@ def show_taskgraph(options):
     default="taskcluster",
     help="Relative path to the root of the Taskgraph definition.",
 )
-@argument(
-    "-t", "--tag", help="tag that the image should be built as.", metavar="name:tag"
-)
 @argument(
     "--context-only",
     help="File name the context tarball should be written to."
@@ -582,19 +579,16 @@ def show_taskgraph(options):
 )
 def build_image(args):
     from taskgraph.config import load_graph_config  # noqa: PLC0415
-    from taskgraph.docker import build_context, build_image  # noqa: PLC0415
+    from taskgraph.docker import build_image  # noqa: PLC0415
 
     validate_docker()
+    graph_config = load_graph_config(args["root"])
 
-    root = args["root"]
-    graph_config = load_graph_config(root)
-
-    if args["context_only"] is None:
-        build_image(args["image_name"], args["tag"], graph_config, os.environ)
-    else:
-        build_context(
-            args["image_name"], args["context_only"], graph_config, os.environ
-        )
+    return build_image(
+        graph_config,
+        args["image_name"],
+        args["context_only"],
+    )
 
 
 @command(
diff --git a/test/test_docker.py b/test/test_docker.py
index 282c33ab..2e7ccc12 100644
--- a/test/test_docker.py
+++ b/test/test_docker.py
@@ -6,6 +6,7 @@
 
 from taskgraph import docker
 from taskgraph.config import GraphConfig
+from taskgraph.transforms.docker_image import IMAGE_BUILDER_IMAGE
 
 
 @pytest.fixture(autouse=True, scope="module")
@@ -27,87 +28,6 @@ def side_effect(topsrcdir, context_dir, out_file, image_name=None, args=None):
     return (m_stream, m_run)
 
 
-def test_build_image(caplog, mock_docker_build):
-    caplog.set_level(logging.DEBUG)
-    m_stream, m_run = mock_docker_build
-    image = "hello-world-tag"
-    tag = f"test/{image}:1.0"
-
-    graph_config = GraphConfig(
-        {
-            "trust-domain": "test-domain",
-            "docker-image-kind": "docker-image",
-        },
-        "test/data/taskcluster",
-    )
-
-    assert docker.build_image(image, None, graph_config=graph_config) is None
-    m_stream.assert_called_once()
-    m_run.assert_called_once_with(
-        ["docker", "image", "build", "--no-cache", f"-t={tag}", "-"],
-        input=b"xyz",
-        check=True,
-    )
-
-    assert f"Successfully built {image} and tagged with {tag}" in caplog.text
-
-
-def test_build_image_no_tag(caplog, mock_docker_build):
-    caplog.set_level(logging.DEBUG)
-    m_stream, m_run = mock_docker_build
-    image = "hello-world"
-
-    graph_config = GraphConfig(
-        {
-            "trust-domain": "test-domain",
-            "docker-image-kind": "docker-image",
-        },
-        "test/data/taskcluster",
-    )
-
-    assert docker.build_image(image, None, graph_config=graph_config) is None
-    m_stream.assert_called_once()
-    m_run.assert_called_once_with(
-        ["docker", "image", "build", "--no-cache", "-"],
-        input=b"xyz",
-        check=True,
-    )
-
-    assert f"Successfully built {image}" in caplog.text
-
-
-def test_build_image_error(caplog, mock_docker_build):
-    caplog.set_level(logging.DEBUG)
-    m_stream, m_run = mock_docker_build
-
-    def mock_run(*popenargs, check=False, **kwargs):
-        if check:
-            raise docker.subprocess.CalledProcessError(1, popenargs)
-        return 1
-
-    m_run.side_effect = mock_run
-    image = "hello-world"
-
-    graph_config = GraphConfig(
-        {
-            "trust-domain": "test-domain",
-            "docker-image-kind": "docker-image",
-        },
-        "test/data/taskcluster",
-    )
-
-    with pytest.raises(Exception):
-        docker.build_image(image, None, graph_config=graph_config)
-    m_stream.assert_called_once()
-    m_run.assert_called_once_with(
-        ["docker", "image", "build", "--no-cache", "-"],
-        input=b"xyz",
-        check=True,
-    )
-
-    assert f"Successfully built {image}" not in caplog.text
-
-
 @pytest.fixture
 def run_load_task(mocker):
     def inner(
@@ -131,7 +51,7 @@ def inner(
 
         mocks = {
             "build_image": mocker.patch.object(
-                docker, "build_image", return_value=None
+                docker, "build_image", return_value="taskcluster/hello-world:latest"
             ),
             "load_image_by_task_id": mocker.patch.object(
                 docker, "load_image_by_task_id", return_value="image/tag"
@@ -515,3 +435,152 @@ def test_load_task_with_custom_image_registry(mocker, run_load_task, task):
     assert ret == 0
     assert not mocks["load_image_by_task_id"].called
     assert not mocks["build_image"].called
+
+
+@pytest.fixture
+def run_build_image(mocker):
+    def inner(image_name, save_image=None, context_file=None):
+        graph_config = GraphConfig(
+            {
+                "trust-domain": "test-domain",
+                "docker-image-kind": "docker-image",
+            },
+            "test/data/taskcluster",
+        )
+
+        # Mock the TemporaryDirectory context manager since the current build_image uses it
+        temp_dir_mock = mocker.MagicMock()
+        temp_dir_str = "/tmp/test_temp_dir"
+        temp_dir_mock.__enter__ = mocker.MagicMock(return_value=temp_dir_str)
+        temp_dir_mock.__exit__ = mocker.MagicMock(return_value=False)
+
+        # Mock Path objects
+        temp_dir_path = mocker.MagicMock()
+        output_dir = mocker.MagicMock()
+
+        # Mock Path constructor
+        original_path = docker.Path
+
+        def mock_path_constructor(path_arg):
+            if str(path_arg) == temp_dir_str:
+                return temp_dir_path
+            elif str(path_arg).endswith(".tar.gz"):
+                # This is the image_context path
+                image_context_mock = mocker.MagicMock()
+                image_context_mock.resolve.return_value = image_context_mock
+                return image_context_mock
+            elif save_image and str(path_arg) == save_image:
+                save_path_mock = mocker.MagicMock()
+                save_path_mock.resolve.return_value = save_path_mock
+                save_path_mock.__str__ = lambda self: save_image
+                return save_path_mock
+            return original_path(path_arg)
+
+        # Set up directory operations
+        temp_dir_path.__truediv__ = mocker.MagicMock(return_value=output_dir)
+        output_dir.mkdir = mocker.MagicMock()
+        output_dir.__truediv__ = mocker.MagicMock()
+
+        # Initialize mocks dictionary
+        mocks = {
+            "TemporaryDirectory": mocker.patch.object(
+                docker.tempfile, "TemporaryDirectory", return_value=temp_dir_mock
+            ),
+            "Path": mocker.patch.object(
+                docker, "Path", side_effect=mock_path_constructor
+            ),
+            "load_tasks_for_kind": mocker.patch.object(docker, "load_tasks_for_kind"),
+            "load_task": mocker.patch.object(docker, "load_task"),
+            "subprocess": mocker.patch.object(docker.subprocess, "run"),
+            "shutil_copy": mocker.patch.object(docker.shutil, "copy"),
+            "shutil_move": mocker.patch.object(docker.shutil, "move"),
+            "status_task": mocker.patch.object(docker, "status_task"),
+            "isdir": mocker.patch.object(docker.os.path, "isdir", return_value=True),
+            "getuid": mocker.patch.object(docker.os, "getuid", return_value=1000),
+            "getgid": mocker.patch.object(docker.os, "getgid", return_value=1000),
+        }
+
+        # Mock image task
+        mocks["task"] = mocker.MagicMock()
+        mocks["task"].task = {"payload": {"env": {}}}
+        mocks["load_tasks_for_kind"].return_value = {
+            f"docker-image-{image_name}": mocks["task"]
+        }
+
+        # Mock subprocess result for docker load
+        mocks["proc_result"] = mocker.MagicMock()
+        mocks["proc_result"].stdout = f"Loaded image: {image_name}:latest"
+        mocks["subprocess"].return_value = mocks["proc_result"]
+
+        # Add convenience references
+        mocks["graph_config"] = graph_config
+        mocks["temp_dir_path"] = temp_dir_path
+        mocks["output_dir"] = output_dir
+
+        # Run the build_image function
+        result = docker.build_image(
+            graph_config, image_name, context_file=context_file, save_image=save_image
+        )
+
+        return result, mocks
+
+    return inner
+
+
+def test_build_image(run_build_image):
+    # Test building image without save_image
+    result, mocks = run_build_image("hello-world")
+
+    # Verify TemporaryDirectory is used for cleanup
+    mocks["TemporaryDirectory"].assert_called_once()
+
+    # Verify the function calls
+    mocks["load_tasks_for_kind"].assert_called_once_with(
+        {"do_not_optimize": ["docker-image-hello-world"]},
+        "docker-image",
+        graph_attr="morphed_task_graph",
+        write_artifacts=True,
+    )
+
+    mocks["load_task"].assert_called_once()
+    call_args = mocks["load_task"].call_args
+    assert call_args[0][0] == mocks["graph_config"]
+    assert call_args[0][1] == mocks["task"].task
+    assert call_args[1]["custom_image"] == IMAGE_BUILDER_IMAGE
+    assert call_args[1]["interactive"] is False
+    assert "volumes" in call_args[1]
+
+    # Verify docker load was called
+    mocks["subprocess"].assert_called_once()
+    docker_load_args = mocks["subprocess"].call_args[0][0]
+    assert docker_load_args[:3] == ["docker", "load", "-i"]
+
+    assert result == "hello-world:latest"
+
+
+def test_build_image_with_save_image(run_build_image):
+    save_path = "/path/to/save.tar"
+
+    # Test building image with save_image option
+    result, mocks = run_build_image("test", save_image=save_path)
+
+    # Verify TemporaryDirectory is used for cleanup
+    mocks["TemporaryDirectory"].assert_called_once()
+
+    # Verify copy was called instead of docker load
+    mocks["shutil_copy"].assert_called_once()
+
+    # Result should be the string representation of the save path
+    assert save_path in str(result)
+
+
+def test_build_image_context_only(run_build_image):
+    context_path = "/path/to/context.tar"
+
+    # Test building only the context file
+    result, mocks = run_build_image("context-test", context_file=context_path)
+
+    # Verify move was called for the context file
+    mocks["shutil_move"].assert_called_once()
+
+    assert result == ""

From 100bc63e3779d4c5321b5223674dc26eddf1754b Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Thu, 25 Sep 2025 11:46:10 -0400
Subject: [PATCH 09/10] fix(load-task): ensure we don't leave docker-contexts
 in cwd

---
 src/taskgraph/docker.py | 41 +++++++++++++++++++++++++----------------
 1 file changed, 25 insertions(+), 16 deletions(-)

diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py
index d3163cd6..bdccf42d 100644
--- a/src/taskgraph/docker.py
+++ b/src/taskgraph/docker.py
@@ -26,7 +26,7 @@
     zstd = e
 
 from taskgraph.config import GraphConfig
-from taskgraph.transforms.docker_image import IMAGE_BUILDER_IMAGE
+from taskgraph.transforms import docker_image
 from taskgraph.util import docker, json
 from taskgraph.util.taskcluster import (
     find_task_id,
@@ -162,22 +162,31 @@ def build_image(
     if not os.path.isdir(image_dir):
         raise Exception(f"image directory does not exist: {image_dir}")
 
-    label = f"docker-image-{name}"
-    image_tasks = load_tasks_for_kind(
-        {"do_not_optimize": [label]},
-        "docker-image",
-        graph_attr="morphed_task_graph",
-        write_artifacts=True,
-    )
-
-    image_context = Path(f"docker-contexts/{name}.tar.gz").resolve()
-    if context_file:
-        shutil.move(image_context, context_file)
-        return ""
-
     with tempfile.TemporaryDirectory() as temp_dir:
         temp_dir = Path(temp_dir)
-        output_dir = temp_dir / "artifacts"
+
+        # Ensure the docker-image transforms save the docker-contexts to our
+        # temp_dir
+        label = f"docker-image-{name}"
+        contexts_dir = temp_dir / "docker-contexts"
+        old_contexts_dir = docker_image.CONTEXTS_DIR
+        try:
+            docker_image.CONTEXTS_DIR = str(contexts_dir)
+            image_tasks = load_tasks_for_kind(
+                {"do_not_optimize": [label]},
+                "docker-image",
+                graph_attr="morphed_task_graph",
+                write_artifacts=True,
+            )
+        finally:
+            docker_image.CONTEXTS_DIR = old_contexts_dir
+
+        image_context = contexts_dir.joinpath(f"{name}.tar.gz").resolve()
+        if context_file:
+            shutil.move(image_context, context_file)
+            return ""
+
+        output_dir = temp_dir / "out"
         output_dir.mkdir()
         volumes = {
             # TODO write artifacts to tmpdir
@@ -208,7 +217,7 @@ def build_image(
         load_task(
             graph_config,
             task_def,
-            custom_image=IMAGE_BUILDER_IMAGE,
+            custom_image=docker_image.IMAGE_BUILDER_IMAGE,
             interactive=False,
             volumes=volumes,
         )

From fd68aa1868b138e055cd47c288b10899642370e9 Mon Sep 17 00:00:00 2001
From: Andrew Halberstadt <ahal@mozilla.com>
Date: Fri, 26 Sep 2025 10:35:33 -0400
Subject: [PATCH 10/10] refactor(load-task): build docker command in the
 logical order

---
 src/taskgraph/docker.py | 24 ++++++++++++------------
 test/test_docker.py     | 16 ++++++++--------
 2 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/src/taskgraph/docker.py b/src/taskgraph/docker.py
index bdccf42d..18232252 100644
--- a/src/taskgraph/docker.py
+++ b/src/taskgraph/docker.py
@@ -540,24 +540,15 @@ def load_task(
         if isatty:
             command.append("-t")
 
-        command.append(image_tag)
-
-        if task_command:
-            command.extend(task_command)
-
-        if volumes:
-            for k, v in volumes.items():
-                command[2:2] = ["-v", f"{k}:{v}"]
-
         if remove:
-            command.insert(2, "--rm")
+            command.append("--rm")
 
         if env:
             envfile = tempfile.NamedTemporaryFile("w+", delete=False)
             envfile.write("\n".join([f"{k}={v}" for k, v in env.items()]))
             envfile.close()
 
-            command.insert(2, f"--env-file={envfile.name}")
+            command.append(f"--env-file={envfile.name}")
 
         if exec_command:
             initfile = tempfile.NamedTemporaryFile("w+", delete=False)
@@ -576,7 +567,16 @@ def load_task(
             )
             initfile.close()
 
-            command[2:2] = ["-v", f"{initfile.name}:/builds/worker/.bashrc"]
+            command.extend(["-v", f"{initfile.name}:/builds/worker/.bashrc"])
+
+        if volumes:
+            for k, v in volumes.items():
+                command.extend(["-v", f"{k}:{v}"])
+
+        command.append(image_tag)
+
+        if task_command:
+            command.extend(task_command)
 
         logger.info(f"Running: {' '.join(command)}")
         proc = subprocess.run(command)
diff --git a/test/test_docker.py b/test/test_docker.py
index 2e7ccc12..1481f95d 100644
--- a/test/test_docker.py
+++ b/test/test_docker.py
@@ -136,15 +136,15 @@ def test_load_task(run_load_task):
     expected = [
         "docker",
         "run",
-        "-v",
-        re.compile(f"{tempfile.gettempdir()}/tmp.*:/builds/worker/.bashrc"),
+        "-i",
+        "-t",
         re.compile(f"--env-file={tempfile.gettempdir()}/tmp.*"),
         "-v",
-        "/another/host:/another/container",
+        re.compile(f"{tempfile.gettempdir()}/tmp.*:/builds/worker/.bashrc"),
         "-v",
         "/host/path:/container/path",
-        "-i",
-        "-t",
+        "-v",
+        "/another/host:/another/container",
         "image/tag",
         "bash",
         "-c",
@@ -239,9 +239,9 @@ def test_load_task_env_init_and_remove(mocker, run_load_task):
     # Verify subprocess was called with the correct env file and init file
     mocks["subprocess_run"].assert_called_once()
     actual = mocks["subprocess_run"].call_args[0][0]
-    assert actual[3] == "/tmp/test_initfile:/builds/worker/.bashrc"
-    assert actual[4] == "--env-file=/tmp/test_envfile"
-    assert actual[5] == "--rm"
+    assert actual[4] == "--rm"
+    assert actual[5] == "--env-file=/tmp/test_envfile"
+    assert actual[6:8] == ["-v", "/tmp/test_initfile:/builds/worker/.bashrc"]
 
 
 @pytest.mark.parametrize(