fix(automation): persist cancelled scheduled runs

coworker/server/manager.py

@@ -4752,58 +4752,69 @@ async def _run_scheduled_task(self, task, trigger: str) -> TaskRun:
             task_id=task.id, trigger=trigger
         )  # __post_init__ sets run.session_id
         self.task_store.add_run(run)  # mark "running"
-        # UX-026: tell every open app window a SCHEDULED run just started (the 5s
-        # top-right toast). Manual runs never come through here — the user is
-        # already watching those live.
-        await self.broadcast_event(
-            {
-                "type": "automation_run_started",
-                "data": {
-                    "task_id": task.id,
-                    "task_title": task.title,
-                    "session_id": run.session_id,
-                    "workspace": task.workspace,
-                    "agent": task.agent,
-                    "trigger": trigger,
-                },
-            }
-        )
-        # Each run is a real, persisted conversation thread: it runs the instructions under its
-        # own session id, then saves the transcript. The user can reopen that session and ask a
-        # follow-up — the scheduled agent is no longer fire-and-forget.
-        engine = self._build_task_engine(task, session_id=run.session_id)
-        # Register the live engine up-front: a parked approval persists the session
-        # mid-run (durable suspend), and resolving from the Inbox must find this engine.
-        self._engines[run.session_id] = engine
-        # The first turn is the task itself. The framing matters: instructions often restate the
-        # schedule ("every day at 5:32pm…"), so make explicit that the schedule already fired and
-        # the job now is to execute, not to (re)schedule.
-        opening = (
-            f"⏰ Scheduled run — {task.title}\n\n"
-            "This automation is due now: carry out the task below immediately and produce the "
-            "result. The schedule already exists — do not create or modify any scheduled tasks.\n\n"
-            f"{task.instructions}"
-        )
+        engine = None
         try:
+            # UX-026: tell every open app window a SCHEDULED run just started (the 5s
+            # top-right toast). Manual runs never come through here — the user is
+            # already watching those live.
+            await self.broadcast_event(
+                {
+                    "type": "automation_run_started",
+                    "data": {
+                        "task_id": task.id,
+                        "task_title": task.title,
+                        "session_id": run.session_id,
+                        "workspace": task.workspace,
+                        "agent": task.agent,
+                        "trigger": trigger,
+                    },
+                }
+            )
+            # Each run is a real, persisted conversation thread: it runs the instructions under
+            # its own session id, then saves the transcript. The user can reopen that session and
+            # ask a follow-up — the scheduled agent is no longer fire-and-forget.
+            engine = self._build_task_engine(task, session_id=run.session_id)
+            # Register the live engine up-front: a parked approval persists the session
+            # mid-run (durable suspend), and resolving from the Inbox must find this engine.
+            self._engines[run.session_id] = engine
+            # Instructions often restate the schedule, so make explicit that it already fired.
+            opening = (
+                f"⏰ Scheduled run — {task.title}\n\n"
+                "This automation is due now: carry out the task below immediately and produce "
+                "the result. The schedule already exists — do not create or modify any "
+                "scheduled tasks.\n\n"
+                f"{task.instructions}"
+            )
             async for _event in engine.run(opening):
                 pass
             run.result_text = _last_assistant_text(engine.messages)
             run.artifacts = _recent_files(task.workspace, since=run.started_at)
             run.status = "ok"
-            if task.notify_on_completion:
-                await self._notify_task_done(task, run)
+        except asyncio.CancelledError:
+            if engine is not None:
+                engine.request_interrupt()
+            run.status, run.error = "error", "cancelled during scheduler shutdown"
+            raise
         except Exception as exc:
             run.status, run.error = "error", str(exc)
         finally:
             run.finished_at = _epoch()
             # Persist the run as a continuable session + keep the live engine for an immediate
             # follow-up; record the run (now carrying its session_id).
+            if engine is not None:
+                try:
+                    self.save(run.session_id, engine)
+                    self._engines[run.session_id] = engine
+                except Exception:
+                    pass
+            self.task_store.add_run(run)
+        if task.notify_on_completion and run.status == "ok":
             try:
-                self.save(run.session_id, engine)
-                self._engines[run.session_id] = engine
-            except Exception:
+                await self._notify_task_done(task, run)
+            except asyncio.CancelledError:
+                # Execution and persistence completed. Preserve success so Scheduler.run_task
+                # advances the task instead of repeating side effects after restart.
                 pass
-            self.task_store.add_run(run)
         return run
 
     async def _notify_task_done(self, task, run: TaskRun) -> None:

tests/test_automation.py

