Keep secrets safe in conditional branch rendering (#4353)

skyvern/forge/sdk/workflow/models/block.py

@@ -4283,9 +4283,11 @@ class BranchEvaluationContext:
         *,
         workflow_run_context: WorkflowRunContext | None = None,
         block_label: str | None = None,
+        template_renderer: Callable[[str], str] | None = None,
     ) -> None:
         self.workflow_run_context = workflow_run_context
         self.block_label = block_label
+        self.template_renderer = template_renderer
 
     def build_template_data(self) -> dict[str, Any]:
         """Build Jinja template data mirroring block parameter rendering context."""
@@ -4402,28 +4404,36 @@ class JinjaBranchCriteria(BranchCriteria):
     criteria_type: Literal["jinja2_template"] = "jinja2_template"
 
     async def evaluate(self, context: BranchEvaluationContext) -> bool:
-        # Build the template context explicitly to avoid surprises in templates.
-        template_data = context.build_template_data()
+        # Prefer the renderer provided by the caller (matches block parameter rendering),
+        # otherwise build a minimal sandboxed renderer using the evaluation context.
+        if context.template_renderer:
+            try:
+                rendered = context.template_renderer(self.expression)
+            except MissingJinjaVariables:
+                # Let upstream MissingJinjaVariables bubble as-is.
+                raise
+            except Exception as exc:  # pragma: no cover - caught for robustness
+                raise FailedToFormatJinjaStyleParameter(self.expression, str(exc)) from exc
+        else:
+            template_data = context.build_template_data()
+            sandbox_env = (
+                SandboxedEnvironment(undefined=StrictUndefined)
+                if settings.WORKFLOW_TEMPLATING_STRICTNESS == "strict"
+                else SandboxedEnvironment()
+            )
 
-        try:
-            template = jinja_sandbox_env.from_string(self.expression)
-        except Exception as exc:
-            raise FailedToFormatJinjaStyleParameter(
-                template=self.expression,
-                msg=str(exc),
-            ) from exc
+            try:
+                missing_vars = get_missing_variables(self.expression, template_data)
+                if missing_vars:
+                    raise MissingJinjaVariables(self.expression, missing_vars)
 
-        if settings.WORKFLOW_TEMPLATING_STRICTNESS == "strict":
-            if missing := get_missing_variables(self.expression, template_data):
-                raise MissingJinjaVariables(template=self.expression, variables=missing)
-
-        try:
-            rendered = template.render(template_data)
-        except Exception as exc:
-            raise FailedToFormatJinjaStyleParameter(
-                template=self.expression,
-                msg=str(exc),
-            ) from exc
+                template = sandbox_env.from_string(self.expression)
+                rendered = template.render(template_data)
+            except MissingJinjaVariables:
+                raise
+            except Exception as exc:
+                # Covers syntax errors and rendering issues
+                raise FailedToFormatJinjaStyleParameter(self.expression, str(exc)) from exc
 
         return _evaluate_truthy_string(rendered)
 
@@ -4521,35 +4531,13 @@ class ConditionalBlock(Block):
             raise ValueError("organization_id is required to evaluate natural language branches")
 
         workflow_run_context = evaluation_context.workflow_run_context
-        template_data = evaluation_context.build_template_data()
 
         rendered_branch_criteria: list[dict[str, Any]] = []
         for idx, branch in enumerate(branches):
             expression = branch.criteria.expression if branch.criteria else ""
-            rendered_expression = expression
-
-            # Allow Jinja templating inside natural language branch expressions so users can
-            # mix free text with dynamic values (e.g., "If response: {{ foo.bar }}").
-            if "{{" in expression and "}}" in expression:
-                try:
-                    template = jinja_sandbox_env.from_string(expression)
-                except Exception as exc:
-                    raise FailedToFormatJinjaStyleParameter(
-                        template=expression,
-                        msg=str(exc),
-                    ) from exc
-
-                if settings.WORKFLOW_TEMPLATING_STRICTNESS == "strict":
-                    if missing := get_missing_variables(expression, template_data):
-                        raise MissingJinjaVariables(template=expression, variables=missing)
-
-                try:
-                    rendered_expression = template.render(template_data)
-                except Exception as exc:
-                    raise FailedToFormatJinjaStyleParameter(
-                        template=expression,
-                        msg=str(exc),
-                    ) from exc
+            rendered_expression = (
+                evaluation_context.template_renderer(expression) if evaluation_context.template_renderer else expression
+            )
 
             rendered_branch_criteria.append({"index": idx, "expression": rendered_expression})
 
@@ -4668,6 +4656,14 @@ class ConditionalBlock(Block):
         evaluation_context = BranchEvaluationContext(
             workflow_run_context=workflow_run_context,
             block_label=self.label,
+            template_renderer=(
+                lambda potential_template: self.format_block_parameter_template_from_workflow_run_context(
+                    potential_template,
+                    workflow_run_context,
+                )
+            )
+            if workflow_run_context
+            else None,
         )
 
         matched_branch = None