fix: count a run completed at the DONE transition, not after it (researchx#28) #29

Merged
jonathan merged 2 commits from po/fix-cost-counter-race into main 2026-10-27 22:36:43 +00:00
2 changed files with 7 additions and 3 deletions
+2 -3
View File
@@ -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)
+5
View File
@@ -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: