[BUG] Pending FutureResult in warmup causes a busy loop and duplicate…

statek/executors/utils.py

@@ -1001,6 +1001,20 @@ async def wrapped(job: Job, *args: Any, **kwargs: Any) -> Any:
     return wrapped
 
 
+def _release_ready_continuation(job: Job) -> bool:
+    """Return whether a job may run, releasing a satisfied continuation once."""
+    awaited_result = job.awaited_result
+    if awaited_result is None:
+        return True
+    if not awaited_result.check_condition():
+        return False
+
+    job.awaited_result = None
+    if job.status == JobStatus.SUSPENDED:
+        job.set_status(JobStatus.STARTED)
+    return True
+
+
 @_with_statek_job_context
 async def run_job_step(job: Job, provider: str = None) -> bool:
     """
@@ -1028,6 +1042,10 @@ async def run_job_step(job: Job, provider: str = None) -> bool:
     if job.status == JobStatus.DONE:
         return True
 
+    # Keep waiting continuations dormant until their recorded future is ready.
+    if not _release_ready_continuation(job):
+        return False
+
     # Build shared local_context for all context-aware functions in this step
     local_context = dict(job.py_env.local_state) if job.py_env.local_state else {}
 
@@ -1378,23 +1396,17 @@ def process_agent_events(
 
 def unsuspend_jobs():
     """
-    Review continuation conditions of suspended jobs and change their status to STARTED
-    where the continuation criteria are satisfied.
+    Review waiting jobs and release continuations whose criteria are satisfied.
 
-    This function finds all jobs with status SUSPENDED, checks if their awaited_result
-    condition is satisfied, and changes their status to STARTED if the condition is met.
+    Ordinary suspended jobs transition to STARTED when released. Warmup jobs retain
+    WARMING_UP so execution can continue in the active warmup block.
 
