fix: Pickle-safe deferred FormatStrings; scope parity for SimpleAction and wrap hooks

Fix three review findings on the run-time-deferred timeout/cancelation
work:

1. Jobs carrying deferred FormatStrings were unpicklable and
   un-deepcopyable. FormatString.__new__ requires a keyword-only parsing
   context, but str.__getnewargs__ supplies only the value, so
   pickle.loads, copy.copy/deepcopy, and model_copy(deep=True) raised
   TypeError - and subclasses that default the context (ArgString,
   CommandString) either failed pickling on the engine-backed parse tree
   or re-parsed EXPR-grammar strings with the legacy parser. FormatString
   now snapshots the parse-relevant context (spec revision + extensions)
   at construction and implements __getnewargs_ex__ (reconstructing
   through __new__ with a context that reproduces the original grammar)
   plus __getstate__ returning None (the parse tree is rebuilt, never
   serialized). EXPR strings round-trip with identical typed evaluation.

2. SimpleAction validated timeout/cancelation at the ambient TASK scope
   while the identical desugared Action form validates them at TEMPLATE
   scope. Add the same _template_field_scopes to SimpleAction, so
   e.g. a bash-sugar timeout referencing Task.Param.* is now rejected
   exactly like the script-form equivalent.

3. Wrap-hook timeout/cancelation could reference WrappedEnv.Name /
   WrappedStep.Name because all per-hook symbols were injected at
   TEMPLATE scope. openjd-rs validates hook timeout/cancelation against
   template symbols + WrappedAction.* only. Split the injection:
   _template_field_inject keeps WrappedAction.* (all scopes) and the new
   _template_field_inject_session carries WrappedEnv.Name /
   WrappedStep.Name at SESSION scope, so hook command/args still accept
   them but timeout/cancelation reject them.

4. §3.4.2 empty-PATH propagation (flagged by the automated PR review):
   PathTaskParameterDefinition rejects literal empty-string range items
   at parse time, but RangeListTaskParameterDefinition - the model that
   RFC 0006 typed whole-field resolution instantiates ranges into - had
   no such check, so `range: "{{Param.Paths}}"` with a LIST[PATH] job
   parameter containing "" flowed into the instantiated Job. Add the
   matching validator on the create-time target model (PATH only;
   STRING ranges legitimately allow ""), mirroring openjd-rs's
   resolve-time check in create_job (ranges.rs).
Also reword three stale comments claiming create-time resolution (the
fields validate at template scope but resolve at run time), and hoist
redundant function-local imports in the FB1 and wrap-actions tests.

Signed-off-by: David Leong <116610336+leongdl@users.noreply.github.com>

Author: leongdl

src/openjd/model/_format_strings/_format_string.py

@@ -7,7 +7,18 @@
 from .._symbol_table import SymbolTable
 from ._dyn_constrained_str import DynamicConstrainedStr
 from ._expression import InterpolationExpression
-from .._types import ModelParsingContextInterface
+from .._types import ModelParsingContextInterface, SpecificationRevision
+
+
+class _ReconstructionContext(ModelParsingContextInterface):
+    """Minimal parsing context used when a FormatString is reconstructed by
+    pickle or the copy module (see ``FormatString.__getnewargs_ex__``).
+
+    It carries only the parse-relevant snapshot (specification revision and
+    extension set) that the FormatString recorded at construction time, so
+    the reconstructed instance is re-parsed under the same grammar (EXPR vs
+    legacy) as the original.
+    """
 
 
 @dataclass
