From d9684392817249e84eb666e304510defe4ff78c9 Mon Sep 17 00:00:00 2001 From: kevin9327 Date: Sun, 30 Aug 2026 15:26:00 +0900 Subject: [PATCH] fix(workflow-runs): take part in the job TTL purge so the headless API doesn't leak MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The /workflow-runs surface shares _jobs, _cancel_events and _completed_at with the /generate endpoints, but never participated in their TTL purge: - create_run_from_image never called _purge_old_jobs(), so terminal job records accumulated indefinitely unless a /generate/from-image call happened to sweep them. That is the opposite of the intended use — /workflow-runs is the headless automation surface, where nothing else triggers the purge. - cancel_run set status="cancelled" but never stamped _completed_at, so _purge_old_jobs() (which only sweeps entries that have a completion time) could never evict a cancelled run — a permanent leak of a JobStatus + threading.Event per cancellation. Mirror what cancel_job and generate_from_image already do: purge on create, and record the completion time on cancel. Collection routing is intentionally left untouched here. Adds api/tests/test_workflow_runs_lifecycle.py (both cases fail before, pass after). Co-Authored-By: Claude Opus 4.8 --- api/routers/workflow_runs.py | 6 ++ api/tests/test_workflow_runs_lifecycle.py | 96 +++++++++++++++++++++++ 2 files changed, 102 insertions(+) create mode 100644 api/tests/test_workflow_runs_lifecycle.py diff --git a/api/routers/workflow_runs.py b/api/routers/workflow_runs.py index cb47bc92..9b78a205 100644 --- a/api/routers/workflow_runs.py +++ b/api/routers/workflow_runs.py @@ -1,5 +1,6 @@ import json import threading +import time import uuid from typing import Optional from fastapi import APIRouter, BackgroundTasks, File, Form, HTTPException, UploadFile @@ -9,7 +10,9 @@ VALID_REMESH_MODES, _cancel_events, _cancelled, + _completed_at, _jobs, + _purge_old_jobs, _run_generation, sanitize_collection, ) @@ -75,6 +78,8 @@ async def create_run_from_image( job_id = str(uuid.uuid4()) image_bytes = await image.read() + _purge_old_jobs() + _jobs[job_id] = JobStatus(job_id=job_id, status="pending", progress=0) _cancel_events[job_id] = threading.Event() @@ -115,6 +120,7 @@ async def cancel_run(run_id: str): _cancel_events[run_id].set() if job.status in ("pending", "running"): job.status = "cancelled" + _completed_at[run_id] = time.monotonic() try: gen = generator_registry._generators.get(generator_registry._active_id) diff --git a/api/tests/test_workflow_runs_lifecycle.py b/api/tests/test_workflow_runs_lifecycle.py new file mode 100644 index 00000000..f8fe7371 --- /dev/null +++ b/api/tests/test_workflow_runs_lifecycle.py @@ -0,0 +1,96 @@ +import asyncio +import threading +import time +import unittest + +from fastapi import BackgroundTasks + +import routers.generation as generation +import routers.workflow_runs as workflow_runs +from schemas.generation import JobStatus + + +class _FakeUpload: + """Minimal UploadFile stand-in: an image content-type and readable bytes.""" + + def __init__(self, content_type: str = "image/png", data: bytes = b"\x89PNG\r\n") -> None: + self.content_type = content_type + self._data = data + + async def read(self) -> bytes: + return self._data + + +class _FakeRegistry: + """Accepts any model id and exposes the attrs cancel_run pokes at.""" + + _generators: dict = {} + _active_id = None + + def get_generator(self, model_id: str) -> object: + return object() + + def switch_model(self, model_id: str) -> None: + pass + + +def _clear_job_stores() -> None: + for store in ( + generation._jobs, + generation._cancel_events, + generation._cancelled, + generation._completed_at, + ): + store.clear() + + +class WorkflowRunJobLifecycleTests(unittest.TestCase): + """The headless /workflow-runs surface shares the job dicts with /generate, + so it must take part in the same TTL purge — otherwise long-running + automation leaks a JobStatus + Event per run forever.""" + + def setUp(self) -> None: + self._prev = workflow_runs.generator_registry + workflow_runs.generator_registry = _FakeRegistry() + _clear_job_stores() + + def tearDown(self) -> None: + workflow_runs.generator_registry = self._prev + _clear_job_stores() + + def test_create_run_purges_terminal_jobs_past_ttl(self) -> None: + stale = "stale-run" + generation._jobs[stale] = JobStatus(job_id=stale, status="done", progress=100) + generation._cancel_events[stale] = threading.Event() + generation._completed_at[stale] = time.monotonic() - generation._JOB_TTL - 1 + + background = BackgroundTasks() + asyncio.run( + workflow_runs.create_run_from_image( + background, + image=_FakeUpload(), + model_id="sf3d", + collection="Default", + params="{}", + ) + ) + + # Before the fix create_run_from_image never purged, so the stale job lingered. + self.assertNotIn(stale, generation._jobs) + self.assertNotIn(stale, generation._completed_at) + self.assertNotIn(stale, generation._cancel_events) + + def test_cancel_run_records_completion_so_it_can_be_purged(self) -> None: + run_id = "run-1" + generation._jobs[run_id] = JobStatus(job_id=run_id, status="running", progress=10) + generation._cancel_events[run_id] = threading.Event() + + asyncio.run(workflow_runs.cancel_run(run_id)) + + self.assertEqual(generation._jobs[run_id].status, "cancelled") + # Without a _completed_at stamp the purge sweep can never evict a cancelled run. + self.assertIn(run_id, generation._completed_at) + + +if __name__ == "__main__": + unittest.main()