Skip to content

Commit 2d39630

Browse files
committed
feat: add cancelled filter to /jobs
1 parent 2c03884 commit 2d39630

File tree

2 files changed

+55
-14
lines changed

2 files changed

+55
-14
lines changed

comfy_execution/jobs.py

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,9 @@ class JobStatus:
1414
IN_PROGRESS = 'in_progress'
1515
COMPLETED = 'completed'
1616
FAILED = 'failed'
17+
CANCELLED = 'cancelled'
1718

18-
ALL = [PENDING, IN_PROGRESS, COMPLETED, FAILED]
19+
ALL = [PENDING, IN_PROGRESS, COMPLETED, FAILED, CANCELLED]
1920

2021

2122
# Media types that can be previewed in the frontend
@@ -94,19 +95,14 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
9495

9596
status_info = history_item.get('status', {})
9697
status_str = status_info.get('status_str') if status_info else None
97-
if status_str == 'success':
98-
status = JobStatus.COMPLETED
99-
elif status_str == 'error':
100-
status = JobStatus.FAILED
101-
else:
102-
status = JobStatus.COMPLETED
10398

10499
outputs = history_item.get('outputs', {})
105100
outputs_count, preview_output = get_outputs_summary(outputs)
106101

107102
execution_error = None
108103
execution_start_time = None
109104
execution_end_time = None
105+
was_interrupted = False
110106
if status_info:
111107
messages = status_info.get('messages', [])
112108
for entry in messages:
@@ -119,6 +115,15 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
119115
execution_end_time = event_data.get('timestamp')
120116
if event_name == 'execution_error':
121117
execution_error = event_data
118+
elif event_name == 'execution_interrupted':
119+
was_interrupted = True
120+
121+
if status_str == 'success':
122+
status = JobStatus.COMPLETED
123+
elif status_str == 'error':
124+
status = JobStatus.CANCELLED if was_interrupted else JobStatus.FAILED
125+
else:
126+
status = JobStatus.COMPLETED
122127

123128
job = prune_dict({
124129
'id': prompt_id,
@@ -268,13 +273,13 @@ def get_all_jobs(
268273
for item in queued:
269274
jobs.append(normalize_queue_item(item, JobStatus.PENDING))
270275

271-
include_completed = JobStatus.COMPLETED in status_filter
272-
include_failed = JobStatus.FAILED in status_filter
273-
if include_completed or include_failed:
276+
history_statuses = {JobStatus.COMPLETED, JobStatus.FAILED, JobStatus.CANCELLED}
277+
requested_history_statuses = history_statuses & set(status_filter)
278+
if requested_history_statuses:
274279
for prompt_id, history_item in history.items():
275-
is_failed = history_item.get('status', {}).get('status_str') == 'error'
276-
if (is_failed and include_failed) or (not is_failed and include_completed):
277-
jobs.append(normalize_history_item(prompt_id, history_item))
280+
job = normalize_history_item(prompt_id, history_item)
281+
if job.get('status') in requested_history_statuses:
282+
jobs.append(job)
278283

279284
if workflow_id:
280285
jobs = [j for j in jobs if j.get('workflow_id') == workflow_id]

tests/execution/test_jobs.py

Lines changed: 37 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,14 +19,16 @@ def test_status_values(self):
1919
assert JobStatus.IN_PROGRESS == 'in_progress'
2020
assert JobStatus.COMPLETED == 'completed'
2121
assert JobStatus.FAILED == 'failed'
22+
assert JobStatus.CANCELLED == 'cancelled'
2223

2324
def test_all_contains_all_statuses(self):
2425
"""ALL should contain all status values."""
2526
assert JobStatus.PENDING in JobStatus.ALL
2627
assert JobStatus.IN_PROGRESS in JobStatus.ALL
2728
assert JobStatus.COMPLETED in JobStatus.ALL
2829
assert JobStatus.FAILED in JobStatus.ALL
29-
assert len(JobStatus.ALL) == 4
30+
assert JobStatus.CANCELLED in JobStatus.ALL
31+
assert len(JobStatus.ALL) == 5
3032

3133

3234
class TestIsPreviewable:
@@ -336,6 +338,40 @@ def test_failed_job(self):
336338
assert job['execution_error']['node_type'] == 'KSampler'
337339
assert job['execution_error']['exception_message'] == 'CUDA out of memory'
338340

341+
def test_cancelled_job(self):
342+
"""Cancelled/interrupted history item should have cancelled status."""
343+
history_item = {
344+
'prompt': (
345+
5,
346+
'prompt-cancelled',
347+
{'nodes': {}},
348+
{'create_time': 1234567890000},
349+
['node1'],
350+
),
351+
'status': {
352+
'status_str': 'error',
353+
'completed': False,
354+
'messages': [
355+
('execution_start', {'prompt_id': 'prompt-cancelled', 'timestamp': 1234567890500}),
356+
('execution_interrupted', {
357+
'prompt_id': 'prompt-cancelled',
358+
'node_id': '5',
359+
'node_type': 'KSampler',
360+
'executed': ['1', '2', '3'],
361+
'timestamp': 1234567891000,
362+
})
363+
]
364+
},
365+
'outputs': {},
366+
}
367+
368+
job = normalize_history_item('prompt-cancelled', history_item)
369+
assert job['status'] == 'cancelled'
370+
assert job['execution_start_time'] == 1234567890500
371+
assert job['execution_end_time'] == 1234567891000
372+
# Cancelled jobs should not have execution_error set
373+
assert 'execution_error' not in job
374+
339375
def test_include_outputs(self):
340376
"""When include_outputs=True, should include full output data."""
341377
history_item = {

0 commit comments

Comments
 (0)