From 8a2bb7c20aecd47683e1c5ca1463a95840dc15ab Mon Sep 17 00:00:00 2001
From: Richard Yu <richard95yu@gmail.com>
Date: Wed, 3 Dec 2025 20:14:55 -0800
Subject: [PATCH] set workflow id

---
 comfy_execution/jobs.py      |  8 ++++++--
 tests/execution/test_jobs.py | 12 ++++++++++--
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/comfy_execution/jobs.py b/comfy_execution/jobs.py
index 01da70517..ac70fb1e4 100644
--- a/comfy_execution/jobs.py
+++ b/comfy_execution/jobs.py
@@ -52,6 +52,8 @@ def normalize_queue_item(item, status):
     """Convert queue item tuple to unified job dict."""
     _, prompt_id, _, extra_data, _ = item[:5]
     create_time = extra_data.get('create_time')
+    extra_pnginfo = extra_data.get('extra_pnginfo', {}) or {}
+    workflow_id = extra_pnginfo.get('workflow', {}).get('id')
 
     return {
         'id': prompt_id,
@@ -63,7 +65,7 @@ def normalize_queue_item(item, status):
         'execution_end_time': None,
         'outputs_count': 0,
         'preview_output': None,
-        'workflow_id': None,
+        'workflow_id': workflow_id,
     }
 
 
@@ -72,6 +74,8 @@ def normalize_history_item(prompt_id, history_item, include_outputs=False):
     prompt_tuple = history_item['prompt']
     _, _, prompt, extra_data, _ = prompt_tuple[:5]
     create_time = extra_data.get('create_time')
+    extra_pnginfo = extra_data.get('extra_pnginfo', {}) or {}
+    workflow_id = extra_pnginfo.get('workflow', {}).get('id')
 
     status_info = history_item.get('status', {})
     status_str = status_info.get('status_str') if status_info else None
@@ -114,7 +118,7 @@ def normalize_history_item(prompt_id, history_item, include_outputs=False):
         'execution_end_time': execution_end_time,
         'outputs_count': outputs_count,
         'preview_output': preview_output,
-        'workflow_id': None,
+        'workflow_id': workflow_id,
     }
 
     if include_outputs:
diff --git a/tests/execution/test_jobs.py b/tests/execution/test_jobs.py
index c08e07b2b..30e64b1e7 100644
--- a/tests/execution/test_jobs.py
+++ b/tests/execution/test_jobs.py
@@ -246,7 +246,10 @@ class TestNormalizeQueueItem:
             10,  # priority/number
             'prompt-123',  # prompt_id
             {'nodes': {}},  # prompt
-            {'create_time': 1234567890},  # extra_data
+            {
+                'create_time': 1234567890,
+                'extra_pnginfo': {'workflow': {'id': 'workflow-abc'}}
+            },  # extra_data
             ['node1'],  # outputs_to_execute
         )
         job = normalize_queue_item(item, JobStatus.PENDING)
@@ -258,6 +261,7 @@ class TestNormalizeQueueItem:
         assert job['execution_end_time'] is None
         assert job['error_message'] is None
         assert job['outputs_count'] == 0
+        assert job['workflow_id'] == 'workflow-abc'
 
 
 class TestNormalizeHistoryItem:
@@ -270,7 +274,10 @@ class TestNormalizeHistoryItem:
                 5,  # priority
                 'prompt-456',
                 {'nodes': {}},
-                {'create_time': 1234567890000},  # milliseconds
+                {
+                    'create_time': 1234567890000,
+                    'extra_pnginfo': {'workflow': {'id': 'workflow-xyz'}}
+                },  # milliseconds
                 ['node1'],
             ),
             'status': {'status_str': 'success', 'completed': True, 'messages': []},
@@ -283,6 +290,7 @@ class TestNormalizeHistoryItem:
         assert job['status'] == 'completed'
         assert job['execution_start_time'] == 1234567890000
         assert job['execution_end_time'] == 1234567890000 + 2500  # +2.5 seconds in ms
+        assert job['workflow_id'] == 'workflow-xyz'
 
     def test_failed_job(self):
         """Failed history item should have failed status and message."""