From 4ecf6c828534c2c1d8094b074db59e0128896434 Mon Sep 17 00:00:00 2001 From: Kiki McArron Date: Thu, 24 Sep 2026 10:49:24 +0200 Subject: [PATCH 1/2] [BUG] Pending FutureResult in warmup causes a busy loop and duplicated warmup history --- statek/executors/utils.py | 124 +++++--- tests/test_run_job_step.py | 29 +- tests/test_warmup_blocks.py | 157 ++++++++++- tests/test_warmup_wait_scheduling.py | 406 +++++++++++++++++++++++++++ 4 files changed, 670 insertions(+), 46 deletions(-) create mode 100644 tests/test_warmup_wait_scheduling.py diff --git a/statek/executors/utils.py b/statek/executors/utils.py index 1147cbf..5b1d0a1 100644 --- a/statek/executors/utils.py +++ b/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,12 @@ 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 + + is_continuation = job.next_instr_num is not None + # 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 {} @@ -1048,17 +1068,24 @@ async def run_job_step(job: Job, provider: str = None) -> bool: # Extract the code string (CodeBlock stores code + tool_calls; exec needs the string) code_str = code.code if isinstance(code, CodeBlock) else code - # Create WarmupLogItem before warmup block execution + warmup_log_item = None if job.status == JobStatus.WARMING_UP: block_num = job.warmup_block_num if job.warmup_block_num is not None else 0 - warmup_log_item = WarmupLogItem( - console_pos=len(job.py_env.console) if job.py_env.console else 0, - warmup_block_num=block_num - ) - job.chat_log.append(warmup_log_item) + if is_continuation: + warmup_log_item = next( + item for item in reversed(job.chat_log) + if isinstance(item, WarmupLogItem) + and item.warmup_block_num == block_num + ) + else: + warmup_log_item = WarmupLogItem( + console_pos=len(job.py_env.console) if job.py_env.console else 0, + warmup_block_num=block_num + ) + job.chat_log.append(warmup_log_item) # Log warmup block before execution - if job.status == JobStatus.WARMING_UP and code_str: + if job.status == JobStatus.WARMING_UP and code_str and not is_continuation: chat_style = get_statek_settings().chat_style log_code = f"```python\n{code_str}\n```" if chat_style in (ChatStyle.MARKDOWN, ChatStyle.MD_DIALOG, ChatStyle.DIRECT) else code_str # pylint: disable=no-member job._log(log_code) # pylint: disable=protected-access @@ -1074,22 +1101,26 @@ async def run_job_step(job: Job, provider: str = None) -> bool: job.console_append(error_msg, error=Error(ErrorKind.EXECUTION, error_msg)) # Step 5: Execute regular tool calls (not python_cli) if present and not a continuation - last_chat_log_item = job.chat_log[-1] if job.chat_log else None + execution_log_item = ( + warmup_log_item + if job.status == JobStatus.WARMING_UP + else job.chat_log[-1] if job.chat_log else None + ) - if isinstance(code, CodeBlock) and code.tool_calls and job.next_instr_num is None: + if isinstance(code, CodeBlock) and code.tool_calls and not is_continuation: # Pre-allocate tool_log slots aligned 1-to-1 with the full # tool_calls order so both regular (Step 5) and python_cli # (Step 6 finally) results can be written directly at their # correct index. Default slot value "" keeps get_tool_result(j) # well-defined even if a call is skipped. - if last_chat_log_item is not None: - last_chat_log_item.tool_log = ["" for _ in code.tool_calls] + if execution_log_item is not None: + execution_log_item.tool_log = ["" for _ in code.tool_calls] for j, call_spec in enumerate(code.tool_calls): if call_spec.func_name == "python_cli": continue # CLI calls are executed and logged in Step 6 result, error = await exec_tool(call_spec, job, local_context=dict(local_context)) - if last_chat_log_item is not None: - last_chat_log_item.tool_log[j] = ( + if execution_log_item is not None: + execution_log_item.tool_log[j] = ( ToolError(err_message=error) if error is not None else result ) job._log_tool_call_result(call_spec, result) # pylint: disable=protected-access @@ -1143,14 +1174,37 @@ def _cli_error_append(cli_idx, text): for cli_idx in range(len(cli_calls)): joined = "\n".join(cli_outputs.get(cli_idx, [])) error = "\n".join(cli_errors.get(cli_idx, [])) - if (last_chat_log_item is not None - and last_chat_log_item.tool_log is not None + if (execution_log_item is not None + and execution_log_item.tool_log is not None and cli_idx < len(cli_tool_log_positions)): j = cli_tool_log_positions[cli_idx] - if j < len(last_chat_log_item.tool_log): - last_chat_log_item.tool_log[j] = ( - ToolError(err_message=error) if error else joined - ) + if j < len(execution_log_item.tool_log): + previous = execution_log_item.tool_log[j] + if error: + previous_text = ( + previous.err_message + if isinstance(previous, ToolError) + else previous + ) + combined_error = ( + f"{previous_text}\n{error}" + if is_continuation and previous_text + else error + ) + execution_log_item.tool_log[j] = ToolError( + err_message=combined_error + ) + elif joined: + if is_continuation and isinstance(previous, ToolError): + execution_log_item.tool_log[j] = ToolError( + err_message=f"{previous.err_message}\n{joined}" + ) + elif is_continuation and previous: + execution_log_item.tool_log[j] = f"{previous}\n{joined}" + else: + execution_log_item.tool_log[j] = joined + elif not is_continuation: + execution_log_item.tool_log[j] = "" # Step 9: If code has finished, transition to DONE and return True if job.py_env.exit_status is not None: @@ -1378,23 +1432,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 +1525,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: diff --git a/tests/test_run_job_step.py b/tests/test_run_job_step.py index 7d889e3..8c0961a 100644 --- a/tests/test_run_job_step.py +++ b/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). @@ -494,15 +497,25 @@ def counting_tool(): return "counted" call_spec = CallSpec(id="STATEK-001", func_name="counting_tool", args=[], kwargs={}) - warmup_code = CodeBlock(code='exit("ok")', tool_calls=[call_spec]) + warmup_code = CodeBlock( + code='result = future_val\nexit("ok")', + tool_calls=[call_spec], + ) job = _make_job_with_tool("role_cont", "counting_tool", counting_tool, warmup_code) - # Simulate a FutureError continuation: status stays WARMING_UP, next_instr_num is set - job.set_status(JobStatus.WARMING_UP) - job.next_instr_num = 0 + future = create_future_not_ready() + job.py_env.local_state["future_val"] = future + await run_job_step(job) + assert call_count == 1 + + future.deps.value = 42 + future.set_complement_functions( + complement=_fetch_result_from_deps, + condition=_check_condition_true, + ) await run_job_step(job) - assert call_count == 0 + assert call_count == 1 @pytest.mark.asyncio async def test_no_tool_calls_leaves_warmup_item_tool_log_none(self, db0_fixture): # pylint: disable=unused-argument diff --git a/tests/test_warmup_blocks.py b/tests/test_warmup_blocks.py index cbe5d8b..8667e07 100644 --- a/tests/test_warmup_blocks.py +++ b/tests/test_warmup_blocks.py @@ -6,6 +6,8 @@ import pytest import dbzero as db0 +from statek.chat_history import ChatRole +from statek.executors.chat_log_item import UserLogItem, WarmupLogItem from statek.executors.job import ( Job, JobStatus, @@ -307,10 +309,19 @@ async def test_run_job_step_multiple_blocks_continuation_after_future_error( assert job.next_instr_num == 0 # Suspended at result = future_val instruction # 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 + warmup_items = [ + item for item in job.chat_log if isinstance(item, WarmupLogItem) + ] + assert len(warmup_items) == 2 + waiting_item = warmup_items[-1] + job.chat_log.append(UserLogItem(message="arrived while waiting")) + + # 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) @@ -318,6 +329,17 @@ async def test_run_job_step_multiple_blocks_continuation_after_future_error( assert job.status == JobStatus.WARMING_UP assert "42" in job.py_env.console[0] assert job.warmup_block_num == 2 + warmup_items = [ + item for item in job.chat_log if isinstance(item, WarmupLogItem) + ] + assert len(warmup_items) == 2 + assert warmup_items[-1] is waiting_item + matching_history_items = [ + item for item in job.get_chat_history() + if item.role == ChatRole.ASSISTANT + and item.content == 'result = future_val\nprint(result)' + ] + assert len(matching_history_items) == 1 # Third block: exit result4 = await run_job_step(job) @@ -325,6 +347,133 @@ async def test_run_job_step_multiple_blocks_continuation_after_future_error( assert job.status == JobStatus.DONE assert job.py_env.exit_status == "done" + +class TestWarmupContinuationHistory: + """Warmup continuation output remains attached to one history item.""" + + @pytest.mark.asyncio + async def test_cli_output_is_merged_across_future_wait( + self, job_def_factory, db0_fixture + ): # pylint: disable=unused-argument + """CLI output before and after suspension shares one positional tool result.""" + regular_call_count = 0 + + def setup_tool(): + nonlocal regular_call_count + regular_call_count += 1 + return "setup-result" + + regular_call = CallSpec( + id="REGULAR-001", + func_name="setup_tool", + args=[], + kwargs={}, + ) + cli_call = CallSpec( + id="CLI-001", + func_name="python_cli", + kwargs={ + "code": ( + 'print("before")\n' + "result = future_val\n" + 'print("after", result)' + ) + }, + ) + job = Job( + job_def=job_def_factory( + warmup_code=[ + CodeBlock(code=None, tool_calls=[regular_call, cli_call]), + 'exit("done")', + ] + ), + job_status=JobStatus.READY, + ) + job.job_def.agent.context["setup_tool"] = setup_tool + future = create_future_not_ready() + job.py_env.local_state["future_val"] = future + + assert await run_job_step(job) is False + warmup_items = [ + item for item in job.chat_log if isinstance(item, WarmupLogItem) + ] + assert len(warmup_items) == 1 + warmup_item = warmup_items[0] + assert "setup-result" in warmup_item.tool_log[0] + assert warmup_item.tool_log[1] == "before" + assert regular_call_count == 1 + job.chat_log.append(UserLogItem(message="arrived while waiting")) + + future.deps.value = 42 + future.set_complement_functions( + complement=_fetch_result_from_deps, + condition=_check_condition_true, + ) + assert await run_job_step(job) is False + + warmup_items = [ + item for item in job.chat_log if isinstance(item, WarmupLogItem) + ] + assert warmup_items == [warmup_item] + assert len(warmup_item.tool_log) == 2 + assert "setup-result" in warmup_item.tool_log[0] + assert warmup_item.tool_log[1] == "before\nafter 42" + assert regular_call_count == 1 + history = list(job.get_chat_history()) + assistant_items = [ + item for item in history + if item.role == ChatRole.ASSISTANT and item.tool_calls + ] + tool_items = [ + item for item in history + if item.role == ChatRole.TOOL and item.tool_calls.id == "CLI-001" + ] + assert len(assistant_items) == 1 + assert len(tool_items) == 1 + assert tool_items[0].content == "before\nafter 42" + + @pytest.mark.asyncio + async def test_cli_resume_without_output_preserves_existing_result( + self, job_def_factory, db0_fixture + ): # pylint: disable=unused-argument + """An output-free continuation does not erase earlier CLI output.""" + cli_call = CallSpec( + id="CLI-002", + func_name="python_cli", + kwargs={"code": 'print("before")\nresult = future_val'}, + ) + job = Job( + job_def=job_def_factory( + warmup_code=[CodeBlock(code=None, tool_calls=[cli_call]), 'exit("done")'] + ), + job_status=JobStatus.READY, + ) + future = create_future_not_ready() + job.py_env.local_state["future_val"] = future + + assert await run_job_step(job) is False + warmup_item = next( + item for item in job.chat_log if isinstance(item, WarmupLogItem) + ) + assert warmup_item.tool_log == ["before"] + + future.deps.value = 42 + future.set_complement_functions( + complement=_fetch_result_from_deps, + condition=_check_condition_true, + ) + assert await run_job_step(job) is False + + warmup_items = [ + item for item in job.chat_log if isinstance(item, WarmupLogItem) + ] + assert warmup_items == [warmup_item] + assert warmup_item.tool_log == ["before"] + + +class TestWarmupExecutionModes: + """Warmup blocks retain their execution semantics across chat styles.""" + @pytest.mark.asyncio async def test_direct_warmup_plain_code_executes( self, job_def_factory, db0_fixture # pylint: disable=unused-argument diff --git a/tests/test_warmup_wait_scheduling.py b/tests/test_warmup_wait_scheduling.py new file mode 100644 index 0000000..9b6cf7e --- /dev/null +++ b/tests/test_warmup_wait_scheduling.py @@ -0,0 +1,406 @@ +"""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 +from statek.utils import CallSpec, CodeBlock +from tests.conftest import DB0_DIR + + +@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) + + +def _reopen_job(job: Job) -> Job: + identifier = db0.uuid(job) + db0.close() + db0.init(DB0_DIR, read_write=True) + db0.open("test_prefix", "rw") + return db0.fetch(identifier) + + +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 + + +class TestPersistedWarmupWait: + """Warmup continuation state remains correct across dbzero reopen.""" + + @pytest.mark.asyncio + async def test_unresolved_wait_survives_reopen_and_resumes_exact_instruction( + self, job_def_factory, db0_fixture + ): # pylint: disable=unused-argument + """Restart keeps the recorded future and skips completed instructions.""" + wait_state = WaitState() + future = _controlled_future(wait_state) + job = Job( + job_def=job_def_factory( + warmup_code=( + "attempts = attempts + 1\n" + "result = future_value\n" + "exit(result)" + ) + ), + 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 + restored = _reopen_job(job) + restored_future = restored.awaited_result + + assert restored.status == JobStatus.WARMING_UP + assert restored.next_instr_num == 1 + assert restored_future is restored.py_env.local_state["future_value"] + assert await run_job_step(restored) is False + assert restored.py_env.local_state["attempts"] == 1 + assert len(_warmup_items(restored)) == 1 + + restored_future.deps.ready = True + assert await run_job_step(restored) is True + + assert restored.status == JobStatus.DONE + assert restored.py_env.exit_status == "done" + assert restored.py_env.local_state["attempts"] == 1 + assert restored.awaited_result is None + assert restored.next_instr_num is None + assert len(_warmup_items(restored)) == 1 + + @pytest.mark.asyncio + async def test_released_continuation_survives_reopen_until_execution( + self, job_def_factory, db0_fixture + ): # pylint: disable=unused-argument + """A persisted release latch retains its exact resume instruction.""" + wait_state = WaitState() + future = _controlled_future(wait_state) + job = Job( + job_def=job_def_factory( + warmup_code=( + "attempts = attempts + 1\n" + "result = future_value\n" + "exit(result)" + ) + ), + 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 + wait_state.ready = True + unsuspend_jobs() + assert job.awaited_result is None + assert job.next_instr_num == 1 + + restored = _reopen_job(job) + + assert restored.status == JobStatus.WARMING_UP + assert restored.awaited_result is None + assert restored.next_instr_num == 1 + assert await run_job_step(restored) is True + assert restored.py_env.local_state["attempts"] == 1 + assert restored.next_instr_num is None + + +class TestWarmupWaitCompatibility: + """Continuation release remains compatible with repeated waits and old history.""" + + @pytest.mark.asyncio + async def test_resumed_instruction_can_wait_on_a_second_future( + self, job_def_factory, db0_fixture + ): # pylint: disable=unused-argument + """A new FutureError replaces the released continuation state.""" + first_state = WaitState(value="first") + second_state = WaitState(value="second") + first_future = _controlled_future(first_state) + second_future = _controlled_future(second_state) + job = Job( + job_def=job_def_factory( + warmup_code=( + "attempts = attempts + 1\n" + "first = first_future\n" + "between = between + 1\n" + "second = second_future\n" + "exit(second)" + ) + ), + job_status=JobStatus.READY, + ) + job.py_env.local_state.update({ + "attempts": 0, + "between": 0, + "first_future": first_future, + "second_future": second_future, + }) + + assert await run_job_step(job) is False + assert job.awaited_result is first_future + assert job.next_instr_num == 1 + + first_state.ready = True + assert await run_job_step(job) is False + + assert job.awaited_result is second_future + assert job.next_instr_num == 3 + assert job.py_env.local_state["attempts"] == 1 + assert job.py_env.local_state["between"] == 1 + assert len(_warmup_items(job)) == 1 + + assert await run_job_step(job) is False + assert job.py_env.local_state["between"] == 1 + + second_state.ready = True + assert await run_job_step(job) is True + assert job.py_env.exit_status == "second" + assert job.awaited_result is None + assert job.next_instr_num is None + assert len(_warmup_items(job)) == 1 + + @pytest.mark.asyncio + async def test_existing_duplicate_history_reuses_latest_item_without_cleanup( + self, job_def_factory, db0_fixture + ): # pylint: disable=unused-argument + """Previously duplicated history stays intact and does not grow again.""" + cli_call = CallSpec( + id="CLI-OLD", + func_name="python_cli", + kwargs={ + "code": ( + 'print("before")\n' + "result = future_value\n" + 'print("after", result)' + ) + }, + ) + wait_state = WaitState(value="ready") + future = _controlled_future(wait_state) + job = Job( + job_def=job_def_factory( + warmup_code=[CodeBlock(code=None, tool_calls=[cli_call]), 'exit("done")'] + ), + job_status=JobStatus.READY, + ) + job.py_env.local_state["future_value"] = future + + assert await run_job_step(job) is False + original = _warmup_items(job)[0] + duplicate = WarmupLogItem( + console_pos=original.console_pos, + warmup_block_num=original.warmup_block_num, + tool_log=["before"], + ) + job.chat_log.append(duplicate) + + wait_state.ready = True + assert await run_job_step(job) is False + + assert _warmup_items(job) == [original, duplicate] + assert original.tool_log == ["before"] + assert duplicate.tool_log == ["before\nafter ready"] + + @pytest.mark.asyncio + async def test_auto_terminate_waits_for_unresolved_warmup( + self, job_def_factory, db0_fixture + ): # pylint: disable=unused-argument + """The jobs loop remains active until a waiting warmup can finish.""" + wait_state = WaitState() + future = _controlled_future(wait_state) + job = Job( + job_def=job_def_factory(warmup_code="result = future_value\nexit(result)"), + job_status=JobStatus.READY, + ) + job.py_env.local_state["future_value"] = future + assert await run_job_step(job) is False + + scheduler_cycles = 0 + waiting_observed = asyncio.Event() + + def record_scheduler_cycle(_available_capacity: int) -> None: + nonlocal scheduler_cycles + scheduler_cycles += 1 + if scheduler_cycles >= 2: + waiting_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, + auto_terminate=True, + ) + ) + try: + await asyncio.wait_for(waiting_observed.wait(), timeout=2) + assert loop_task.done() is False + assert job.status == JobStatus.WARMING_UP + assert len(_warmup_items(job)) == 1 + + wait_state.ready = True + await asyncio.wait_for(loop_task, timeout=3) + + assert job.status == JobStatus.DONE + assert job.py_env.exit_status == "done" + assert len(_warmup_items(job)) == 1 + finally: + if not loop_task.done(): + loop_task.cancel() + with suppress(asyncio.CancelledError): + await loop_task From a043417e58788682244ba9595f99c86ba6ca05e4 Mon Sep 17 00:00:00 2001 From: Kiki McArron Date: Thu, 24 Sep 2026 16:52:56 +0200 Subject: [PATCH 2/2] Remove unnecessary job history --- statek/executors/utils.py | 74 +++------ tests/test_run_job_step.py | 20 +-- tests/test_warmup_blocks.py | 146 ----------------- tests/test_warmup_wait_scheduling.py | 236 --------------------------- 4 files changed, 24 insertions(+), 452 deletions(-) diff --git a/statek/executors/utils.py b/statek/executors/utils.py index 5b1d0a1..865ed7f 100644 --- a/statek/executors/utils.py +++ b/statek/executors/utils.py @@ -1046,8 +1046,6 @@ async def run_job_step(job: Job, provider: str = None) -> bool: if not _release_ready_continuation(job): return False - is_continuation = job.next_instr_num is not None - # 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 {} @@ -1068,24 +1066,17 @@ async def run_job_step(job: Job, provider: str = None) -> bool: # Extract the code string (CodeBlock stores code + tool_calls; exec needs the string) code_str = code.code if isinstance(code, CodeBlock) else code - warmup_log_item = None + # Create WarmupLogItem before warmup block execution if job.status == JobStatus.WARMING_UP: block_num = job.warmup_block_num if job.warmup_block_num is not None else 0 - if is_continuation: - warmup_log_item = next( - item for item in reversed(job.chat_log) - if isinstance(item, WarmupLogItem) - and item.warmup_block_num == block_num - ) - else: - warmup_log_item = WarmupLogItem( - console_pos=len(job.py_env.console) if job.py_env.console else 0, - warmup_block_num=block_num - ) - job.chat_log.append(warmup_log_item) + warmup_log_item = WarmupLogItem( + console_pos=len(job.py_env.console) if job.py_env.console else 0, + warmup_block_num=block_num + ) + job.chat_log.append(warmup_log_item) # Log warmup block before execution - if job.status == JobStatus.WARMING_UP and code_str and not is_continuation: + if job.status == JobStatus.WARMING_UP and code_str: chat_style = get_statek_settings().chat_style log_code = f"```python\n{code_str}\n```" if chat_style in (ChatStyle.MARKDOWN, ChatStyle.MD_DIALOG, ChatStyle.DIRECT) else code_str # pylint: disable=no-member job._log(log_code) # pylint: disable=protected-access @@ -1101,26 +1092,22 @@ async def run_job_step(job: Job, provider: str = None) -> bool: job.console_append(error_msg, error=Error(ErrorKind.EXECUTION, error_msg)) # Step 5: Execute regular tool calls (not python_cli) if present and not a continuation - execution_log_item = ( - warmup_log_item - if job.status == JobStatus.WARMING_UP - else job.chat_log[-1] if job.chat_log else None - ) + last_chat_log_item = job.chat_log[-1] if job.chat_log else None - if isinstance(code, CodeBlock) and code.tool_calls and not is_continuation: + if isinstance(code, CodeBlock) and code.tool_calls and job.next_instr_num is None: # Pre-allocate tool_log slots aligned 1-to-1 with the full # tool_calls order so both regular (Step 5) and python_cli # (Step 6 finally) results can be written directly at their # correct index. Default slot value "" keeps get_tool_result(j) # well-defined even if a call is skipped. - if execution_log_item is not None: - execution_log_item.tool_log = ["" for _ in code.tool_calls] + if last_chat_log_item is not None: + last_chat_log_item.tool_log = ["" for _ in code.tool_calls] for j, call_spec in enumerate(code.tool_calls): if call_spec.func_name == "python_cli": continue # CLI calls are executed and logged in Step 6 result, error = await exec_tool(call_spec, job, local_context=dict(local_context)) - if execution_log_item is not None: - execution_log_item.tool_log[j] = ( + if last_chat_log_item is not None: + last_chat_log_item.tool_log[j] = ( ToolError(err_message=error) if error is not None else result ) job._log_tool_call_result(call_spec, result) # pylint: disable=protected-access @@ -1174,37 +1161,14 @@ def _cli_error_append(cli_idx, text): for cli_idx in range(len(cli_calls)): joined = "\n".join(cli_outputs.get(cli_idx, [])) error = "\n".join(cli_errors.get(cli_idx, [])) - if (execution_log_item is not None - and execution_log_item.tool_log is not None + if (last_chat_log_item is not None + and last_chat_log_item.tool_log is not None and cli_idx < len(cli_tool_log_positions)): j = cli_tool_log_positions[cli_idx] - if j < len(execution_log_item.tool_log): - previous = execution_log_item.tool_log[j] - if error: - previous_text = ( - previous.err_message - if isinstance(previous, ToolError) - else previous - ) - combined_error = ( - f"{previous_text}\n{error}" - if is_continuation and previous_text - else error - ) - execution_log_item.tool_log[j] = ToolError( - err_message=combined_error - ) - elif joined: - if is_continuation and isinstance(previous, ToolError): - execution_log_item.tool_log[j] = ToolError( - err_message=f"{previous.err_message}\n{joined}" - ) - elif is_continuation and previous: - execution_log_item.tool_log[j] = f"{previous}\n{joined}" - else: - execution_log_item.tool_log[j] = joined - elif not is_continuation: - execution_log_item.tool_log[j] = "" + if j < len(last_chat_log_item.tool_log): + last_chat_log_item.tool_log[j] = ( + ToolError(err_message=error) if error else joined + ) # Step 9: If code has finished, transition to DONE and return True if job.py_env.exit_status is not None: diff --git a/tests/test_run_job_step.py b/tests/test_run_job_step.py index 8c0961a..248abe9 100644 --- a/tests/test_run_job_step.py +++ b/tests/test_run_job_step.py @@ -497,25 +497,15 @@ def counting_tool(): return "counted" call_spec = CallSpec(id="STATEK-001", func_name="counting_tool", args=[], kwargs={}) - warmup_code = CodeBlock( - code='result = future_val\nexit("ok")', - tool_calls=[call_spec], - ) + warmup_code = CodeBlock(code='exit("ok")', tool_calls=[call_spec]) job = _make_job_with_tool("role_cont", "counting_tool", counting_tool, warmup_code) - future = create_future_not_ready() - job.py_env.local_state["future_val"] = future - - await run_job_step(job) - assert call_count == 1 + # Simulate a FutureError continuation: status stays WARMING_UP, next_instr_num is set + job.set_status(JobStatus.WARMING_UP) + job.next_instr_num = 0 - future.deps.value = 42 - future.set_complement_functions( - complement=_fetch_result_from_deps, - condition=_check_condition_true, - ) await run_job_step(job) - assert call_count == 1 + assert call_count == 0 @pytest.mark.asyncio async def test_no_tool_calls_leaves_warmup_item_tool_log_none(self, db0_fixture): # pylint: disable=unused-argument diff --git a/tests/test_warmup_blocks.py b/tests/test_warmup_blocks.py index 8667e07..9e399d3 100644 --- a/tests/test_warmup_blocks.py +++ b/tests/test_warmup_blocks.py @@ -6,8 +6,6 @@ import pytest import dbzero as db0 -from statek.chat_history import ChatRole -from statek.executors.chat_log_item import UserLogItem, WarmupLogItem from statek.executors.job import ( Job, JobStatus, @@ -309,12 +307,6 @@ async def test_run_job_step_multiple_blocks_continuation_after_future_error( assert job.next_instr_num == 0 # Suspended at result = future_val instruction # warmup_block_num should NOT advance since block didn't complete assert job.warmup_block_num == 1 - warmup_items = [ - item for item in job.chat_log if isinstance(item, WarmupLogItem) - ] - assert len(warmup_items) == 2 - waiting_item = warmup_items[-1] - job.chat_log.append(UserLogItem(message="arrived while waiting")) # Resolve the exact future recorded on the job. future_not_ready.deps.value = 42 @@ -329,17 +321,6 @@ async def test_run_job_step_multiple_blocks_continuation_after_future_error( assert job.status == JobStatus.WARMING_UP assert "42" in job.py_env.console[0] assert job.warmup_block_num == 2 - warmup_items = [ - item for item in job.chat_log if isinstance(item, WarmupLogItem) - ] - assert len(warmup_items) == 2 - assert warmup_items[-1] is waiting_item - matching_history_items = [ - item for item in job.get_chat_history() - if item.role == ChatRole.ASSISTANT - and item.content == 'result = future_val\nprint(result)' - ] - assert len(matching_history_items) == 1 # Third block: exit result4 = await run_job_step(job) @@ -347,133 +328,6 @@ async def test_run_job_step_multiple_blocks_continuation_after_future_error( assert job.status == JobStatus.DONE assert job.py_env.exit_status == "done" - -class TestWarmupContinuationHistory: - """Warmup continuation output remains attached to one history item.""" - - @pytest.mark.asyncio - async def test_cli_output_is_merged_across_future_wait( - self, job_def_factory, db0_fixture - ): # pylint: disable=unused-argument - """CLI output before and after suspension shares one positional tool result.""" - regular_call_count = 0 - - def setup_tool(): - nonlocal regular_call_count - regular_call_count += 1 - return "setup-result" - - regular_call = CallSpec( - id="REGULAR-001", - func_name="setup_tool", - args=[], - kwargs={}, - ) - cli_call = CallSpec( - id="CLI-001", - func_name="python_cli", - kwargs={ - "code": ( - 'print("before")\n' - "result = future_val\n" - 'print("after", result)' - ) - }, - ) - job = Job( - job_def=job_def_factory( - warmup_code=[ - CodeBlock(code=None, tool_calls=[regular_call, cli_call]), - 'exit("done")', - ] - ), - job_status=JobStatus.READY, - ) - job.job_def.agent.context["setup_tool"] = setup_tool - future = create_future_not_ready() - job.py_env.local_state["future_val"] = future - - assert await run_job_step(job) is False - warmup_items = [ - item for item in job.chat_log if isinstance(item, WarmupLogItem) - ] - assert len(warmup_items) == 1 - warmup_item = warmup_items[0] - assert "setup-result" in warmup_item.tool_log[0] - assert warmup_item.tool_log[1] == "before" - assert regular_call_count == 1 - job.chat_log.append(UserLogItem(message="arrived while waiting")) - - future.deps.value = 42 - future.set_complement_functions( - complement=_fetch_result_from_deps, - condition=_check_condition_true, - ) - assert await run_job_step(job) is False - - warmup_items = [ - item for item in job.chat_log if isinstance(item, WarmupLogItem) - ] - assert warmup_items == [warmup_item] - assert len(warmup_item.tool_log) == 2 - assert "setup-result" in warmup_item.tool_log[0] - assert warmup_item.tool_log[1] == "before\nafter 42" - assert regular_call_count == 1 - history = list(job.get_chat_history()) - assistant_items = [ - item for item in history - if item.role == ChatRole.ASSISTANT and item.tool_calls - ] - tool_items = [ - item for item in history - if item.role == ChatRole.TOOL and item.tool_calls.id == "CLI-001" - ] - assert len(assistant_items) == 1 - assert len(tool_items) == 1 - assert tool_items[0].content == "before\nafter 42" - - @pytest.mark.asyncio - async def test_cli_resume_without_output_preserves_existing_result( - self, job_def_factory, db0_fixture - ): # pylint: disable=unused-argument - """An output-free continuation does not erase earlier CLI output.""" - cli_call = CallSpec( - id="CLI-002", - func_name="python_cli", - kwargs={"code": 'print("before")\nresult = future_val'}, - ) - job = Job( - job_def=job_def_factory( - warmup_code=[CodeBlock(code=None, tool_calls=[cli_call]), 'exit("done")'] - ), - job_status=JobStatus.READY, - ) - future = create_future_not_ready() - job.py_env.local_state["future_val"] = future - - assert await run_job_step(job) is False - warmup_item = next( - item for item in job.chat_log if isinstance(item, WarmupLogItem) - ) - assert warmup_item.tool_log == ["before"] - - future.deps.value = 42 - future.set_complement_functions( - complement=_fetch_result_from_deps, - condition=_check_condition_true, - ) - assert await run_job_step(job) is False - - warmup_items = [ - item for item in job.chat_log if isinstance(item, WarmupLogItem) - ] - assert warmup_items == [warmup_item] - assert warmup_item.tool_log == ["before"] - - -class TestWarmupExecutionModes: - """Warmup blocks retain their execution semantics across chat styles.""" - @pytest.mark.asyncio async def test_direct_warmup_plain_code_executes( self, job_def_factory, db0_fixture # pylint: disable=unused-argument diff --git a/tests/test_warmup_wait_scheduling.py b/tests/test_warmup_wait_scheduling.py index 9b6cf7e..0f49ac4 100644 --- a/tests/test_warmup_wait_scheduling.py +++ b/tests/test_warmup_wait_scheduling.py @@ -13,8 +13,6 @@ 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 -from statek.utils import CallSpec, CodeBlock -from tests.conftest import DB0_DIR @db0.memo @@ -56,14 +54,6 @@ async def _wait_until_done(job: Job) -> None: await asyncio.sleep(0.01) -def _reopen_job(job: Job) -> Job: - identifier = db0.uuid(job) - db0.close() - db0.init(DB0_DIR, read_write=True) - db0.open("test_prefix", "rw") - return db0.fetch(identifier) - - class TestWarmupWaitReadiness: """Waiting warmups remain dormant until their recorded future is ready.""" @@ -178,229 +168,3 @@ def test_ready_suspended_job_is_released_and_clears_awaited_result( assert job.awaited_result is None assert job.next_instr_num == 0 assert wait_state.condition_checks == 1 - - -class TestPersistedWarmupWait: - """Warmup continuation state remains correct across dbzero reopen.""" - - @pytest.mark.asyncio - async def test_unresolved_wait_survives_reopen_and_resumes_exact_instruction( - self, job_def_factory, db0_fixture - ): # pylint: disable=unused-argument - """Restart keeps the recorded future and skips completed instructions.""" - wait_state = WaitState() - future = _controlled_future(wait_state) - job = Job( - job_def=job_def_factory( - warmup_code=( - "attempts = attempts + 1\n" - "result = future_value\n" - "exit(result)" - ) - ), - 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 - restored = _reopen_job(job) - restored_future = restored.awaited_result - - assert restored.status == JobStatus.WARMING_UP - assert restored.next_instr_num == 1 - assert restored_future is restored.py_env.local_state["future_value"] - assert await run_job_step(restored) is False - assert restored.py_env.local_state["attempts"] == 1 - assert len(_warmup_items(restored)) == 1 - - restored_future.deps.ready = True - assert await run_job_step(restored) is True - - assert restored.status == JobStatus.DONE - assert restored.py_env.exit_status == "done" - assert restored.py_env.local_state["attempts"] == 1 - assert restored.awaited_result is None - assert restored.next_instr_num is None - assert len(_warmup_items(restored)) == 1 - - @pytest.mark.asyncio - async def test_released_continuation_survives_reopen_until_execution( - self, job_def_factory, db0_fixture - ): # pylint: disable=unused-argument - """A persisted release latch retains its exact resume instruction.""" - wait_state = WaitState() - future = _controlled_future(wait_state) - job = Job( - job_def=job_def_factory( - warmup_code=( - "attempts = attempts + 1\n" - "result = future_value\n" - "exit(result)" - ) - ), - 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 - wait_state.ready = True - unsuspend_jobs() - assert job.awaited_result is None - assert job.next_instr_num == 1 - - restored = _reopen_job(job) - - assert restored.status == JobStatus.WARMING_UP - assert restored.awaited_result is None - assert restored.next_instr_num == 1 - assert await run_job_step(restored) is True - assert restored.py_env.local_state["attempts"] == 1 - assert restored.next_instr_num is None - - -class TestWarmupWaitCompatibility: - """Continuation release remains compatible with repeated waits and old history.""" - - @pytest.mark.asyncio - async def test_resumed_instruction_can_wait_on_a_second_future( - self, job_def_factory, db0_fixture - ): # pylint: disable=unused-argument - """A new FutureError replaces the released continuation state.""" - first_state = WaitState(value="first") - second_state = WaitState(value="second") - first_future = _controlled_future(first_state) - second_future = _controlled_future(second_state) - job = Job( - job_def=job_def_factory( - warmup_code=( - "attempts = attempts + 1\n" - "first = first_future\n" - "between = between + 1\n" - "second = second_future\n" - "exit(second)" - ) - ), - job_status=JobStatus.READY, - ) - job.py_env.local_state.update({ - "attempts": 0, - "between": 0, - "first_future": first_future, - "second_future": second_future, - }) - - assert await run_job_step(job) is False - assert job.awaited_result is first_future - assert job.next_instr_num == 1 - - first_state.ready = True - assert await run_job_step(job) is False - - assert job.awaited_result is second_future - assert job.next_instr_num == 3 - assert job.py_env.local_state["attempts"] == 1 - assert job.py_env.local_state["between"] == 1 - assert len(_warmup_items(job)) == 1 - - assert await run_job_step(job) is False - assert job.py_env.local_state["between"] == 1 - - second_state.ready = True - assert await run_job_step(job) is True - assert job.py_env.exit_status == "second" - assert job.awaited_result is None - assert job.next_instr_num is None - assert len(_warmup_items(job)) == 1 - - @pytest.mark.asyncio - async def test_existing_duplicate_history_reuses_latest_item_without_cleanup( - self, job_def_factory, db0_fixture - ): # pylint: disable=unused-argument - """Previously duplicated history stays intact and does not grow again.""" - cli_call = CallSpec( - id="CLI-OLD", - func_name="python_cli", - kwargs={ - "code": ( - 'print("before")\n' - "result = future_value\n" - 'print("after", result)' - ) - }, - ) - wait_state = WaitState(value="ready") - future = _controlled_future(wait_state) - job = Job( - job_def=job_def_factory( - warmup_code=[CodeBlock(code=None, tool_calls=[cli_call]), 'exit("done")'] - ), - job_status=JobStatus.READY, - ) - job.py_env.local_state["future_value"] = future - - assert await run_job_step(job) is False - original = _warmup_items(job)[0] - duplicate = WarmupLogItem( - console_pos=original.console_pos, - warmup_block_num=original.warmup_block_num, - tool_log=["before"], - ) - job.chat_log.append(duplicate) - - wait_state.ready = True - assert await run_job_step(job) is False - - assert _warmup_items(job) == [original, duplicate] - assert original.tool_log == ["before"] - assert duplicate.tool_log == ["before\nafter ready"] - - @pytest.mark.asyncio - async def test_auto_terminate_waits_for_unresolved_warmup( - self, job_def_factory, db0_fixture - ): # pylint: disable=unused-argument - """The jobs loop remains active until a waiting warmup can finish.""" - wait_state = WaitState() - future = _controlled_future(wait_state) - job = Job( - job_def=job_def_factory(warmup_code="result = future_value\nexit(result)"), - job_status=JobStatus.READY, - ) - job.py_env.local_state["future_value"] = future - assert await run_job_step(job) is False - - scheduler_cycles = 0 - waiting_observed = asyncio.Event() - - def record_scheduler_cycle(_available_capacity: int) -> None: - nonlocal scheduler_cycles - scheduler_cycles += 1 - if scheduler_cycles >= 2: - waiting_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, - auto_terminate=True, - ) - ) - try: - await asyncio.wait_for(waiting_observed.wait(), timeout=2) - assert loop_task.done() is False - assert job.status == JobStatus.WARMING_UP - assert len(_warmup_items(job)) == 1 - - wait_state.ready = True - await asyncio.wait_for(loop_task, timeout=3) - - assert job.status == JobStatus.DONE - assert job.py_env.exit_status == "done" - assert len(_warmup_items(job)) == 1 - finally: - if not loop_task.done(): - loop_task.cancel() - with suppress(asyncio.CancelledError): - await loop_task