diff --git a/dk-installer.py b/dk-installer.py
index 02f2ba5..01b602d 100755
--- a/dk-installer.py
+++ b/dk-installer.py
@@ -67,6 +67,10 @@
 TESTGEN_LOG_FILE_PATH = pathlib.Path.home() / ".testgen" / "logs" / "app.log"
 TESTGEN_CONFIG_ENV_PATH = pathlib.Path.home() / ".testgen" / "config.env"
 TESTGEN_APP_READY_TIMEOUT = 120
+# Seconds TestGen is given to stop. Must exceed TG_JOB_SHUTDOWN_TIMEOUT (default 60s) plus
+# the time the process needs to record what stopped, or the scheduler is killed mid-wait and
+# a running job is cut instead of stopping at a checkpoint.
+TESTGEN_STOP_GRACE_PERIOD = 90
 INSTALL_MARKER_FILE = "dk-{}-install.json"
 INSTALL_MODE_DOCKER = "docker"
 INSTALL_MODE_PIP = "pip"
@@ -2079,6 +2083,32 @@ def execute(self, args):
             CONSOLE.msg("Observability Heartbeat demo stopped.")
 
 
+def find_in_block(contents: str, block: str, key: str) -> typing.Optional[re.Match]:
+    """Find the ``key:`` line inside the compose ``block:`` mapping, or None.
+
+    Scans by indentation rather than parsing YAML — enough for the block-style files the
+    installer writes, and it avoids a runtime dependency. Offsets on the returned match
+    are absolute, so callers can splice around it; group 1 is the key's indent.
+
+    Scoping to a block is the point: the same key can appear on several services, and
+    only ``engine`` runs the scheduler. Deliberately says nothing about *which* image a
+    service uses — ``tg install --image`` accepts any registry.
+    """
+    headers = list(re.finditer(rf"^([ \t]*){re.escape(block)}:[ \t]*$", contents, flags=re.M))
+    if not headers:
+        return None
+    # Shallowest wins: a name like ``postgres`` is both a service and a nested key under
+    # another service's ``depends_on``, and it's the service the caller means.
+    header = min(headers, key=lambda match: len(match.group(1)))
+    # The block body ends at the first line indented no deeper than the block key itself.
+    end = len(contents)
+    for line in re.finditer(r"^([ \t]*)\S.*$", contents[header.end() :], flags=re.M):
+        if len(line.group(1)) <= len(header.group(1)):
+            end = header.end() + line.start()
+            break
+    return re.compile(rf"^([ \t]+){re.escape(key)}:.*$", flags=re.M).search(contents, header.end(), end)
+
+
 class UpdateComposeFileStep(Step):
     label = "Updating the Docker compose file"
 
@@ -2088,6 +2118,7 @@ def __init__(self):
         self.update_token = False
         self.update_base_url = False
         self.update_api_port = False
+        self.update_stop_grace = False
         super().__init__()
 
     def pre_execute(self, action, args):
@@ -2149,6 +2180,17 @@ def pre_execute(self, action, args):
             and not re.search(rf"- \d+:{TESTGEN_DEFAULT_API_PORT}\b", contents)
         )
 
+        # Compose files written before the grace period was added stop the engine after
+        # docker's 10s default, cutting a running job instead of letting it checkpoint.
+        # Only count it as a pending change if we can actually place it, or the step
+        # would report success having silently rewritten the file unchanged.
+        engine_image = find_in_block(contents, "engine", "image")
+        if engine_image is None:
+            LOG.info("No image line in the compose 'engine' service; leaving stop_grace_period alone")
+        self.update_stop_grace = (
+            engine_image is not None and find_in_block(contents, "engine", "stop_grace_period") is None
+        )
+
         if not any(
             (
                 self.update_version,
@@ -2156,6 +2198,7 @@ def pre_execute(self, action, args):
                 self.update_token,
                 self.update_base_url,
                 self.update_api_port,
+                self.update_stop_grace,
             )
         ):
             CONSOLE.msg("No changes will be applied.")
@@ -2169,6 +2212,7 @@ def execute(self, action, args):
                 self.update_token,
                 self.update_base_url,
                 self.update_api_port,
+                self.update_stop_grace,
             )
         ):
             raise SkipStep
@@ -2210,6 +2254,15 @@ def execute(self, action, args):
             new_mapping = f"\n{match.group(1)}- {TESTGEN_DEFAULT_API_PORT}:{TESTGEN_DEFAULT_API_PORT}"
             contents = contents[0 : match.end()] + new_mapping + contents[match.end() :]
 
+        if self.update_stop_grace and (image := find_in_block(contents, "engine", "image")):
+            indent = image.group(1)
+            grace = (
+                f"\n{indent}# Must exceed TG_JOB_SHUTDOWN_TIMEOUT (default 60s), or docker kills the scheduler"
+                f"\n{indent}# mid-wait and running jobs are cut instead of stopping at a checkpoint."
+                f"\n{indent}stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s"
+            )
+            contents = contents[: image.end()] + grace + contents[image.end() :]
+
         action.get_compose_file_path(args).write_text(contents)
 
 
@@ -2309,6 +2362,9 @@ def get_compose_file_contents(self, action, args):
             services:
               engine:
                 image: {args.image}
+                # Must exceed TG_JOB_SHUTDOWN_TIMEOUT (default 60s), or docker kills the scheduler
+                # mid-wait and running jobs are cut instead of stopping at a checkpoint.
+                stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s
                 container_name: testgen
                 environment: *common-variables
                 volumes:
@@ -2524,17 +2580,25 @@ def read_testgen_config_env() -> dict[str, str]:
     return config
 
 
-def stop_app_tree(proc: subprocess.Popen, timeout: int = 10) -> None:
-    """Terminate ``proc`` and all of its descendants.
+def supports_graceful_stop() -> bool:
+    """Whether we can ask the app to stop rather than kill it outright.
 
-    Plain ``proc.terminate()`` only kills the parent — pixeltable-pgserver
-    spawns ``postgres`` children that get orphaned otherwise. Cross-platform:
-    on Windows we shell out to ``taskkill /F /T``; on POSIX we send SIGTERM
-    to the whole process group (the parent was started with
-    ``start_new_session=True``).
+    POSIX only. On Windows the app is started without ``CREATE_NEW_PROCESS_GROUP``,
+    so there is no catchable signal we can deliver to it and ``taskkill /F`` is the
+    only reliable stop — a running job is cut there whatever grace period we would
+    nominally allow. Callers use this to avoid promising a wait that can't happen.
+    """
+    return platform.system() != "Windows"
+
+
+def force_kill_app_tree(proc: subprocess.Popen, timeout: int = 5) -> None:
+    """Kill ``proc`` and every descendant, including those outside its process group.
+
+    ``testgen run-app all`` starts its ui/scheduler children in their own sessions, so
+    they sit outside ``proc``'s process group and outlive a ``killpg`` — killing only
+    the parent leaves the UI holding its port and postgres holding the data directory,
+    which then breaks the next ``tg start``. Hence the orphan sweep to finish the job.
     """
-    if proc.poll() is not None:
-        return
     if platform.system() == "Windows":
         with contextlib.suppress(Exception):
             subprocess.run(
@@ -2546,16 +2610,51 @@ def stop_app_tree(proc: subprocess.Popen, timeout: int = 10) -> None:
             )
     else:
         with contextlib.suppress(Exception):
-            os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
+            os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
+    # Backstop in case the platform kill above didn't land (e.g. taskkill denied).
+    proc.kill()
+    with contextlib.suppress(subprocess.TimeoutExpired):
+        proc.wait(timeout=timeout)
+    stop_standalone_orphans()
+
+
+def stop_app_tree(proc: subprocess.Popen, timeout: int = 10) -> bool:
+    """Terminate ``proc`` and all of its descendants. Returns whether it stopped on its own.
+
+    Plain ``proc.terminate()`` only kills the parent — pixeltable-pgserver
+    spawns ``postgres`` children that get orphaned otherwise. Cross-platform:
+    on Windows we shell out to ``taskkill /F /T``; on POSIX we send SIGTERM
+    to the whole process group (the parent was started with
+    ``start_new_session=True``) and let it forward the signal to its children.
+
+    ``timeout`` is how long the tree is given to shut down cooperatively before
+    it is force-killed — the caller passes ``TESTGEN_STOP_GRACE_PERIOD`` when a
+    running job may need to reach a checkpoint first. A second Ctrl+C during
+    that wait is taken as "stop now" and skips straight to the force-kill.
+
+    Returns ``True`` when the tree exited within ``timeout`` (or was already
+    gone), ``False`` when it had to be force-killed. See ``supports_graceful_stop``
+    for why Windows always reports ``False`` when there was a live process.
+    """
+    if proc.poll() is not None:
+        return True
+
+    if not supports_graceful_stop():
+        force_kill_app_tree(proc, timeout=timeout)
+        return False
+
+    with contextlib.suppress(Exception):
+        os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
     try:
         proc.wait(timeout=timeout)
-    except subprocess.TimeoutExpired:
-        if platform.system() != "Windows":
-            with contextlib.suppress(Exception):
-                os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
-        proc.kill()
-        with contextlib.suppress(subprocess.TimeoutExpired):
-            proc.wait(timeout=5)
+    except (subprocess.TimeoutExpired, KeyboardInterrupt):
+        # KeyboardInterrupt here is a second Ctrl+C while we were waiting. Swallow it
+        # rather than letting it unwind into the caller's ``finally``, which would send
+        # the tree another SIGTERM — TestGen reads a second signal as "hurry up" and
+        # force-kills the job mid-pause, losing exactly the checkpoint we were waiting for.
+        force_kill_app_tree(proc)
+        return False
+    return True
 
 
 def stop_standalone_orphans() -> None:
@@ -2687,8 +2786,21 @@ def start_testgen_app(action, args) -> None:
             # Reset the cursor to column 0 — the terminal echoed `^C` mid-line.
             print("")
             CONSOLE.msg("Stopping TestGen...")
-            stop_app_tree(proc, timeout=10)
-            CONSOLE.msg("TestGen stopped.")
+            graceful = supports_graceful_stop()
+            if graceful:
+                # A running profiling job stops at its next checkpoint rather than being cut,
+                # but that can take up to TG_JOB_SHUTDOWN_TIMEOUT — say so, or the wait reads
+                # as a hang and the user reaches for a second Ctrl+C.
+                CONSOLE.msg(
+                    f"Waiting up to {TESTGEN_STOP_GRACE_PERIOD} seconds for running jobs to reach a checkpoint..."
+                )
+            stopped_cleanly = stop_app_tree(proc, timeout=TESTGEN_STOP_GRACE_PERIOD)
+            # Only worth flagging where we actually offered a grace period. Windows always
+            # force-kills, so the warning would fire on every stop and mean nothing.
+            if graceful and not stopped_cleanly:
+                CONSOLE.msg("TestGen stopped. A job that was still running will restart from the beginning.")
+            else:
+                CONSOLE.msg("TestGen stopped.")
             CONSOLE.msg(f"To start it again, {command_hint(args.prod, 'start', 'Start TestGen')}.")
     finally:
         stop_app_tree(proc, timeout=5)
diff --git a/tests/conftest.py b/tests/conftest.py
index 3d57a0d..41b737a 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -33,6 +33,17 @@ def _no_real_process_group_signals():
         yield killpg_mock
 
 
+@pytest.fixture(autouse=True)
+def _no_real_orphan_sweep():
+    """``force_kill_app_tree`` finishes with ``stop_standalone_orphans``, which shells out
+    to a real ``pkill -9 -f 'testgen.*run-app'``. Unpatched, a test exercising the
+    force-kill path would kill the developer's own running TestGen. Tests that assert on
+    the sweep override this inside their own ``with patch(...)``.
+    """
+    with patch("tests.installer.stop_standalone_orphans") as mock:
+        yield mock
+
+
 @pytest.fixture
 def stdout_mock():
     return Mock(return_value=[])
diff --git a/tests/test_tg_install.py b/tests/test_tg_install.py
index ce2cf5f..85228d8 100644
--- a/tests/test_tg_install.py
+++ b/tests/test_tg_install.py
@@ -10,6 +10,7 @@
     AbortAction,
     TestGenCreateDockerComposeFileStep,
     ComposeVerifyExistingInstallStep,
+    TESTGEN_STOP_GRACE_PERIOD,
 )
 
 
@@ -114,6 +115,24 @@ def test_tg_compose_base_url_custom_port(tg_install_action, start_cmd_mock, stdo
     assert "TG_UI_BASE_URL: http://localhost:9000" in contents
 
 
+@pytest.mark.integration
+def test_tg_compose_sets_engine_stop_grace_period(tg_install_action, start_cmd_mock, stdout_mock, compose_path):
+    """Docker's 10s default kills the scheduler mid-shutdown, cutting a running job
+    instead of letting it stop at a checkpoint."""
+    tg_install_action.execute()
+    compose_content = compose_path.read_text()
+
+    # Only the engine needs it — postgres shuts down on its own quickly.
+    assert compose_content.count("stop_grace_period") == 1
+    lines = compose_content.splitlines()
+    image_idx = next(i for i, line in enumerate(lines) if "image: datakitchen/dataops-testgen" in line)
+    grace_idx = next(i for i, line in enumerate(lines) if "stop_grace_period" in line)
+    # Inside the engine service, right under its image: two comment lines, then the key.
+    assert grace_idx == image_idx + 3
+    indent = lines[image_idx][: -len(lines[image_idx].lstrip())]
+    assert lines[grace_idx] == f"{indent}stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s"
+
+
 @pytest.mark.integration
 def test_tg_compose_base_url_ssl(tg_install_action, start_cmd_mock, stdout_mock, args_mock, compose_path):
     args_mock.ssl_cert_file = "/path/to/cert.crt"
diff --git a/tests/test_tg_start.py b/tests/test_tg_start.py
index ecd86ba..a798764 100644
--- a/tests/test_tg_start.py
+++ b/tests/test_tg_start.py
@@ -1,3 +1,4 @@
+import signal
 from pathlib import Path
 from unittest.mock import MagicMock, patch
 
@@ -11,7 +12,9 @@
     TestgenStartAction,
     start_testgen_app,
     stop_app_tree,
+    force_kill_app_tree,
     InstallMarker,
+    TESTGEN_STOP_GRACE_PERIOD,
 )
 
 
@@ -117,18 +120,66 @@ def test_start_testgen_app_handles_keyboard_interrupt(app_action, args_mock, con
         patch("tests.installer.resolve_testgen_path", return_value="/bin/testgen"),
         patch("tests.installer.subprocess.Popen", return_value=proc),
         patch("tests.installer.wait_for_tcp_port", return_value=True),
-        patch("tests.installer.stop_app_tree") as stop_mock,
+        patch("tests.installer.stop_app_tree", return_value=True) as stop_mock,
     ):
         start_testgen_app(app_action, args_mock)
 
-    # Called once for the keyboard-interrupt branch (timeout=10) and again in
-    # the ``finally`` cleanup (timeout=5; no-op since proc already stopped).
+    # Called once for the keyboard-interrupt branch and again in the ``finally``
+    # cleanup (timeout=5; no-op since proc already stopped). The first wait gets the
+    # full grace period so a running job can reach a checkpoint before being killed.
     assert stop_mock.call_args_list[0].args[0] is proc
-    assert stop_mock.call_args_list[0].kwargs == {"timeout": 10}
-    console_msg_mock.assert_any_msg_contains("TestGen stopped")
+    assert stop_mock.call_args_list[0].kwargs == {"timeout": TESTGEN_STOP_GRACE_PERIOD}
+    console_msg_mock.assert_any_msg_contains("reach a checkpoint")
+    console_msg_mock.assert_any_msg_contains("TestGen stopped.")
     console_msg_mock.assert_any_msg_contains("tg start")
 
 
+@pytest.mark.unit
+def test_start_testgen_app_warns_when_grace_period_expires(app_action, args_mock, console_msg_mock, empty_tg_config):
+    """A tree that had to be force-killed lost its checkpoint — say so rather than
+    reporting a clean stop the user can't trust."""
+    args_mock.prod = "tg"
+
+    proc = MagicMock()
+    proc.poll.return_value = None
+    proc.wait.side_effect = [KeyboardInterrupt(), 0]
+
+    with (
+        patch("tests.installer.resolve_testgen_path", return_value="/bin/testgen"),
+        patch("tests.installer.subprocess.Popen", return_value=proc),
+        patch("tests.installer.wait_for_tcp_port", return_value=True),
+        patch("tests.installer.stop_app_tree", return_value=False),
+    ):
+        start_testgen_app(app_action, args_mock)
+
+    console_msg_mock.assert_any_msg_contains("will restart from the beginning")
+
+
+@pytest.mark.unit
+def test_start_testgen_app_makes_no_grace_promise_on_windows(app_action, args_mock, console_msg_mock, empty_tg_config):
+    """Windows always force-kills, so promising a wait would be a lie and the
+    force-kill warning would fire on every single stop, saying nothing."""
+    args_mock.prod = "tg"
+
+    proc = MagicMock()
+    proc.poll.return_value = None
+    proc.wait.side_effect = [KeyboardInterrupt(), 0]
+
+    with (
+        patch("tests.installer.resolve_testgen_path", return_value="/bin/testgen"),
+        patch("tests.installer.subprocess.Popen", return_value=proc),
+        patch("tests.installer.wait_for_tcp_port", return_value=True),
+        patch("tests.installer.supports_graceful_stop", return_value=False),
+        patch("tests.installer.stop_app_tree", return_value=False),
+    ):
+        start_testgen_app(app_action, args_mock)
+
+    printed = " ".join(str(c) for c in console_msg_mock.call_args_list)
+    assert "reach a checkpoint" not in printed
+    assert "will restart from the beginning" not in printed
+    console_msg_mock.assert_any_msg_contains("TestGen stopped.")
+
+
 # --- stop_app_tree ------------------------------------------------------------
 
 
