From ed36ca80d35660aa59759a57563dc50e092b2edf Mon Sep 17 00:00:00 2001
From: Raphael Fakhri <153192858+RaphaelFakhri@users.noreply.github.com>
Date: Tue, 29 Sep 2026 10:02:22 +0000
Subject: [PATCH] fix(healing): read YAML workflow files in
 process_workflow_file_with_markers

---
 .../healing/tests/test_variable_utils.py      | 53 +++++++++++++++++++
 .../workflow_use/healing/variable_utils.py    | 13 +++--
 2 files changed, 62 insertions(+), 4 deletions(-)
 create mode 100644 workflows/workflow_use/healing/tests/test_variable_utils.py

diff --git a/workflows/workflow_use/healing/tests/test_variable_utils.py b/workflows/workflow_use/healing/tests/test_variable_utils.py
new file mode 100644
index 00000000..ccbe1134
--- /dev/null
+++ b/workflows/workflow_use/healing/tests/test_variable_utils.py
@@ -0,0 +1,53 @@
+"""Tests for the variable marker file utilities."""
+
+import json
+
+import yaml
+
+from workflow_use.healing.variable_utils import process_workflow_file_with_markers
+
+WORKFLOW = {
+	'name': 'Test Workflow',
+	'description': 'Test',
+	'version': '1.0',
+	'steps': [
+		{'type': 'input', 'target_text': 'Email', 'value': 'VAR:user_email:test@example.com'},
+		{'type': 'extract', 'extractionGoal': 'Extract the result'},
+	],
+	'input_schema': [],
+}
+
+
+def test_process_yaml_workflow_file(tmp_path):
+	"""A .workflow.yaml file, the format the storage service writes, is processed in place."""
+	path = tmp_path / 'test.workflow.yaml'
+	path.write_text(yaml.dump(WORKFLOW, sort_keys=False))
+
+	process_workflow_file_with_markers(path)
+
+	saved = yaml.safe_load(path.read_text())
+	assert saved['steps'][0]['value'] == '{user_email}'
+	assert [inp['name'] for inp in saved['input_schema']] == ['user_email']
+
+
+def test_process_json_workflow_file(tmp_path):
+	"""A JSON workflow file still works and stays JSON."""
+	path = tmp_path / 'test.workflow.json'
+	path.write_text(json.dumps(WORKFLOW))
+
+	process_workflow_file_with_markers(path)
+
+	saved = json.loads(path.read_text())
+	assert saved['steps'][0]['value'] == '{user_email}'
+
+
+def test_yaml_input_to_json_output(tmp_path):
+	"""The output format follows the output path extension."""
+	source = tmp_path / 'test.workflow.yaml'
+	source.write_text(yaml.dump(WORKFLOW, sort_keys=False))
+	output = tmp_path / 'out.json'
+
+	process_workflow_file_with_markers(source, output)
+
+	saved = json.loads(output.read_text())
+	assert saved['steps'][0]['value'] == '{user_email}'
diff --git a/workflows/workflow_use/healing/variable_utils.py b/workflows/workflow_use/healing/variable_utils.py
index 298f5451..b1826ef7 100644
--- a/workflows/workflow_use/healing/variable_utils.py
+++ b/workflows/workflow_use/healing/variable_utils.py
@@ -4,6 +4,8 @@
 from pathlib import Path
 from typing import Optional
 
+import yaml
+
 from workflow_use.healing.variable_extractor import VariableExtractor
 from workflow_use.schema.views import WorkflowDefinitionSchema
 
@@ -57,9 +59,9 @@ def process_workflow_file_with_markers(
 	input_path = Path(input_path)
 	output_path = Path(output_path) if output_path else input_path
 
-	# Load the workflow
+	# Load the workflow. YAML is a superset of JSON, so this reads both formats.
 	with open(input_path, 'r') as f:
-		workflow_data = json.load(f)
+		workflow_data = yaml.safe_load(f)
 
 	workflow = WorkflowDefinitionSchema(**workflow_data)
 
@@ -67,9 +69,12 @@ def process_workflow_file_with_markers(
 	extractor = VariableExtractor()
 	updated_workflow, extracted_inputs = extractor.process_workflow_with_markers(workflow)
 
-	# Save the updated workflow
+	# Save the updated workflow in the format that matches the output file extension
 	with open(output_path, 'w') as f:
-		json.dump(updated_workflow.model_dump(), f, indent=2)
+		if output_path.suffix.lower() in ('.yaml', '.yml'):
+			yaml.dump(updated_workflow.model_dump(mode='json'), f, default_flow_style=False, sort_keys=False)
+		else:
+			json.dump(updated_workflow.model_dump(), f, indent=2)
 
 	print(f'Processed workflow: {input_path}')
 	if extracted_inputs: