diff --git a/backend/src/researchx/api/app.py b/backend/src/researchx/api/app.py index ffc4271..e5488b7 100644 --- a/backend/src/researchx/api/app.py +++ b/backend/src/researchx/api/app.py @@ -226,6 +226,7 @@ def create_app( storage=storage, registry=registry, event_bus=bus, + on_complete=runs.note_completed, ) probe_data: dict[str, Any] | None = None if supports_probe(model): @@ -240,9 +241,7 @@ def create_app( runs.add(run_id, actor, task, bus) async def _watch() -> None: - state = await task - if state.state == RunStateName.DONE: - runs.note_completed() + await task asyncio.create_task(_watch()) return _summary(run_id) diff --git a/backend/src/researchx/engine/actor.py b/backend/src/researchx/engine/actor.py index 7b54f9e..edbed81 100644 --- a/backend/src/researchx/engine/actor.py +++ b/backend/src/researchx/engine/actor.py @@ -1,4 +1,5 @@ import asyncio +from collections.abc import Callable from datetime import UTC, datetime from typing import Any @@ -34,6 +35,7 @@ class RunActor: storage: LocalStorageAdapter, registry: MethodologyRegistry, event_bus: EventBus, + on_complete: Callable[[], None] | None = None, ) -> None: self._config = run_config self._model = model @@ -41,6 +43,7 @@ class RunActor: self._storage = storage self._registry = registry self._bus = event_bus + self._on_complete = on_complete self._store = EvidenceStore() self.state = RunState( run_id=run_config.run_id, @@ -83,6 +86,8 @@ class RunActor: async def _checkpoint(self, target: RunStateName) -> None: if self.state.state != target: self.state.transition(target) + if target is RunStateName.DONE and self._on_complete is not None: + self._on_complete() await self._persist() async def _wait_if_paused(self) -> None: