From ba07d0914ac44867fd76ddc6cfc7395c892ef7a0 Mon Sep 17 00:00:00 2001
From: David Leong <116610336+leongdl@users.noreply.github.com>
Date: Sat, 9 May 2026 15:10:35 -0700
Subject: [PATCH] feat: Support RFC extensions in environment templates

---
 src/openjd/cli/_check/_check_command.py     |  7 ++++++-
 src/openjd/cli/_common/_extensions.py       | 13 +++++++++++--
 src/openjd/cli/_common/_validation_utils.py | 20 +++++++++++++++++---
 src/openjd/cli/_run/_run_command.py         |  8 +++++++-
 4 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/src/openjd/cli/_check/_check_command.py b/src/openjd/cli/_check/_check_command.py
index b3297ea..7e82e5b 100644
--- a/src/openjd/cli/_check/_check_command.py
+++ b/src/openjd/cli/_check/_check_command.py
@@ -41,7 +41,12 @@ def do_check(args: Namespace) -> OpenJDCliResult:
         if TemplateSpecificationVersion.is_job_template(template_version):
             decode_job_template(template=template_object, supported_extensions=extensions)
         elif TemplateSpecificationVersion.is_environment_template(template_version):
-            decode_environment_template(template=template_object)
+            # Pass `extensions` so env templates can declare e.g.
+            # WRAP_ACTIONS (RFC 0008) — without this the model rejects
+            # the extension declaration even though the CLI accepts it.
+            decode_environment_template(
+                template=template_object, supported_extensions=extensions
+            )
         else:
             return OpenJDCliResult(
                 status="error",
diff --git a/src/openjd/cli/_common/_extensions.py b/src/openjd/cli/_common/_extensions.py
index b414a11..ba6e4cd 100644
--- a/src/openjd/cli/_common/_extensions.py
+++ b/src/openjd/cli/_common/_extensions.py
@@ -3,8 +3,17 @@
 from argparse import ArgumentParser
 from typing import Optional
 
-# This is the list of Open Job Description extensions with implemented support
-SUPPORTED_EXTENSIONS = ["TASK_CHUNKING", "REDACTED_ENV_VARS", "FEATURE_BUNDLE_1"]
+# This is the list of Open Job Description extensions with implemented support.
+# EXPR (RFCs 0005/0006) and WRAP_ACTIONS (RFC 0008) are added so the CLI can run
+# templates that declare them. WRAP_ACTIONS implies EXPR; the model library
+# handles that implication when it parses the template.
+SUPPORTED_EXTENSIONS = [
+    "TASK_CHUNKING",
+    "REDACTED_ENV_VARS",
+    "FEATURE_BUNDLE_1",
+    "EXPR",
+    "WRAP_ACTIONS",
+]
 
 
 def add_extensions_argument(run_parser: ArgumentParser):
diff --git a/src/openjd/cli/_common/_validation_utils.py b/src/openjd/cli/_common/_validation_utils.py
index 19a0f62..ef9879a 100644
--- a/src/openjd/cli/_common/_validation_utils.py
+++ b/src/openjd/cli/_common/_validation_utils.py
@@ -1,6 +1,6 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 
-from typing import Any
+from typing import Any, Optional
 from pathlib import Path
 
 from openjd.model import (
@@ -70,15 +70,29 @@ def read_job_template(template_file: Path, *, supported_extensions: list[str]) -
     return template
 
 
-def read_environment_template(template_file: Path) -> EnvironmentTemplate:
+def read_environment_template(
+    template_file: Path,
+    *,
+    supported_extensions: Optional[list[str]] = None,
+) -> EnvironmentTemplate:
     """Open a JSON or YAML-formatted file and attempt to parse it into an EnvironmentTemplate object.
     Raises a RuntimeError if the file doesn't exist or can't be opened, and raises a
     DecodeValidationError if its contents can't be parsed into a valid EnvironmentTemplate.
+
+    Environment templates may declare extensions just like job templates do
+    (e.g. ``WRAP_ACTIONS`` for RFC 0008's wrap hooks). Pass the CLI's allow-list
+    via ``supported_extensions`` so the model accepts those declarations;
+    omitting the argument preserves the legacy behavior of accepting only the
+    default model surface.
     """
     # Raises RuntimeError
     template_object = read_template(template_file)
 
+    decode_kwargs: dict[str, list[str]] = {}
+    if supported_extensions is not None:
+        decode_kwargs["supported_extensions"] = supported_extensions
+
     # Raises: DecodeValidationError
-    template = decode_environment_template(template=template_object)
+    template = decode_environment_template(template=template_object, **decode_kwargs)
 
     return template
diff --git a/src/openjd/cli/_run/_run_command.py b/src/openjd/cli/_run/_run_command.py
index 7b59079..c8817d0 100644
--- a/src/openjd/cli/_run/_run_command.py
+++ b/src/openjd/cli/_run/_run_command.py
@@ -415,7 +415,13 @@ def do_run(args: Namespace) -> OpenJDCliResult:
             filename = Path(env).expanduser()
             try:
                 # Raises: RuntimeError, DecodeValidationError
-                env_template = read_environment_template(filename)
+                # Pass `extensions` so env templates can declare e.g.
+                # WRAP_ACTIONS (RFC 0008) — without this they're parsed
+                # against the default model surface and reject the
+                # extension declaration.
+                env_template = read_environment_template(
+                    filename, supported_extensions=extensions
+                )
                 environments.append(env_template)
             except (RuntimeError, DecodeValidationError) as e:
                 return OpenJDCliResult(status="error", message=str(e))