@@ -197,7 +248,139 @@ def test_stop_app_tree_falls_through_to_kill_on_timeout():
         patch("tests.installer.os.killpg"),
         patch("tests.installer.os.getpgid", return_value=4242),
     ):
-        stop_app_tree(proc, timeout=1)
+        forced = stop_app_tree(proc, timeout=1)
+
+    proc.kill.assert_called_once()
+    assert forced is False  # had to be force-killed
+
+
+@pytest.mark.unit
+def test_stop_app_tree_reports_graceful_stop():
+    """Exiting within the grace period is the signal the caller uses to decide
+    whether a running job got to checkpoint."""
+    proc = MagicMock()
+    proc.poll.return_value = None
+    proc.pid = 4242
+    proc.wait.return_value = 0
+
+    with (
+        patch("tests.installer.platform.system", return_value="Linux"),
+        patch("tests.installer.os.killpg"),
+        patch("tests.installer.os.getpgid", return_value=4242),
+    ):
+        assert stop_app_tree(proc, timeout=3) is True
+
+
+@pytest.mark.unit
+def test_stop_app_tree_swallows_second_interrupt():
+    """A second Ctrl+C while we wait means 'stop now'. It must force-kill here rather
+    than escaping to the caller's ``finally``, which would re-signal the tree — TestGen
+    reads a second signal as 'hurry up' and cuts the job mid-pause."""
+    proc = MagicMock()
+    proc.poll.return_value = None
+    proc.pid = 4242
+    proc.wait.side_effect = [KeyboardInterrupt(), 0]
+
+    with (
+        patch("tests.installer.platform.system", return_value="Linux"),
+        patch("tests.installer.os.killpg") as killpg_mock,
+        patch("tests.installer.os.getpgid", return_value=4242),
+    ):
+        forced = stop_app_tree(proc, timeout=90)
+
+    assert forced is False
+    proc.kill.assert_called_once()
+    # SIGTERM first, then the escalation to SIGKILL.
+    assert [c.args[1] for c in killpg_mock.call_args_list] == [signal.SIGTERM, signal.SIGKILL]
+
+
+@pytest.mark.unit
+def test_stop_app_tree_windows_stop_is_always_forced():
+    """``taskkill /F`` gives the tree no chance to checkpoint, so the caller must not be
+    told the stop was graceful."""
+    proc = MagicMock()
+    proc.poll.return_value = None
+    proc.pid = 4242
+    proc.wait.return_value = 0
+
+    with (
+        patch("tests.installer.platform.system", return_value="Windows"),
+        patch("tests.installer.subprocess.run"),
+    ):
+        assert stop_app_tree(proc, timeout=90) is False
+
+
+@pytest.mark.unit
+def test_force_kill_sweeps_orphans_outside_the_process_group():
+    """``run-app all`` starts its ui/scheduler children in their own sessions, so killpg
+    on the parent leaves them holding the port and the data dir — breaking the next
+    ``tg start``. The sweep is what actually finishes them off."""
+    proc = MagicMock()
+    proc.poll.return_value = None
+    proc.pid = 4242
+
+    with (
+        patch("tests.installer.platform.system", return_value="Linux"),
+        patch("tests.installer.os.killpg"),
+        patch("tests.installer.os.getpgid", return_value=4242),
+        patch("tests.installer.stop_standalone_orphans") as sweep_mock,
+    ):
+        force_kill_app_tree(proc)
+
+    sweep_mock.assert_called_once()
+
+
+@pytest.mark.unit
+def test_second_interrupt_still_sweeps_orphans():
+    """The second-Ctrl+C path force-kills, so it owes the same cleanup."""
+    proc = MagicMock()
+    proc.poll.return_value = None
+    proc.pid = 4242
+    proc.wait.side_effect = [KeyboardInterrupt(), 0]
+
+    with (
+        patch("tests.installer.platform.system", return_value="Linux"),
+        patch("tests.installer.os.killpg"),
+        patch("tests.installer.os.getpgid", return_value=4242),
+        patch("tests.installer.stop_standalone_orphans") as sweep_mock,
+    ):
+        assert stop_app_tree(proc, timeout=90) is False
+
+    sweep_mock.assert_called_once()
+
+
+@pytest.mark.unit
+def test_graceful_stop_does_not_sweep_orphans():
+    """A tree that stopped on its own has nothing left behind — don't reach for pkill."""
+    proc = MagicMock()
+    proc.poll.return_value = None
+    proc.pid = 4242
+    proc.wait.return_value = 0
+
+    with (
+        patch("tests.installer.platform.system", return_value="Linux"),
+        patch("tests.installer.os.killpg"),
+        patch("tests.installer.os.getpgid", return_value=4242),
+        patch("tests.installer.stop_standalone_orphans") as sweep_mock,
+    ):
+        assert stop_app_tree(proc, timeout=3) is True
+
+    sweep_mock.assert_not_called()
+
+
+@pytest.mark.unit
+def test_force_kill_falls_back_when_taskkill_fails():
+    """taskkill can be denied (elevated child). Without the proc.kill() backstop the
+    installer would report a stop that never happened."""
+    proc = MagicMock()
+    proc.poll.return_value = None
+    proc.pid = 4242
+
+    with (
+        patch("tests.installer.platform.system", return_value="Windows"),
+        patch("tests.installer.subprocess.run", side_effect=OSError("Access denied")),
+    ):
+        force_kill_app_tree(proc, timeout=3)
 
     proc.kill.assert_called_once()
 
