fix(cfn-lang-ext): preserve template-time intrinsics with unresolved param Refs in PARTIAL mode (#9004) (#9010)

* fix: preserve Fn::FindInMap with unresolved param Ref keys in PARTIAL mode (#9004)

* fix: preserve Fn::Join with unresolved param Ref args in PARTIAL mode

* fix: widen Fn::Join resolve() return type to Any for PARTIAL preservation

* fix: preserve Fn::Select with unresolved param Ref index in PARTIAL mode

* fix: preserve Fn::Base64 with unresolved param Ref arg in PARTIAL mode

* test: end-to-end regression test for issue #9004

* test: rename new test classes in test_fn_find_in_map.py to avoid name collision

* refactor: consolidate is_unresolved_param_or_pseudo_ref into utils

Hoist the helper duplicated in four resolvers (fn_find_in_map, fn_join,
fn_select, fn_base64) up to samcli/lib/cfn_language_extensions/utils.py
next to PSEUDO_PARAMETERS. Each resolver now imports the shared helper
and drops its private copy.

Adds focused unit tests in test_utils.py covering parameter Refs,
pseudo-parameter Refs, resource Refs, Fn::GetAtt, multi-key dicts,
non-string Ref targets, and contexts without a parsed_template.

* style: apply black formatting

Author: bnusunny

samcli/lib/cfn_language_extensions/resolvers/fn_base64.py

@@ -9,7 +9,9 @@
 from typing import Any, Dict
 
 from samcli.lib.cfn_language_extensions.exceptions import InvalidTemplateException
+from samcli.lib.cfn_language_extensions.models import ResolutionMode
 from samcli.lib.cfn_language_extensions.resolvers.base import IntrinsicFunctionResolver
+from samcli.lib.cfn_language_extensions.utils import is_unresolved_param_or_pseudo_ref
 
 
 class FnBase64Resolver(IntrinsicFunctionResolver):
@@ -32,7 +34,7 @@ class FnBase64Resolver(IntrinsicFunctionResolver):
 
     FUNCTION_NAMES = ["Fn::Base64"]
 
-    def resolve(self, value: Dict[str, Any]) -> str:
+    def resolve(self, value: Dict[str, Any]) -> Any:
         """
         Resolve the Fn::Base64 intrinsic function.
 
@@ -46,7 +48,10 @@ def resolve(self, value: Dict[str, Any]) -> str:
                    {"Fn::Base64": {"Ref": "MyStringParam"}}
 
         Returns:
-            The base64-encoded string.
+            The base64-encoded string. In PARTIAL resolution mode, if the
+            argument is an unresolved Ref to a declared template parameter or a
+            pseudo-parameter, the original Fn::Base64 call is returned verbatim
+            so CloudFormation can resolve it at deploy time.
 
         Raises:
             InvalidTemplateException: If the resolved value is not a string.
@@ -64,6 +69,15 @@ def resolve(self, value: Dict[str, Any]) -> str:
             # If no parent resolver, use args as-is (for testing)
             resolved_args = args
 
+        # In PARTIAL mode, if the argument resolved to a deferred parameter Ref
+        # (a Ref to a declared parameter without a default/override, or to a
+        # pseudo-parameter), preserve the call so CloudFormation can resolve it
+        # at deploy time. Resource refs / GetAtt / etc. still raise.
+        if is_unresolved_param_or_pseudo_ref(resolved_args, self.context):
+            if self.context.resolution_mode == ResolutionMode.PARTIAL:
+                return {"Fn::Base64": resolved_args}
+            raise InvalidTemplateException("Fn::Base64 layout is incorrect")
+
         # Validate that the resolved value is a string
         if not isinstance(resolved_args, str):
             raise InvalidTemplateException("Fn::Base64 layout is incorrect")

samcli/lib/cfn_language_extensions/resolvers/fn_find_in_map.py

@@ -12,7 +12,9 @@
 from typing import Any, Dict
 
 from samcli.lib.cfn_language_extensions.exceptions import InvalidTemplateException
+from samcli.lib.cfn_language_extensions.models import ResolutionMode
 from samcli.lib.cfn_language_extensions.resolvers.base import IntrinsicFunctionResolver
+from samcli.lib.cfn_language_extensions.utils import is_unresolved_param_or_pseudo_ref
 
 
 class FnFindInMapResolver(IntrinsicFunctionResolver):
@@ -95,12 +97,23 @@ def resolve(self, value: Dict[str, Any]) -> Any:
             top_key = top_key_arg
             second_key = second_key_arg
 
-        # Validate resolved keys are strings
-        if not isinstance(map_name, str):
-            raise InvalidTemplateException("Fn::FindInMap layout is incorrect")
-        if not isinstance(top_key, str):
-            raise InvalidTemplateException("Fn::FindInMap layout is incorrect")
-        if not isinstance(second_key, str):
+        # In PARTIAL mode, if a key didn't resolve to a string but did resolve
+        # to a deferred parameter/pseudo-parameter Ref (Ref to a declared
+        # parameter without a default value and no override, or to a
+        # pseudo-parameter without a provided value), preserve the call so
+        # CloudFormation can resolve it at deploy time. See GitHub issue #9004.
+        # Resource refs and other intrinsics (Fn::GetAtt, etc.) continue to
+        # raise because they can't be resolved by the template-time
+        # Fn::FindInMap function at any stage — matching Kotlin compat.
+        keys = (map_name, top_key, second_key)
+        if not all(isinstance(k, str) for k in keys):
+            if self.context.resolution_mode == ResolutionMode.PARTIAL and all(
+                isinstance(k, str) or is_unresolved_param_or_pseudo_ref(k, self.context) for k in keys
+            ):
+                preserved = [map_name, top_key, second_key]
+                if len(args) >= self._ARGS_WITH_DEFAULT:
+                    preserved.append(args[3])
+                return {"Fn::FindInMap": preserved}
             raise InvalidTemplateException("Fn::FindInMap layout is incorrect")
 
         # Check for DefaultValue option (4th argument)

samcli/lib/cfn_language_extensions/resolvers/fn_join.py

@@ -10,7 +10,9 @@
 from typing import Any, Dict
 
 from samcli.lib.cfn_language_extensions.exceptions import InvalidTemplateException
+from samcli.lib.cfn_language_extensions.models import ResolutionMode
 from samcli.lib.cfn_language_extensions.resolvers.base import IntrinsicFunctionResolver
+from samcli.lib.cfn_language_extensions.utils import is_unresolved_param_or_pseudo_ref
 
 
 class FnJoinResolver(IntrinsicFunctionResolver):
@@ -37,7 +39,7 @@ class FnJoinResolver(IntrinsicFunctionResolver):
 
     _EXPECTED_ARGS = 2
 
-    def resolve(self, value: Dict[str, Any]) -> str:
+    def resolve(self, value: Dict[str, Any]) -> Any:
         """
         Resolve the Fn::Join intrinsic function.
 
@@ -50,7 +52,9 @@ def resolve(self, value: Dict[str, Any]) -> str:
                    E.g., {"Fn::Join": [",", ["a", "b", "c"]]}
 
         Returns:
-            A string with all list elements joined by the delimiter.
+            A string with all list elements joined by the delimiter, or - in
+            PARTIAL mode when an arg is an unresolved parameter Ref - the
+            preserved Fn::Join call as a dict.
 
         Raises:
             InvalidTemplateException: If the layout is incorrect.
@@ -70,14 +74,25 @@ def resolve(self, value: Dict[str, Any]) -> str:
         if self.parent is not None:
             delimiter = self.parent.resolve_value(delimiter)
 
-        # Validate delimiter is a string
-        if not isinstance(delimiter, str):
-            raise InvalidTemplateException("Fn::Join layout is incorrect")
-
         # Resolve any nested intrinsic functions in the list
         if self.parent is not None:
             list_to_join = self.parent.resolve_value(list_to_join)
 
+        # In PARTIAL mode, preserve the call when either argument is still an
+        # unresolved Ref to a declared template parameter or a pseudo-parameter
+        # (i.e. CloudFormation will resolve it at deploy time). Resource refs
+        # and other intrinsics still raise — they aren't valid Fn::Join inputs.
+        delim_is_param_ref = is_unresolved_param_or_pseudo_ref(delimiter, self.context)
+        list_is_param_ref = is_unresolved_param_or_pseudo_ref(list_to_join, self.context)
+        if delim_is_param_ref or list_is_param_ref:
+            if self.context.resolution_mode == ResolutionMode.PARTIAL:
+                return {"Fn::Join": [delimiter, list_to_join]}
+            raise InvalidTemplateException("Fn::Join layout is incorrect")
+
+        # Validate delimiter is a string
+        if not isinstance(delimiter, str):
+            raise InvalidTemplateException("Fn::Join layout is incorrect")
+
         # Validate the list
         if not isinstance(list_to_join, list):
             raise InvalidTemplateException("Fn::Join layout is incorrect")

samcli/lib/cfn_language_extensions/resolvers/fn_select.py

@@ -10,7 +10,9 @@
 from typing import Any, Dict
 
 from samcli.lib.cfn_language_extensions.exceptions import InvalidTemplateException
+from samcli.lib.cfn_language_extensions.models import ResolutionMode
 from samcli.lib.cfn_language_extensions.resolvers.base import IntrinsicFunctionResolver
+from samcli.lib.cfn_language_extensions.utils import is_unresolved_param_or_pseudo_ref
 
 
 class FnSelectResolver(IntrinsicFunctionResolver):
@@ -73,6 +75,18 @@ def resolve(self, value: Dict[str, Any]) -> Any:
         if self.parent is not None:
             index = self.parent.resolve_value(index)
 
+        # In PARTIAL mode, if the index resolved to a deferred parameter Ref,
+        # preserve the call. Resource refs / GetAtt / etc. still raise.
+        # We must still resolve the source list below in case it contains
+        # resolvable intrinsics — that way, partial expansion makes maximal
+        # progress.
+        if is_unresolved_param_or_pseudo_ref(index, self.context):
+            if self.context.resolution_mode != ResolutionMode.PARTIAL:
+                raise InvalidTemplateException("Fn::Select layout is incorrect")
+            if self.parent is not None:
+                source_list = self.parent.resolve_value(source_list)
+            return {"Fn::Select": [index, source_list]}
+
         # Validate index is an integer (or can be converted to one)
         if isinstance(index, str):
             try:

samcli/lib/cfn_language_extensions/utils.py

@@ -5,7 +5,10 @@
 for working with templates that may contain Fn::ForEach blocks.
 """
 
-from typing import Dict, Iterator, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Iterator, Tuple
+
+if TYPE_CHECKING:
+    from samcli.lib.cfn_language_extensions.models import TemplateProcessingContext
 
 FOREACH_PREFIX = "Fn::ForEach::"
 
@@ -65,6 +68,31 @@ def is_intrinsic_key(key: str) -> bool:
     return key.startswith("Fn::") or key in _INTRINSIC_SINGLE_KEYS
 
 
+def is_unresolved_param_or_pseudo_ref(value: Any, context: "TemplateProcessingContext") -> bool:
+    """Return True if *value* is ``{"Ref": <name>}`` where ``<name>`` is a declared
+    template parameter or a pseudo-parameter — i.e. an unresolved reference that
+    CloudFormation will resolve at deploy time.
+
+    Used by template-time intrinsic resolvers (Fn::FindInMap, Fn::Join, Fn::Select,
+    Fn::Base64) to decide, in PARTIAL resolution mode, whether to preserve the
+    enclosing call instead of raising. Resource refs and other intrinsics return
+    False so they continue to raise — matching Kotlin compatibility for
+    template-time intrinsics that genuinely cannot accept deploy-time inputs.
+    """
+    if not isinstance(value, dict) or len(value) != 1:
+        return False
+    if "Ref" not in value:
+        return False
+    ref_target = value["Ref"]
+    if not isinstance(ref_target, str):
+        return False
+    if ref_target in PSEUDO_PARAMETERS:
+        return True
+    if context.parsed_template is not None and ref_target in context.parsed_template.parameters:
+        return True
+    return False
+
+
 # Mapping-name prefixes that SAM CLI emits for dynamic Fn::ForEach handling:
 #   - SAM + <packageable artifact property> + <nesting path> + [resource suffix]
 #     (see language_extensions_packaging._compute_mapping_name)

tests/unit/lib/cfn_language_extensions/test_fn_base64.py

@@ -25,6 +25,7 @@
     IntrinsicResolver,
 )
 from samcli.lib.cfn_language_extensions.resolvers.fn_base64 import FnBase64Resolver
+from samcli.lib.cfn_language_extensions.resolvers.fn_ref import FnRefResolver
 from samcli.lib.cfn_language_extensions.exceptions import InvalidTemplateException
 
 
@@ -478,6 +479,22 @@ def test_fn_base64_with_preserved_intrinsic(self, orchestrator: IntrinsicResolve
             "preserved": {"Fn::GetAtt": ["MyBucket", "Arn"]},
         }
 
+    def test_unresolved_arg_is_preserved_in_partial_mode(self):
+        from samcli.lib.cfn_language_extensions.models import ParsedTemplate
+
+        ctx = TemplateProcessingContext(
+            fragment={"Resources": {}},
+            resolution_mode=ResolutionMode.PARTIAL,
+            parsed_template=ParsedTemplate(parameters={"UserData": {"Type": "String"}}),
+        )
+        orch = IntrinsicResolver(ctx)
+        orch.register_resolver(FnRefResolver)
+        orch.register_resolver(FnBase64Resolver)
+
+        value = {"Fn::Base64": {"Ref": "UserData"}}
+        result = orch.resolve_value(value)
+        assert result == {"Fn::Base64": {"Ref": "UserData"}}
+
 
 class TestFnBase64ResolverRealWorldScenarios:
     """Tests for real-world CloudFormation scenarios using Fn::Base64."""