Add cancel support for running reports (#1819)

Author: ghostiee-11

lumen/ai/report.py

@@ -81,7 +81,7 @@ class Task(Viewer):
     running = param.Boolean(doc="""
         Whether the task is currently running.""")
 
-    status = param.Selector(objects=["idle", "running", "success", "error"], default="idle", doc="""
+    status = param.Selector(objects=["idle", "running", "success", "error", "cancelled"], default="idle", doc="""
         The current status of the task.""")
 
     steps_layout = param.ClassSelector(default=None, class_=(ListLike, NamedListLike), allow_None=True, doc="""
@@ -230,11 +230,14 @@ async def execute(self, context: TContext | None = None, **kwargs) -> tuple[list
             if not self._prepared:
                 await self.prepare(context)
             views, out_context = await self._execute(context, **kwargs)
+        except asyncio.CancelledError:
+            self.status = "cancelled"
+            raise
         except Exception:
             self.status = "error"
             raise
         finally:
-            if self.status != "error":
+            if self.status not in ("error", "cancelled"):
                 self.status = "success"
             self.running = False
         self.out_context = out_context
@@ -411,6 +414,9 @@ async def _execute(self, context: TContext, **kwargs):
             new = []
             try:
                 new, new_context = await self._run_task(i, task, context, **kwargs)
+            except asyncio.CancelledError:
+                self.status = "cancelled"
+                raise
             except MissingContextError:
                 # Re-raise MissingContextError to allow retry logic at Plan level
                 raise
@@ -432,8 +438,8 @@ async def _execute(self, context: TContext, **kwargs):
                         break
                 views += new
             finally:
-                self._current = i + (0 if task.status == "error" else 1)
-        if self.status != "error":
+                self._current = i + (0 if task.status in ("error", "cancelled") else 1)
+        if self.status not in ("error", "cancelled"):
             self.status = "success"
         contexts = [self.context] if self.context else []
         contexts += [task.out_context for task in self]
@@ -845,6 +851,8 @@ class Report(TaskGroup):
     auto_execute = param.Boolean(default=False, doc="""
         If True, automatically execute the report on initialization.""")
 
+    _active_task = param.ClassSelector(class_=asyncio.Task, default=None, allow_None=True)
+
     _tasks = param.List(item_type=Section)
 
     level = 1
@@ -875,6 +883,12 @@ def _init_view(self):
         self._run = IconButton(
             icon="play_arrow", on_click=self._execute_event, margin=0, size="large",
             description="Execute Report", loading=self.param.running,
+            visible=self.param._active_task.rx.is_(None),
+        )
+        self._stop = IconButton(
+            icon="stop", on_click=self._handle_cancel, margin=0, size="large",
+            description="Stop Report", color="error",
+            visible=self.param._active_task.rx.is_not(None),
         )
         self._clear = IconButton(
             icon="clear", on_click=lambda _: self.reset(), margin=0, size="large",
@@ -926,6 +940,7 @@ def _init_view(self):
         self._menu = Row(
             self._header_title,
             self._run,
+            self._stop,
             self._clear,
             self._collapse,
             self._export,
@@ -935,6 +950,7 @@ def _init_view(self):
         self._dial = SpeedDial(
             items=[
                 {"label": "Execute Report", "icon": "play_arrow"},
+                {"label": "Stop Report", "icon": "stop"},
                 {"label": "Clear Report", "icon": "clear"},
                 {"label": "Export as Notebook", "icon": "description", "format": "ipynb"},
                 {"label": "Export as HTML", "icon": "language", "format": "html"},
@@ -972,7 +988,12 @@ def _init_view(self):
     async def _trigger_event(self, item: dict):
         icon = item["icon"]
         if icon == "play_arrow":
-            await self._execute_event()
+            if self._active_task is not None and not self._active_task.done():
+                self._handle_cancel()
+            else:
+                await self._execute_event()
+        elif icon == "stop":
+            self._handle_cancel()
         elif icon == "clear":
             self.reset()
         elif icon in ("description", "language"):
@@ -990,7 +1011,7 @@ def _update_run_state(self):
     @param.depends('status', 'views', watch=True)
     def _update_icon_visibility(self):
         """Show/hide icons based on whether report has outputs."""
-        has_outputs = self.status in ("success", "error") or bool(self.views)
+        has_outputs = self.status in ("success", "error", "cancelled") or bool(self.views)
         self._clear.visible = has_outputs
         self._collapse.visible = has_outputs
         self._export.visible = has_outputs
@@ -1008,11 +1029,27 @@ def _update_icon_visibility(self):
             }
 
     async def _execute_event(self, event=None):
+        if self._active_task is not None and not self._active_task.done():
+            # Already running; the stop button handles cancellation
+            return
         await asyncio.sleep(0.01)  # yield the event loop to allow button loading state to update
-        await self.execute()
+        task = asyncio.create_task(self.execute())
+        self._active_task = task
+        task.add_done_callback(self._on_execute_done)
+
+    def _on_execute_done(self, task):
+        """Clear the active task reference so UI controls flip back to idle."""
+        if task is not self._active_task:
+            return
+        self._active_task = None
+
+    def _handle_cancel(self, event=None):
+        """Cancel the in-flight report execution, if any."""
+        if self._active_task is not None and not self._active_task.done():
+            self._active_task.cancel()
 
     async def _export_report(self, item=None):
-        if len(self) and self.status != "success":
+        if len(self) and self.status not in ("success", "cancelled", "error"):
             await self.execute()
         fmt = item.get("format", "ipynb") if isinstance(item, dict) else "ipynb"
         title = self.title or "Report"
@@ -1035,7 +1072,7 @@ def _open_settings(self, event=None):
 
     def _populate_view(self):
         self._view[:] = objects = [(task.title, task) for task in self]
-        has_outputs = self.status in ("success", "error") or bool(self.views)
+        has_outputs = self.status in ("success", "error", "cancelled") or bool(self.views)
         if has_outputs:
             self._view.active = list(range(len(objects)))
         else:

lumen/tests/ai/test_report.py

@@ -392,3 +392,98 @@ async def test_report_to_html():
     assert "<html" in html_string.lower()
     assert "Hello Report" in html_string
     assert "Hello" in html_string
+
+
+class BlockingAction(Action):
+    """Action that waits on a signal before returning, so tests can deterministically cancel mid-execution."""
+
+    def __init__(self, started, release, **params):
+        super().__init__(**params)
+        self._started = started
+        self._release = release
+
+    async def _execute(self, context, **kwargs):
+        self._started.set()
+        await self._release.wait()
+        return [Markdown("should not render")], {"text": "never"}
+
+
+async def test_task_execute_marks_cancelled():
+    started = asyncio.Event()
+    release = asyncio.Event()
+    action = BlockingAction(started, release, title="Blocks")
+
+    task = asyncio.create_task(action.execute())
+    await started.wait()
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert action.status == "cancelled"
+    assert action.running is False
+
+
+async def test_taskgroup_cancel_propagates():
+    started = asyncio.Event()
+    release = asyncio.Event()
+    tg = TaskGroup(BlockingAction(started, release, title="Blocks"), B(), title="Group")
+
+    task = asyncio.create_task(tg.execute())
+    await started.wait()
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert tg.status == "cancelled"
+    assert tg[0].status == "cancelled"
+    assert tg[1].status == "idle"
+
+
+async def test_report_cancel_mid_execution():
+    started = asyncio.Event()
+    release = asyncio.Event()
+    report = Report(
+        Section(BlockingAction(started, release, title="Blocks"), title="S1"),
+        Section(B(), title="S2"),
+        title="Cancellable",
+    )
+
+    await report._execute_event()
+    await started.wait()
+    assert report._active_task is not None and not report._active_task.done()
+
+    report._handle_cancel()
+    # Wait for the task to finish reacting to the cancel
+    while report._active_task is not None:
+        await asyncio.sleep(0.01)
+
+    assert report.status == "cancelled"
+    assert report[0].status == "cancelled"
+    assert report[1].status == "idle"
+
+
+async def test_report_handle_cancel_idempotent():
+    report = Report(Section(HelloAction(), title="S1"), title="NoOp")
+    # No task running: _handle_cancel must not raise
+    report._handle_cancel()
+    assert report._active_task is None
+
+
+async def test_report_completed_sections_preserved_on_cancel():
+    started = asyncio.Event()
+    release = asyncio.Event()
+    report = Report(
+        Section(HelloAction(), title="Done"),
+        Section(BlockingAction(started, release, title="Blocks"), title="S2"),
+        title="Partial",
+    )
+
+    await report._execute_event()
+    await started.wait()
+    report._handle_cancel()
+    while report._active_task is not None:
+        await asyncio.sleep(0.01)
+
+    assert report[0].status == "success"
+    assert report[1].status == "cancelled"
+    assert report.status == "cancelled"