diff --git a/coworker/server/manager.py b/coworker/server/manager.py index 1486d9064..1cf41112a 100644 --- a/coworker/server/manager.py +++ b/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: diff --git a/tests/test_automation.py b/tests/test_automation.py index 4503f22ca..3a0134935 100644 --- a/tests/test_automation.py +++ b/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"