-    Note: This implementation might be slow for a large number of suspended jobs.
+    Note: This implementation might be slow for a large number of waiting jobs.
     Future versions will introduce a more robust Notifier engine and job expiration conditions.
     """
-    # Find all suspended jobs
-    suspended_jobs = db0.find(Job, JobStatus.SUSPENDED)
-    for job in suspended_jobs:
-        # Check if the job has an awaited_result and if its condition is satisfied
-        condition_met = job.awaited_result.check_condition()
-        if job.awaited_result is not None and condition_met:
-            # Change status from SUSPENDED to STARTED
-            job.set_status(JobStatus.STARTED)
+    waiting_jobs = db0.find(Job, [JobStatus.SUSPENDED, JobStatus.WARMING_UP])
+    for job in waiting_jobs:
+        _release_ready_continuation(job)
 
 
 def handle_critical_error(error: Optional[Exception] = None) -> None:
@@ -1477,8 +1489,16 @@ async def run_jobs_loop(max_concurrency: int = 100, provider: str = None,
             available_capacity = max_concurrency - len(pending_tasks)
             start_jobs_func(available_capacity)
 
-        ready_or_started_jobs = db0.filter(lambda found_job: found_job not in pending_tasks,
-                                            db0.find(Job, [JobStatus.READY, JobStatus.WARMING_UP, JobStatus.STARTED]))
+        ready_or_started_jobs = db0.filter(
+            lambda found_job: (
+                found_job not in pending_tasks
+                and not (
+                    found_job.status == JobStatus.WARMING_UP
+                    and found_job.awaited_result is not None
+                )
+            ),
+            db0.find(Job, [JobStatus.READY, JobStatus.WARMING_UP, JobStatus.STARTED]),
+        )
         # Make sure not to exceed max_concurrency
         if len(pending_tasks) < max_concurrency:
             for job in ready_or_started_jobs:

tests/test_run_job_step.py

@@ -216,9 +216,12 @@ async def test_resumes_from_future_error_line(
         assert job.next_instr_num == 2  # suspended at idx 2
         assert job.warmup_block_num is None  # block did not advance (None means block 0)
 
-        # Make the future ready and replace it in env
-        future_ready = create_future_ready(99)
-        job.py_env.local_state['future_val'] = future_ready
+        # Resolve the exact future recorded on the job.
+        future_not_ready.deps.value = 99
+        future_not_ready.set_complement_functions(
+            complement=_fetch_result_from_deps,
+            condition=_check_condition_true,
+        )
 
         # Second run: resumes from idx 2, skipping idx 0 and 1.
         # counter must stay 1 (not incremented again).

tests/test_warmup_blocks.py

@@ -308,9 +308,12 @@ async def test_run_job_step_multiple_blocks_continuation_after_future_error(
         # warmup_block_num should NOT advance since block didn't complete
         assert job.warmup_block_num == 1
 
-        # Make the future ready
-        future_ready = create_future_ready(42)
-        job.py_env.local_state['future_val'] = future_ready
+        # Resolve the exact future recorded on the job.
+        future_not_ready.deps.value = 42
+        future_not_ready.set_complement_functions(
+            complement=_fetch_result_from_deps,
+            condition=_check_condition_true,
+        )
 
         # Resume second block - should complete and advance
         result3 = await run_job_step(job)

tests/test_warmup_wait_scheduling.py

@@ -0,0 +1,170 @@
+"""Regression tests for scheduling warmups that await a FutureResult."""
+# pylint: disable=no-member
+
+import asyncio
+from contextlib import suppress
+from dataclasses import dataclass
+
+import dbzero as db0
+import pytest
+
+from statek.exceptions import FutureError
+from statek.executors.chat_log_item import WarmupLogItem
+from statek.executors.job import Job, JobStatus
+from statek.executors.utils import run_job_step, run_jobs_loop, unsuspend_jobs
+from statek.future import FutureResult
+
+
+@db0.memo
+@dataclass
+class WaitState:
+    """Mutable dependency used to control and count future readiness checks."""
+
+    ready: bool = False
+    value: str = "done"
+    condition_checks: int = 0
+
+
+def _check_wait_state(future_result: FutureResult) -> bool:
+    future_result.deps.condition_checks += 1
+    return future_result.deps.ready
+
+
+def _fetch_wait_state(future_result: FutureResult) -> str:
+    if not future_result.deps.ready:
+        raise FutureError(future_result=future_result)
+    return future_result.deps.value
+
+
+def _controlled_future(wait_state: WaitState) -> FutureResult:
+    future = FutureResult(deps=wait_state, state_num=0)
+    future.set_complement_functions(
+        complement=_fetch_wait_state,
+        condition=_check_wait_state,
+    )
+    return future
+
+
+def _warmup_items(job: Job) -> list[WarmupLogItem]:
+    return [item for item in job.chat_log if isinstance(item, WarmupLogItem)]
+
+
+async def _wait_until_done(job: Job) -> None:
+    while job.status != JobStatus.DONE:
+        await asyncio.sleep(0.01)
+
+
+class TestWarmupWaitReadiness:
+    """Waiting warmups remain dormant until their recorded future is ready."""
+
+    @pytest.mark.asyncio
+    async def test_direct_calls_do_not_reexecute_unresolved_warmup(
+        self, job_def_factory, db0_fixture
+    ):  # pylint: disable=unused-argument
+        """Direct callers cannot bypass the continuation readiness gate."""
+        wait_state = WaitState()
+        future = _controlled_future(wait_state)
+        job = Job(
+            job_def=job_def_factory(
+                warmup_code="attempts = attempts + 1\nresult = future_value"
+            ),
+            job_status=JobStatus.READY,
+        )
+        job.py_env.local_state["attempts"] = 0
+        job.py_env.local_state["future_value"] = future
+
+        assert await run_job_step(job) is False
+        initial_history = list(job.chat_log)
+        initial_console = list(job.py_env.console or [])
+
+        assert await run_job_step(job) is False
+        assert await run_job_step(job) is False
+
+        assert job.status == JobStatus.WARMING_UP
+        assert job.awaited_result is future
+        assert job.next_instr_num == 1
+        assert job.py_env.local_state["attempts"] == 1
+        assert list(job.chat_log) == initial_history
+        assert list(job.py_env.console or []) == initial_console
+        assert wait_state.condition_checks == 2
+
+    @pytest.mark.asyncio
+    async def test_scheduler_skips_blocked_warmup_and_checks_release_once(
+        self, job_def_factory, db0_fixture
+    ):  # pylint: disable=unused-argument
+        """A blocked warmup does not consume a worker or starve another ready job."""
+        wait_state = WaitState()
+        future = _controlled_future(wait_state)
+        blocked_job = Job(
+            job_def=job_def_factory(
+                warmup_code="attempts = attempts + 1\nresult = future_value\nexit(result)"
+            ),
+            job_status=JobStatus.READY,
+        )
+        blocked_job.py_env.local_state["attempts"] = 0
+        blocked_job.py_env.local_state["future_value"] = future
+        assert await run_job_step(blocked_job) is False
+
+        ready_job = Job(
+            job_def=job_def_factory(warmup_code='exit("ready job completed")'),
+            job_status=JobStatus.READY,
+        )
+        initial_history_length = len(_warmup_items(blocked_job))
+        scheduler_cycles = 0
+        repeated_wait_observed = asyncio.Event()
+
+        def record_scheduler_cycle(_available_capacity: int) -> None:
+            nonlocal scheduler_cycles
+            scheduler_cycles += 1
+            if scheduler_cycles >= 3:
+                repeated_wait_observed.set()
+
+        loop_task = asyncio.create_task(
+            run_jobs_loop(
+                max_concurrency=1,
+                queue_prefixes=[db0.get_current_prefix().name],
+                start_jobs_func=record_scheduler_cycle,
+            )
+        )
+        try:
+            await asyncio.wait_for(repeated_wait_observed.wait(), timeout=2)
+
+            assert ready_job.status == JobStatus.DONE
+            assert blocked_job.status == JobStatus.WARMING_UP
+            assert blocked_job.awaited_result is future
+            assert blocked_job.py_env.local_state["attempts"] == 1
+            assert len(_warmup_items(blocked_job)) == initial_history_length
+
+            checks_before_release = wait_state.condition_checks
+            wait_state.ready = True
+            await asyncio.wait_for(_wait_until_done(blocked_job), timeout=2)
+
+            assert blocked_job.py_env.exit_status == "done"
+            assert blocked_job.py_env.local_state["attempts"] == 1
+            assert blocked_job.awaited_result is None
+            assert blocked_job.next_instr_num is None
+            assert wait_state.condition_checks == checks_before_release + 1
+        finally:
+            loop_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await loop_task
+
+    def test_ready_suspended_job_is_released_and_clears_awaited_result(
+        self, job_def_factory, db0_fixture
+    ):  # pylint: disable=unused-argument
+        """Ordinary suspended jobs retain their existing transition to STARTED."""
+        wait_state = WaitState(ready=True)
+        future = _controlled_future(wait_state)
+        job = Job(
+            job_def=job_def_factory(warmup_code=None),
+            job_status=JobStatus.SUSPENDED,
+        )
+        job.awaited_result = future
+        job.next_instr_num = 0
+
+        unsuspend_jobs()
+
+        assert job.status == JobStatus.STARTED
+        assert job.awaited_result is None
+        assert job.next_instr_num == 0
+        assert wait_state.condition_checks == 1