@@ -31,6 +42,13 @@ def __init__(self, *, string: str, start: int, end: int, expr: str = "", details
 
 class FormatString(DynamicConstrainedStr):
     _processed_list: list[Union[str, ExpressionInfo]]
+    # Parse-relevant snapshot of the construction context, recorded so that
+    # pickle/copy reconstruction re-parses the string under the same grammar
+    # (EXPR vs legacy). The context object itself is not retained: it is
+    # shared and mutable (its extension set is narrowed while the template's
+    # `extensions` field is validated), so we snapshot at construction time.
+    _parse_spec_rev: SpecificationRevision
+    _parse_extensions: frozenset[str]
 
     def __new__(cls, value: str, *, context: ModelParsingContextInterface):
         """
@@ -53,9 +71,35 @@ def __new__(cls, value: str, *, context: ModelParsingContextInterface):
         FormatStringError: if the original string is nonvalid.
         """
         self = super().__new__(cls, value, context=context)
+        self._parse_spec_rev = context.spec_rev
+        self._parse_extensions = frozenset(context.extensions)
         self._processed_list = self._preprocess(context=context)
         return self
 
+    def __getnewargs_ex__(self) -> tuple[tuple[str], dict[str, Any]]:
+        """Support for pickling and copying (``pickle``, ``copy.copy``,
+        ``copy.deepcopy``, and pydantic's ``model_copy(deep=True)``).
+
+        ``str.__getnewargs__`` supplies only the string value, which cannot
+        satisfy the keyword-only ``context`` argument of ``__new__`` — and for
+        subclasses that default the argument, it would re-parse an
+        EXPR-grammar string with the legacy parser. Reconstruct through
+        ``__new__`` with a context carrying the recorded parse snapshot so
+        the copy is parsed exactly as the original was.
+        """
+        context = _ReconstructionContext(
+            spec_rev=self._parse_spec_rev,
+            supported_extensions=self._parse_extensions,
+        )
+        return ((str(self),), {"context": context})
+
+    def __getstate__(self) -> None:
+        """No instance state is serialized: ``_processed_list`` holds
+        engine-backed parse trees that cannot be pickled, and ``__new__``
+        fully rebuilds it (and the parse snapshot) from the string value and
+        the reconstruction context (see ``__getnewargs_ex__``)."""
+        return None
+
     @property
     def original_value(self) -> str:
         """

src/openjd/model/_internal/_variable_reference_validation.py

@@ -415,15 +415,26 @@ def _validate_model_template_variable_references(
         # Per-field extra symbols (e.g. the RFC 0008 WrappedAction.* variables
         # within their wrap hook's action). Added at TEMPLATE scope so they
         # are visible in every scope within the field's subtree, including
-        # creation-time fields such as the action's timeout — a wrap hook's
-        # timeout may forward "{{WrappedAction.Timeout}}".
+        # the template-scoped timeout/cancelation fields — a wrap hook's
+        # timeout may forward "{{WrappedAction.Timeout}}" for run-time
+        # resolution.
         for symbol in model._template_field_inject.get(field_name, set()):
             symbol_name = symbol[1:] if symbol.startswith("|") else f"{symbol_prefix}{symbol}"
             _add_symbol(validation_symbols, ResolutionScope.TEMPLATE, symbol_name)
 
-        # Per-field scope override: fields that resolve at job creation (e.g.
-        # an Action's timeout/cancelation) validate at their declared scope
-        # rather than the model's ambient scope.
+        # Per-field SESSION-scoped symbols (e.g. the RFC 0008
+        # WrappedEnv.Name / WrappedStep.Name variables): visible to the
+        # field subtree's session/task-scoped fields (a hook's command/args)
+        # but not to its template-scoped fields (timeout/cancelation),
+        # matching openjd-rs's wrap-hook validation.
+        for symbol in model._template_field_inject_session.get(field_name, set()):
+            symbol_name = symbol[1:] if symbol.startswith("|") else f"{symbol_prefix}{symbol}"
+            _add_symbol(validation_symbols, ResolutionScope.SESSION, symbol_name)
+
+        # Per-field scope override: fields such as an Action's
+        # timeout/cancelation validate at their declared (template) scope
+        # rather than the model's ambient scope, even though their values are
+        # carried through job creation unresolved and resolve at run time.
         field_scope = model._template_field_scopes.get(field_name, current_scope)
 
         errors.extend(

src/openjd/model/_types.py

@@ -377,21 +377,31 @@ class OpenJDModel(BaseModel):
 
     # Per-field variable-scope overrides: fields listed here (and their
     # submodels) validate their format-string references at the given scope
-    # instead of the model's ambient scope. Used for fields that resolve at
-    # job creation while their siblings resolve at run time — e.g. an
-    # Action's `timeout` and `cancelation` are creation-time fields
-    # (template scope: no Session.*, no Env.File.*/Task.File.*, no host
-    # functions) while its `command`/`args` resolve in the session.
+    # instead of the model's ambient scope. Used for fields that VALIDATE at
+    # template scope (no Session.*, no Env.File.*/Task.File.*, no host
+    # functions) while their siblings validate at the ambient session/task
+    # scope — e.g. an Action's `timeout` and `cancelation`, which are carried
+    # through job creation unresolved and resolve at run time, but may only
+    # reference template-scope symbols (plus per-field injections below).
     _template_field_scopes: ClassVar[dict[str, ResolutionScope]] = {}
 
     # Per-field extra symbol injection: symbols listed here are visible in
     # every scope, but only within the named field's subtree. Names use the
     # DefinesTemplateVariables.inject spelling (a "|" prefix discards the
-    # parent scope prefix). Used for the RFC 0008 wrap hooks, whose
-    # WrappedAction.* / WrappedEnv.* / WrappedStep.* variables exist only
-    # within their hook's action.
+    # parent scope prefix). Used for the RFC 0008 wrap hooks' WrappedAction.*
+    # variables, which exist only within their hook's action but must be
+    # visible even to the hook's template-scoped fields (timeout/cancelation)
+    # for round-trip forwarding.
     _template_field_inject: ClassVar[dict[str, set[str]]] = {}
 
+    # As _template_field_inject, but the symbols are injected at SESSION
+    # scope: visible to the field subtree's session/task-scoped fields
+    # (e.g. a wrap hook's command/args) but NOT to its template-scoped
+    # fields (timeout/cancelation). Used for the RFC 0008 WrappedEnv.Name /
+    # WrappedStep.Name variables, which openjd-rs excludes from hook
+    # timeout/cancelation validation.
+    _template_field_inject_session: ClassVar[dict[str, set[str]]] = {}
+
     # ----
     # Metadata used in the creation of a Job from a Job Template
 

src/openjd/model/v2023_09/_model.py

@@ -683,10 +683,16 @@ class EnvironmentActions(OpenJDModel_v2023_09):
     # hook's action, seeded by the runtime when the hook runs in place of the
     # wrapped action. WrappedAction.* is available in all three hooks;
     # WrappedEnv.Name only in the env-enter/exit hooks; WrappedStep.Name only
-    # in the task-run hook. Injected per-field at every scope so a hook's
-    # creation-scoped fields (timeout/cancelation) can round-trip forward
-    # the wrapped action's values ("{{WrappedAction.Timeout}}",
-    # "{{WrappedAction.Cancelation.Mode}}").
+    # in the task-run hook.
+    #
+    # WrappedAction.* is injected at every scope so a hook's template-scoped
+    # fields (timeout/cancelation) can round-trip forward the wrapped
+    # action's values ("{{WrappedAction.Timeout}}",
+    # "{{WrappedAction.Cancelation.Mode}}"). WrappedEnv.Name /
+    # WrappedStep.Name are injected at SESSION scope only: a hook's
+    # command/args may reference them, but its timeout/cancelation may not —
+    # matching openjd-rs (format_strings.rs validates hook
+    # timeout/cancelation against template symbols + WrappedAction.* only).
     _WRAPPED_ACTION_SYMBOLS: ClassVar[set[str]] = {
         "|WrappedAction.Command",
         "|WrappedAction.Args",
@@ -696,9 +702,14 @@ class EnvironmentActions(OpenJDModel_v2023_09):
         "|WrappedAction.Cancelation.NotifyPeriodInSeconds",
     }
     _template_field_inject = {
-        "onWrapEnvEnter": _WRAPPED_ACTION_SYMBOLS | {"|WrappedEnv.Name"},
-        "onWrapEnvExit": _WRAPPED_ACTION_SYMBOLS | {"|WrappedEnv.Name"},
-        "onWrapTaskRun": _WRAPPED_ACTION_SYMBOLS | {"|WrappedStep.Name"},
+        "onWrapEnvEnter": _WRAPPED_ACTION_SYMBOLS,
+        "onWrapEnvExit": _WRAPPED_ACTION_SYMBOLS,
+        "onWrapTaskRun": _WRAPPED_ACTION_SYMBOLS,
+    }
+    _template_field_inject_session = {
+        "onWrapEnvEnter": {"|WrappedEnv.Name"},
+        "onWrapEnvExit": {"|WrappedEnv.Name"},
+        "onWrapTaskRun": {"|WrappedStep.Name"},
     }
 
     @model_validator(mode="before")
@@ -1040,6 +1051,17 @@ class SimpleAction(OpenJDModel_v2023_09):
         "args": {"__self__"},
     }
 
+    # Parity with Action: `timeout` and `cancelation` VALIDATE at template
+    # scope (no Session.*, no Env.File.*/Task.File.*, no host-context
+    # functions) even though the desugared SimpleAction is otherwise
+    # TASK-scoped. The desugared form is an ordinary Action, whose identical
+    # fields validate at template scope — matching openjd-rs
+    # (format_strings.rs validates these fields against the template symtab).
+    _template_field_scopes = {
+        "timeout": ResolutionScope.TEMPLATE,
+        "cancelation": ResolutionScope.TEMPLATE,
+    }
+
     @field_validator("let")
     @classmethod
     def _validate_let(cls, v: Any, info: ValidationInfo) -> Any:
@@ -1243,6 +1265,22 @@ def _validate_range_len(cls, value: Any) -> Any:
         # resolve-time checks in create_job (ranges.rs).
         return validate_task_param_range_list_len(value)
 
+    @field_validator("range")
+    @classmethod
+    def _validate_no_empty_path_values(cls, value: Any, info: ValidationInfo) -> Any:
+        # §3.4.2: an empty string is not a valid path on any OS, so a PATH
+        # task parameter's range may not contain one. The template model's
+        # parse-time validator cannot see values that arrive via RFC 0006
+        # typed whole-field resolution (e.g. `range: "{{Param.Paths}}"` with
+        # a LIST[PATH] job parameter), so the check is enforced on the
+        # instantiation target as well — matching openjd-rs's resolve-time
+        # check in create_job (ranges.rs).
+        if info.data.get("type") == TaskParameterType.PATH and isinstance(value, list):
+            for i, item in enumerate(value):
+                if str(item) == "":
+                    raise ValueError(f"range[{i}] must not be an empty string.")
+        return value
+
 
 class RangeExpressionTaskParameterDefinition(OpenJDModel_v2023_09):
     # element type of items in the range

test/openjd/model_v0/test_pickle.py

@@ -0,0 +1,165 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+"""Regression tests: Jobs carrying deferred (run-time-resolved) FormatStrings
+must be picklable and copyable.
+
+FormatString.__new__ requires a keyword-only parsing context, but
+str.__getnewargs__ supplies only the string value, so pickle.loads,
+copy.copy/deepcopy, and model_copy(deep=True) raised TypeError — and
+subclasses that default the context (e.g. ArgString) silently re-parsed
+EXPR-grammar strings with the legacy parser (or failed pickling on the
+engine-backed parse tree). FormatString now records a parse-context snapshot
+and reconstructs through __getnewargs_ex__, re-parsing under the same grammar
+(see FormatString.__getstate__/__getnewargs_ex__).
+"""
+
+import copy
+import pickle
+
+from openjd.model import SymbolTable, create_job, decode_job_template
+from openjd.model._format_strings import FormatString
+from openjd.model._types import ParameterValue, ParameterValueType
+from openjd.model.v2023_09 import ArgString, ModelParsingContext
+
+
+def _job_with_deferred_fields(extensions: list[str], *, timeout: str):
+    """An instantiated Job whose onRun action carries a deferred (unresolved)
+    FormatString timeout and a deferred cancelation (mode + notifyPeriod)."""
+    template = decode_job_template(
+        template={
+            "specificationVersion": "jobtemplate-2023-09",
+            "extensions": list(extensions),
+            "name": "PickleTest",
+            "parameterDefinitions": [
+                {"name": "TO", "type": "INT", "default": 30},
+                {"name": "Mode", "type": "STRING", "default": "TERMINATE"},
+            ],
+            "steps": [
+                {
+                    "name": "Step1",
+                    "script": {
+                        "actions": {
+                            "onRun": {
+                                "command": "echo",
+                                "timeout": timeout,
+                                "cancelation": {
+                                    "mode": "{{Param.Mode}}",
+                                    "notifyPeriodInSeconds": "{{Param.TO}}",
+                                },
+                            }
+                        }
+                    },
+                }
+            ],
+        },
+        supported_extensions=extensions,
+    )
+    return create_job(
+        job_template=template,
+        job_parameter_values={
+            "TO": ParameterValue(type=ParameterValueType.INT, value="30"),
+            "Mode": ParameterValue(type=ParameterValueType.STRING, value="TERMINATE"),
+        },
+    )
+
+
+def _deferred_symtab() -> SymbolTable:
+    symtab = SymbolTable()
+    symtab["Param.TO"] = 30
+    symtab["Param.Mode"] = "TERMINATE"
+    return symtab
+
+
+def _assert_deferred_fields_equivalent(original, reconstructed) -> None:
+    """The reconstructed action's deferred FormatStrings must be FormatStrings
+    that evaluate identically to the originals against the same symtab."""
+    symtab = _deferred_symtab()
+    for attr_path in ("timeout", "cancelation.mode", "cancelation.notifyPeriodInSeconds"):
+        orig_value = original
+        recon_value = reconstructed
+        for part in attr_path.split("."):
+            orig_value = getattr(orig_value, part)
+            recon_value = getattr(recon_value, part)
+        assert isinstance(recon_value, FormatString), attr_path
+        assert type(recon_value) is type(orig_value), attr_path
+        assert str(recon_value) == str(orig_value), attr_path
+        assert recon_value.resolve(symtab=symtab) == orig_value.resolve(symtab=symtab), attr_path
+
+
+class TestJobPickleRoundTrip:
+    def test_expr_job_pickle_round_trip(self) -> None:
+        # GIVEN an EXPR + FEATURE_BUNDLE_1 job with deferred timeout and
+        # cancelation, where the timeout uses EXPR-only arithmetic (the
+        # legacy grammar cannot parse it — a reconstruction under a default
+        # context would fail or change semantics).
+        job = _job_with_deferred_fields(["EXPR", "FEATURE_BUNDLE_1"], timeout="{{ Param.TO * 2 }}")
+
+        # WHEN
+        reconstructed = pickle.loads(pickle.dumps(job))
+
+        # THEN
+        original_action = job.steps[0].script.actions.onRun
+        recon_action = reconstructed.steps[0].script.actions.onRun
+        _assert_deferred_fields_equivalent(original_action, recon_action)
+        symtab = _deferred_symtab()
+        assert recon_action.timeout.resolve(symtab=symtab) == "60"
+
+    def test_legacy_fb1_job_pickle_round_trip(self) -> None:
+        # GIVEN a legacy (non-EXPR) FEATURE_BUNDLE_1 job with a deferred
+        # FormatString timeout.
+        job = _job_with_deferred_fields(["FEATURE_BUNDLE_1"], timeout="{{Param.TO}}")
+
+        # WHEN
+        reconstructed = pickle.loads(pickle.dumps(job))
+
+        # THEN
+        original_action = job.steps[0].script.actions.onRun
+        recon_action = reconstructed.steps[0].script.actions.onRun
+        _assert_deferred_fields_equivalent(original_action, recon_action)
+        symtab = _deferred_symtab()
+        assert recon_action.timeout.resolve(symtab=symtab) == "30"
+
+    def test_expr_job_deepcopy_and_model_copy(self) -> None:
+        # GIVEN
+        job = _job_with_deferred_fields(["EXPR", "FEATURE_BUNDLE_1"], timeout="{{ Param.TO * 2 }}")
+        original_action = job.steps[0].script.actions.onRun
+
+        # WHEN / THEN: both deep-copy paths reconstruct working FormatStrings.
+        for job_copy in (copy.deepcopy(job), job.model_copy(deep=True)):
+            copied_action = job_copy.steps[0].script.actions.onRun
+            _assert_deferred_fields_equivalent(original_action, copied_action)
+
+
+class TestFormatStringPickleFaithfulness:
+    """The reconstruction context must reproduce the original parse mode:
+    re-parsing an EXPR-grammar string under a default (legacy) context would
+    reject it outright or silently change its evaluation semantics."""
+
+    def test_expr_arg_string_typed_round_trip(self) -> None:
+        # GIVEN an ArgString whose whole-field EXPR list expression only
+        # parses under the EXPR grammar.
+        context = ModelParsingContext(supported_extensions=["EXPR", "FEATURE_BUNDLE_1"])
+        original = ArgString('{{ ["a","b c"] }}', context=context)
+
+        # WHEN
+        reconstructed = pickle.loads(pickle.dumps(original))
+
+        # THEN: still an EXPR whole-field expression with identical typed
+        # (RFC 0005) evaluation, not a legacy re-parse.
+        assert type(reconstructed) is ArgString
+        assert reconstructed.whole_field_expression() is not None
+        symtab = SymbolTable()
+        original_value = original.resolve_value(symtab=symtab)
+        reconstructed_value = reconstructed.resolve_value(symtab=symtab)
+        assert type(reconstructed_value) is type(original_value)
+        assert str(reconstructed_value) == str(original_value)
+        assert reconstructed.resolve(symtab=symtab) == original.resolve(symtab=symtab)
+
+    def test_legacy_format_string_copy(self) -> None:
+        # A plain legacy-grammar FormatString survives copy.copy and deepcopy.
+        context = ModelParsingContext()
+        original = FormatString("prefix {{ Task.Param.Frame }} suffix", context=context)
+        symtab = SymbolTable()
+        symtab["Task.Param.Frame"] = 7
+        for reconstructed in (copy.copy(original), copy.deepcopy(original)):
+            assert str(reconstructed) == str(original)
+            assert reconstructed.resolve(symtab=symtab) == "prefix 7 suffix"