diff --git a/tests/test_tg_upgrade.py b/tests/test_tg_upgrade.py
index dfa5f31..21d57d6 100644
--- a/tests/test_tg_upgrade.py
+++ b/tests/test_tg_upgrade.py
@@ -9,7 +9,9 @@
     AbortAction,
     CommandFailed,
     TESTGEN_MAJOR_VERSION,
+    TESTGEN_STOP_GRACE_PERIOD,
     TestgenUpgradeAction,
+    find_in_block,
     InstallMarker,
 )
 
@@ -45,7 +47,12 @@ def tg_upgrade_stdout_side_effect(stdout_mock):
     yield side_effect
 
 
-def get_compose_content(*extra_vars):
+def get_compose_content(*extra_vars, stop_grace=False):
+    """A compose file as an older installer would have written it.
+
+    ``stop_grace`` opts into the engine grace period, i.e. a file already current
+    in that respect — leave it off to model the installs the upgrade has to patch.
+    """
     template = textwrap.dedent("""
         name: testgen
 
@@ -63,10 +70,11 @@ def get_compose_content(*extra_vars):
         services:
           engine:
             image: datakitchen/dataops-testgen:v2.14.5
-
+        {}
     """)
 
-    return template.format(textwrap.indent("\n".join(extra_vars), "  "))
+    grace = f"    stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s\n" if stop_grace else ""
+    return template.format(textwrap.indent("\n".join(extra_vars), "  "), grace)
 
 
 def set_version_check_mock(version_check_mock, latest_version):
