diff --git a/document-parser/services/analysis_service.py b/document-parser/services/analysis_service.py index 36c2f5a..f003285 100644 --- a/document-parser/services/analysis_service.py +++ b/document-parser/services/analysis_service.py @@ -53,6 +53,7 @@ class AnalysisService: self._chunker = chunker self._conversion_timeout = conversion_timeout self._semaphore = asyncio.Semaphore(max_concurrent) + self._running_tasks: dict[str, asyncio.Task] = {} async def create( self, @@ -79,7 +80,8 @@ class AnalysisService: chunking_options, ) ) - task.add_done_callback(functools.partial(_on_task_done, job_id=job.id)) + self._running_tasks[job.id] = task + task.add_done_callback(functools.partial(self._on_task_done, job_id=job.id)) return job @@ -92,7 +94,11 @@ class AnalysisService: return await analysis_repo.find_by_id(job_id) async def delete(self, job_id: str) -> bool: - """Delete an analysis job. Returns True if it existed.""" + """Delete an analysis job, cancelling any running task. Returns True if it existed.""" + task = self._running_tasks.pop(job_id, None) + if task and not task.done(): + task.cancel() + logger.info("Cancelled running task for job %s", job_id) return await analysis_repo.delete(job_id) async def rechunk(self, job_id: str, chunking_options: dict) -> list[ChunkResult]: @@ -115,6 +121,11 @@ class AnalysisService: return chunks + def _on_task_done(self, task: asyncio.Task, *, job_id: str) -> None: + """Cleanup running tasks and delegate to module-level handler.""" + self._running_tasks.pop(job_id, None) + _on_task_done(task, job_id=job_id) + async def _run_analysis( self, job_id: str, @@ -235,6 +246,9 @@ def _on_task_done(task: asyncio.Task, *, job_id: str) -> None: _schedule_mark_failed(job_id, _classify_error(exc)) +# Keep the module-level function as the default, but AnalysisService uses its own method. + + def _schedule_mark_failed(job_id: str, error: str) -> None: """Schedule _mark_failed as a tracked background task.""" t = asyncio.ensure_future(_mark_failed(job_id, error)) diff --git a/document-parser/tests/test_analysis_service.py b/document-parser/tests/test_analysis_service.py index 66cbaac..a211dee 100644 --- a/document-parser/tests/test_analysis_service.py +++ b/document-parser/tests/test_analysis_service.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import functools from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -90,6 +91,60 @@ class TestOnTaskDone: mock_mark.assert_not_called() +class TestAnalysisServiceCancellation: + """Verify delete cancels running tasks.""" + + @pytest.mark.asyncio + async def test_delete_cancels_running_task(self): + """Deleting a job while running should cancel its task.""" + converter = MagicMock() + service = AnalysisService(converter=converter) + + blocker = asyncio.Event() + + async def slow_analysis(): + await blocker.wait() + + task = asyncio.create_task(slow_analysis()) + service._running_tasks["j1"] = task + + with patch("services.analysis_service.analysis_repo") as mock_repo: + mock_repo.delete = AsyncMock(return_value=True) + result = await service.delete("j1") + + assert result is True + assert task.cancelling() or task.cancelled() + assert "j1" not in service._running_tasks + + @pytest.mark.asyncio + async def test_delete_completed_job_no_error(self): + """Deleting a completed job should not raise even if no task tracked.""" + converter = MagicMock() + service = AnalysisService(converter=converter) + + with patch("services.analysis_service.analysis_repo") as mock_repo: + mock_repo.delete = AsyncMock(return_value=True) + result = await service.delete("j-gone") + + assert result is True + + @pytest.mark.asyncio + async def test_task_cleaned_from_running_on_completion(self): + """After a task completes, it should be removed from _running_tasks.""" + converter = MagicMock() + service = AnalysisService(converter=converter) + + async def instant(): + pass + + task = asyncio.create_task(instant()) + service._running_tasks["j1"] = task + task.add_done_callback(functools.partial(service._on_task_done, job_id="j1")) + await task + + assert "j1" not in service._running_tasks + + class TestAnalysisServiceConcurrency: """Verify that the semaphore limits concurrent analysis jobs."""