@@ -476,3 +476,128 @@ async def dead(message):
     assert event["data"]["session_id"] == run.session_id
     assert event["data"]["trigger"] == "schedule"
     assert dead not in manager._event_clients  # dropped, not fatal
+
+
+@pytest.mark.asyncio
+async def test_cancelled_scheduled_run_is_not_left_running(tmp_path, monkeypatch):
+    from coworker.server.manager import SessionManager
+
+    monkeypatch.setenv("COWORKER_STATE_DIR", str(tmp_path / "state"))
+    workspace = tmp_path / "ws"
+    workspace.mkdir()
+    manager = SessionManager(data_dir=tmp_path / "data")
+    task = manager.task_store.save(_task(workspace=str(workspace)))
+    started = asyncio.Event()
+    interrupted = False
+
+    class BlockingEngine:
+        messages = []
+
+        async def run(self, _opening):
+            started.set()
+            await asyncio.Event().wait()
+            yield
+
+        def request_interrupt(self):
+            nonlocal interrupted
+            interrupted = True
+
+    monkeypatch.setattr(
+        manager,
+        "_build_task_engine",
+        lambda _task, *, session_id: BlockingEngine(),
+    )
+
+    pending = asyncio.create_task(
+        manager._run_scheduled_task(task, trigger="schedule")
+    )
+    await asyncio.wait_for(started.wait(), timeout=1.0)
+    (initial,) = manager.task_store.runs(task.id)
+    assert initial.status == "running"
+
+    pending.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await pending
+
+    (stored,) = manager.task_store.runs(task.id)
+    assert stored.run_id == initial.run_id
+    assert stored.status == "error"
+    assert stored.finished_at is not None
+    assert stored.error == "cancelled during scheduler shutdown"
+    assert interrupted is True
+
+
+@pytest.mark.asyncio
+async def test_cancelled_start_broadcast_is_not_left_running(tmp_path, monkeypatch):
+    from coworker.server.manager import SessionManager
+
+    monkeypatch.setenv("COWORKER_STATE_DIR", str(tmp_path / "state"))
+    workspace = tmp_path / "ws"
+    workspace.mkdir()
+    manager = SessionManager(data_dir=tmp_path / "data")
+    task = manager.task_store.save(_task(workspace=str(workspace)))
+    broadcast_started = asyncio.Event()
+
+    async def blocking_listener(_message):
+        broadcast_started.set()
+        await asyncio.Event().wait()
+
+    manager.register_event_client(blocking_listener)
+    pending = asyncio.create_task(
+        manager._run_scheduled_task(task, trigger="schedule")
+    )
+    await asyncio.wait_for(broadcast_started.wait(), timeout=1.0)
+    pending.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await pending
+
+    (stored,) = manager.task_store.runs(task.id)
+    assert stored.status == "error"
+    assert stored.finished_at is not None
+    assert stored.error == "cancelled during scheduler shutdown"
+
+
+@pytest.mark.asyncio
+async def test_cancelled_completion_notification_preserves_success(
+    tmp_path, monkeypatch
+):
+    from coworker.server.manager import SessionManager
+
+    monkeypatch.setenv("COWORKER_STATE_DIR", str(tmp_path / "state"))
+    workspace = tmp_path / "ws"
+    workspace.mkdir()
+    manager = SessionManager(data_dir=tmp_path / "data")
+    task = manager.task_store.save(_task(workspace=str(workspace)))
+    notify_started = asyncio.Event()
+
+    class CompletingEngine:
+        messages = [{"role": "assistant", "content": "done"}]
+
+        async def run(self, _opening):
+            if False:
+                yield
+
+    async def blocking_notify(_task, _run):
+        notify_started.set()
+        await asyncio.Event().wait()
+
+    monkeypatch.setattr(
+        manager,
+        "_build_task_engine",
+        lambda _task, *, session_id: CompletingEngine(),
+    )
+    monkeypatch.setattr(manager, "_notify_task_done", blocking_notify)
+    scheduler = Scheduler(manager.task_store, manager._run_scheduled_task)
+    pending = asyncio.create_task(scheduler.run_task(task, trigger="schedule"))
+
+    await asyncio.wait_for(notify_started.wait(), timeout=1.0)
+    pending.cancel()
+    run = await pending
+
+    assert run is not None and run.status == "ok"
+    stored = manager.task_store.find_run(run.run_id)
+    assert stored is not None and stored.status == "ok"
+    updated = manager.task_store.get(task.id)
+    assert updated is not None
+    assert updated.run_count == 1
+    assert updated.last_status == "ok"