@@ -137,7 +145,7 @@ def test_tg_upgrade_abort(
     args_mock.skip_verify = False
     set_version_check_mock(version_check_mock, "1.0.0")
     initial_compose_content = get_compose_content(
-        "TG_INSTANCE_ID: test-instance-id", "TG_UI_BASE_URL: http://localhost:8501"
+        "TG_INSTANCE_ID: test-instance-id", "TG_UI_BASE_URL: http://localhost:8501", stop_grace=True
     )
     compose_path.write_text(initial_compose_content)
 
@@ -239,3 +247,158 @@ def test_tg_upgrade_preserves_existing_base_url(
     compose_content = compose_path.read_text()
     assert "TG_UI_BASE_URL: https://custom.example.com" in compose_content
     assert compose_content.count("TG_UI_BASE_URL") == 1
+
+
+@pytest.mark.integration
+def test_tg_upgrade_adds_stop_grace_period(
+    tg_upgrade_action,
+    compose_path,
+    start_cmd_mock,
+    tg_upgrade_stdout_side_effect,
+    args_mock,
+    version_check_mock,
+):
+    """Existing installs keep their compose file forever — the upgrade is the only
+    chance to give them the grace period a running job needs to checkpoint."""
+    set_version_check_mock(version_check_mock, "1.0.0")
+    compose_path.write_text(get_compose_content("TG_INSTANCE_ID: test-instance-id"))
+
+    tg_upgrade_action.execute(args_mock)
+
+    compose_content = compose_path.read_text()
+    lines = compose_content.splitlines()
+    image_idx = next(i for i, line in enumerate(lines) if "image: datakitchen/dataops-testgen" in line)
+    grace_idx = next(i for i, line in enumerate(lines) if "stop_grace_period" in line)
+    # Inside the engine service, right under its image: two comment lines, then the key.
+    assert grace_idx == image_idx + 3
+    indent = lines[image_idx][: -len(lines[image_idx].lstrip())]
+    assert lines[grace_idx] == f"{indent}stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s"
+
+
+@pytest.mark.integration
+def test_tg_upgrade_preserves_existing_stop_grace_period(
+    tg_upgrade_action,
+    compose_path,
+    start_cmd_mock,
+    tg_upgrade_stdout_side_effect,
+    args_mock,
+    version_check_mock,
+):
+    """A user who tuned the value keeps it, and repeated upgrades don't stack duplicates."""
+    args_mock.skip_verify = True
+    set_version_check_mock(version_check_mock, "1.1.0")
+    compose_path.write_text(
+        get_compose_content("TG_INSTANCE_ID: test-instance-id").replace(
+            "image: datakitchen/dataops-testgen:v2.14.5",
+            "image: datakitchen/dataops-testgen:v2.14.5\n    stop_grace_period: 300s",
+        )
+    )
+
+    tg_upgrade_action.execute(args_mock)
+
+    compose_content = compose_path.read_text()
+    assert "stop_grace_period: 300s" in compose_content
+    assert compose_content.count("stop_grace_period") == 1
+
+
+@pytest.mark.integration
+def test_tg_upgrade_adds_stop_grace_period_to_custom_image(
+    tg_upgrade_action,
+    compose_path,
+    start_cmd_mock,
+    tg_upgrade_stdout_side_effect,
+    args_mock,
+    version_check_mock,
+):
+    """``tg install --image`` accepts a private mirror, so the anchor can't assume the
+    image is a datakitchen one — those installs need the grace period just as much."""
+    args_mock.skip_verify = True
+    set_version_check_mock(version_check_mock, "1.1.0")
+    compose_path.write_text(
+        get_compose_content("TG_INSTANCE_ID: test-instance-id").replace(
+            "datakitchen/dataops-testgen:v2.14.5", "registry.internal.example.com/mirror/testgen:v2.14.5"
+        )
+    )
+
+    tg_upgrade_action.execute(args_mock)
+
+    compose_content = compose_path.read_text()
+    lines = compose_content.splitlines()
+    image_idx = next(i for i, line in enumerate(lines) if "image:" in line)
+    grace_idx = next(i for i, line in enumerate(lines) if "stop_grace_period" in line)
+    assert grace_idx == image_idx + 3
+    assert lines[grace_idx].strip() == f"stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s"
+
+
+@pytest.mark.integration
+def test_tg_upgrade_ignores_stop_grace_period_on_another_service(
+    tg_upgrade_action,
+    compose_path,
+    start_cmd_mock,
+    tg_upgrade_stdout_side_effect,
+    args_mock,
+    version_check_mock,
+):
+    """A grace period set on postgres says nothing about the service that runs the
+    scheduler — the engine must still get its own."""
+    args_mock.skip_verify = True
+    set_version_check_mock(version_check_mock, "1.1.0")
+    compose_path.write_text(
+        # The stray comment matters too: a mention anywhere else in the file must not make
+        # the engine look already-patched.
+        "# note: stop_grace_period is managed by the installer\n"
+        + get_compose_content("TG_INSTANCE_ID: test-instance-id")
+        + "\n  postgres:\n    image: postgres:14.1-alpine\n    stop_grace_period: 30s\n"
+    )
+
+    tg_upgrade_action.execute(args_mock)
+
+    compose_content = compose_path.read_text()
+    engine_block, _, postgres_block = compose_content.partition("  postgres:")
+    assert f"stop_grace_period: {TESTGEN_STOP_GRACE_PERIOD}s" in engine_block
+    # The user's postgres value is left exactly as they set it.
+    assert "stop_grace_period: 30s" in postgres_block
+    # engine's, postgres', and the stray comment.
+    assert compose_content.count("stop_grace_period") == 3
+
+
+COMPOSE_TWO_SERVICES = """name: testgen
+# a stray mention of stop_grace_period above the services section
+services:
+  engine:
+    image: datakitchen/dataops-testgen:v5
+    depends_on:
+      postgres:
+        condition: service_healthy
+
+  postgres:
+    image: postgres:14.1-alpine
+    stop_grace_period: 30s
+"""
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "block, key, expected",
+    (
+        ("engine", "image", "image: datakitchen/dataops-testgen:v5"),
+        ("postgres", "image", "image: postgres:14.1-alpine"),
+        ("postgres", "stop_grace_period", "stop_grace_period: 30s"),
+        # Scoping is the whole point: postgres' grace period is not the engine's, and a
+        # mention in a comment above `services:` is not a setting on any service.
+        ("engine", "stop_grace_period", None),
+        ("engine", "nonexistent", None),
+        ("nonexistent", "image", None),
+    ),
+)
+def test_find_in_block_is_scoped_to_the_block(block, key, expected):
+    match = find_in_block(COMPOSE_TWO_SERVICES, block, key)
+    assert (match.group(0).strip() if match else None) == expected
+
+
+@pytest.mark.unit
+def test_find_in_block_offsets_are_absolute():
+    """Callers splice around the match, so its offsets must index the whole file."""
+    match = find_in_block(COMPOSE_TWO_SERVICES, "postgres", "image")
+    assert COMPOSE_TWO_SERVICES[match.start() : match.end()] == "    image: postgres:14.1-alpine"
+    assert match.group(1) == "    "