From 008cc26a15820f3efc7f93398fafde8ab07e3a69 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Sun, 23 Mar 2025 18:20:20 -0700
Subject: [PATCH] improve block output jinja reference (#2006)

---
 skyvern/forge/sdk/workflow/context_manager.py | 27 +++++++++++++++++++
 skyvern/forge/sdk/workflow/models/block.py    | 12 ++++++++-
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/skyvern/forge/sdk/workflow/context_manager.py b/skyvern/forge/sdk/workflow/context_manager.py
index 1323b881..573eb339 100644
--- a/skyvern/forge/sdk/workflow/context_manager.py
+++ b/skyvern/forge/sdk/workflow/context_manager.py
@@ -1,3 +1,4 @@
+import copy
 import uuid
 from typing import TYPE_CHECKING, Any, Self
 
@@ -584,8 +585,34 @@ class WorkflowRunContext:
             LOG.warning(f"Output parameter {parameter.output_parameter_id} already has a registered value, overwriting")
 
         self.values[parameter.key] = value
+        self.register_block_reference_variable_from_output_parameter(parameter, value)
+
         await self.set_parameter_values_for_output_parameter_dependent_blocks(parameter, value)
 
+    def register_block_reference_variable_from_output_parameter(
+        self,
+        parameter: OutputParameter,
+        value: dict[str, Any] | list | str | None,
+    ) -> None:
+        # output parameter key is formatted as `<block_label>_output`
+        if not parameter.key.endswith("_output"):
+            return
+        block_label = parameter.key.removesuffix("_output")
+
+        block_reference_value = copy.deepcopy(value)
+        if isinstance(block_reference_value, dict) and "extracted_information" in block_reference_value:
+            block_reference_value.update({"output": block_reference_value.get("extracted_information")})
+
+        if block_label in self.values:
+            current_value = self.values[block_label]
+            # only able to merge the value when the current value and the pending value are both dicts
+            if isinstance(current_value, dict) and isinstance(block_reference_value, dict):
+                block_reference_value.update(current_value)
+            else:
+                LOG.warning(f"Parameter {block_label} already has a value in workflow run context, overwriting")
+
+        self.values[block_label] = block_reference_value
+
     async def set_parameter_values_for_output_parameter_dependent_blocks(
         self,
         output_parameter: OutputParameter,
diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py
index c6ebddc6..82b15d96 100644
--- a/skyvern/forge/sdk/workflow/models/block.py
+++ b/skyvern/forge/sdk/workflow/models/block.py
@@ -184,8 +184,18 @@ class Block(BaseModel, abc.ABC):
             return potential_template
         template = Template(potential_template)
 
+        block_reference_data: dict[str, Any] = workflow_run_context.get_block_metadata(self.label)
         template_data = workflow_run_context.values.copy()
-        template_data[self.label] = workflow_run_context.get_block_metadata(self.label)
+        if self.label in template_data:
+            current_value = template_data[self.label]
+            if isinstance(current_value, dict):
+                block_reference_data.update(current_value)
+            else:
+                LOG.warning(
+                    f"Parameter {self.label} has a registered reference value, going to overwrite it by block metadata"
+                )
+
+        template_data[self.label] = block_reference_data
         return template.render(template_data)
 
     @classmethod