diff --git a/statek/executors/chat_log_item.py b/statek/executors/chat_log_item.py index b313b16..7d16b5d 100644 --- a/statek/executors/chat_log_item.py +++ b/statek/executors/chat_log_item.py @@ -93,6 +93,18 @@ class LLM_LogItem(ChatLogItem): llm_reasoning_payload: Optional[Any] = None # Non-default difficulty used to build this request; None means the job's static default. request_difficulty: Optional[TaskDifficulty] = None + # Cumulative approximate usage including this response; None for placeholders. + cumulative_token_usage: Optional[int] = None + + +@db0.memo(no_default_tags=True) +@dataclass(kw_only=True) +class DifficultyChangeLogItem(ChatLogItem): + """An upward effective difficulty transition, excluded from conversation history.""" + + previous_difficulty: TaskDifficulty + difficulty: TaskDifficulty + cumulative_token_usage: int @db0.memo(no_default_tags=True) diff --git a/statek/executors/job.py b/statek/executors/job.py index fc1da78..e912e13 100644 --- a/statek/executors/job.py +++ b/statek/executors/job.py @@ -40,6 +40,7 @@ ) from statek.executors.llm_usage import LLM_Usage from statek.executors.chat_log_item import ( + DifficultyChangeLogItem, ChatLogItem, LLM_LogItem, NotificationLogItem, @@ -717,7 +718,7 @@ def _chat_log_finalized_for_notifications(self) -> bool: """Return whether pending notifications can be safely appended now.""" if not self.chat_log: return True - last = self.chat_log[-1] + last = self.last_chat_log_item if not isinstance(last, LLM_LogItem): return True if last.llm_resp is None: @@ -1097,6 +1098,16 @@ def _notification_tool_call(idx: int) -> CallSpecWrapper: kwargs={"code": "print(get_pending_notifications())"}, ) + def iter_notification_log_items(self) -> Iterable[Tuple[str, NotificationLogItem]]: + """Yield logged notifications paired with their synthetic history call IDs. + + IDs are scoped to this job's append-only chat log. Pending notifications + are excluded; reading this iterator does not deliver or consume them. + """ + for index, item in enumerate(self.chat_log): + if isinstance(item, NotificationLogItem): + yield self._notification_tool_call(index).id, item + def get_next_prompt(self) -> str: """ Generate the next prompt to be included in the LLM chat. @@ -1748,13 +1759,13 @@ def get_current_difficulty(self) -> TaskDifficulty: # update the dynamically resolved difficulty if difficulty is not None: if last_difficulty is None: - self.__last_difficulty = difficulty + self._set_dynamic_difficulty(difficulty) last_difficulty = difficulty else: max_difficulty = max_task_difficulty(last_difficulty, difficulty) # only update if changed if last_difficulty != max_difficulty: - self.__last_difficulty = max_difficulty + self._set_dynamic_difficulty(max_difficulty) last_difficulty = max_difficulty if last_difficulty is not None: @@ -1762,16 +1773,46 @@ def get_current_difficulty(self) -> TaskDifficulty: return _get_static_task_difficulty(self.job_def.metadata) - def panic(self): - """Escalate the current job difficulty to the next higher level.""" + def _set_dynamic_difficulty(self, difficulty: TaskDifficulty) -> None: + """Record each effective increase once, with usage at the transition.""" + previous = (getattr(self, "_Job__last_difficulty", None) + or _get_static_task_difficulty(self.job_def.metadata)) + if (previous != difficulty + and max_task_difficulty(previous, difficulty) == difficulty): + # Model pricing resolves difficulty before usage exists during initialization. + usage = self.approx_token_usage if hasattr(self, "usage") else 0 + self.chat_log.append(DifficultyChangeLogItem( + console_pos=len(self.py_env.console) if self.py_env.console else 0, + previous_difficulty=previous, + difficulty=difficulty, + cumulative_token_usage=usage, + )) + self.__last_difficulty = difficulty + + @property + def difficulty_token_extension(self) -> int: + """Derive the additional token budget from immutable escalation snapshots.""" + return sum(item.cumulative_token_usage for item in self.chat_log + if isinstance(item, DifficultyChangeLogItem)) + + @property + def last_chat_log_item(self) -> Optional[Union[str, ChatLogItem, UserLogItem]]: + """Return the current conversation/execution item, ignoring budget metadata.""" + return next((item for item in reversed(self.chat_log) + if not isinstance(item, DifficultyChangeLogItem)), None) + + def panic(self) -> None: + """Escalate difficulty and extend this job's token budget by usage so far.""" current_difficulty = self.get_current_difficulty() - self.__last_difficulty = _next_task_difficulty(current_difficulty) + next_difficulty = _next_task_difficulty(current_difficulty) + self._set_dynamic_difficulty(next_difficulty) def append_chat_log( self, request: Dict, llm_resp: LLM_Response, request_difficulty: Optional[TaskDifficulty] = None, + cumulative_token_usage: Optional[int] = None, ): """ Register the LLM response in the Job's chat_log container. @@ -1786,6 +1827,8 @@ def append_chat_log( llm_resp: The LLM's response as an LLM_Response object request_difficulty: Concrete difficulty captured when the request was built, or None for the static default. + cumulative_token_usage: Usage captured after response accounting, + before post-processing. Defaults to current usage for direct callers. The console_pos is set to len(console), marking the position past the end of the current console output. @@ -1835,6 +1878,8 @@ def append_chat_log( llm_resp=stored_resp, llm_reasoning_payload=step_data.reasoning_payload, request_difficulty=request_difficulty, + cumulative_token_usage=(self.approx_token_usage if cumulative_token_usage is None + else cumulative_token_usage), ) self.chat_log.append(chat_item) @@ -1927,7 +1972,7 @@ def last_response(self) -> Union[str, CodeBlock, None]: """ if not self.chat_log: return None - last = self.chat_log[-1] + last = self.last_chat_log_item if isinstance(last, LLM_LogItem): return last.llm_resp return None @@ -2278,9 +2323,9 @@ def _warmup_end_positions(self) -> List[int]: for i, item in enumerate(self.chat_log): if isinstance(item, WarmupLogItem): end_pos = ( - self.chat_log[i + 1].console_pos - if i + 1 < len(self.chat_log) - else console_len + next((following.console_pos for following in self.chat_log[i + 1:] + if isinstance(following, ChatLogItem) + and not isinstance(following, DifficultyChangeLogItem)), console_len) ) pairs.append((item.warmup_block_num, end_pos)) pairs.sort(key=lambda x: x[0]) diff --git a/statek/executors/utils.py b/statek/executors/utils.py index 309f53f..b7cdca1 100644 --- a/statek/executors/utils.py +++ b/statek/executors/utils.py @@ -907,32 +907,9 @@ def _private_print(*args, sep=' ', end='\n', **kwargs): def _log_pending_console(job: Job): - """Log console output accumulated since the last logged boundary. - - Determines from_pos based on job state: - - After LLM turns: from last chat_log item's console_pos - - After warmup: from the next element's console_pos (or current console length) - - Initial state: from 0 - """ - if job.chat_log: - last = job.chat_log[-1] - if isinstance(last, WarmupLogItem): - # If a following element exists its console_pos marks where the - # warmup batch ended (it was already logged by _log_console_batch). - # Otherwise the batch has not been logged yet (e.g. early exit - # during warmup), so start from the WarmupLogItem's own position. - last_idx = len(job.chat_log) - 1 - from_pos = ( - job.chat_log[last_idx + 1].console_pos - if last_idx + 1 < len(job.chat_log) - else last.console_pos - ) - elif isinstance(last, ChatLogItem): - from_pos = last.console_pos - else: - from_pos = 0 - else: - from_pos = 0 + """Log console output since the last conversation boundary, ignoring budget metadata.""" + last = job.last_chat_log_item + from_pos = last.console_pos if isinstance(last, ChatLogItem) else 0 to_pos = len(job.py_env.console) if job.py_env.console else 0 job._log_console_batch(from_pos, to_pos) # pylint: disable=protected-access @@ -1092,7 +1069,7 @@ 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 + last_chat_log_item = job.last_chat_log_item 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 @@ -1232,6 +1209,7 @@ def _cli_error_append(cli_idx, text): job.usage.total_cached_tokens += response.stats.cached_tokens if response.stats.cost is not None: job.usage.total_reported_cost = (job.usage.total_reported_cost or 0.0) + response.stats.cost + cumulative_token_usage = job.approx_token_usage # Step 14: Apply post-processors before committing the response to chat_log. processed_step = response.step_data @@ -1251,6 +1229,7 @@ def _cli_error_append(cli_idx, text): request, processed_response, request_difficulty=request_difficulty, + cumulative_token_usage=cumulative_token_usage, ) # Step 15: MD_DIALOG/DIRECT — dispatch LLM response text to user via send_message diff --git a/statek/llm_harness.py b/statek/llm_harness.py index 22dbd19..fcbf828 100644 --- a/statek/llm_harness.py +++ b/statek/llm_harness.py @@ -101,6 +101,8 @@ def check_after_step(self, job): """ nc = getattr(job, 'num_completions', None) eff_tokens = self._effective_limit(self.max_token_usage, nc) + if eff_tokens is not None: + eff_tokens += job.difficulty_token_extension if eff_tokens is not None and job.approx_token_usage > eff_tokens: raise LLM_HarnessError( diff --git a/statek/system.py b/statek/system.py index cccf8ba..bd5f95c 100644 --- a/statek/system.py +++ b/statek/system.py @@ -795,6 +795,8 @@ def panic(**kwargs): # pylint: disable=unused-argument Use this when the current task is too difficult for the current model. This raises the job difficulty by one level and prints the new difficulty. + A successful escalation also extends this job's finite token limit by its + approximate token usage so far. Unlimited token limits remain unlimited. Returns: None. Updates the current job and prints the current difficulty. diff --git a/tests/test_job.py b/tests/test_job.py index 59a4a3e..2eda426 100644 --- a/tests/test_job.py +++ b/tests/test_job.py @@ -21,7 +21,7 @@ ) from statek.llm_api import LLM_Response, LLM_StepData, LLM_Stats, OpenRouter_API from statek.model_pricing import set_model_pricing -from statek.pyenv import Error, ErrorKind +from statek.pyenv import Error, ErrorKind, PyEnv from statek.model_name import ( ModelName, ensure_model_name, @@ -32,6 +32,7 @@ from statek.chat_history import ChatRole, ContentSource, format_chat_history_item from statek.agents.dialog_agent import RecurringReminder from statek.executors.chat_log_item import ( + DifficultyChangeLogItem, LLM_LogItem, NotificationLogItem, PostProcessedItem, @@ -43,6 +44,8 @@ from statek.executors.post_processor import PostProcessor from statek.executors.utils import _log_pending_console from statek.future import FutureResult +from statek.exceptions import LLM_HarnessError +from statek.llm_harness import LLM_Harness from statek.settings import ChatStyle, LLM_API_Settings from statek.locale import StatekLocale, StatekLangCode, StatekCountryCode from statek.prompt_config import make_system_prompt, parse_system_prompt @@ -1026,6 +1029,187 @@ def test_get_current_difficulty_ignores_non_job_task_difficulty_attribute( assert _run_with_current_job(job, job.get_current_difficulty) == TaskDifficulty.low +def test_panic_budget_accumulates_and_survives_reopen(job_def_factory) -> None: + """Each successful panic adds total usage without affecting another job's budget.""" + definition = job_def_factory(metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "low"}) + job = Job(job_def=definition, job_status=JobStatus.READY) + sibling = Job(job_def=definition, job_status=JobStatus.READY) + job.usage.total_bytes_sent = 30000 + job.usage.total_bytes_received = 10000 + harness = LLM_Harness(None, 100, 100, 50000) + + job.panic() + assert job.difficulty_token_extension == 10000 + assert job.get_current_difficulty() == TaskDifficulty.medium + assert job.approx_token_usage == 10000 + job.usage.total_bytes_received = 70000 + job.panic() + assert job.difficulty_token_extension == 35000 + assert job.get_current_difficulty() == TaskDifficulty.high + with pytest.raises(RuntimeError, match="already at high difficulty"): + job.panic() + assert job.difficulty_token_extension == 35000 + events = [item for item in job.chat_log if isinstance(item, DifficultyChangeLogItem)] + assert [item.cumulative_token_usage for item in events] == [10000, 25000] + assert [(item.previous_difficulty, item.difficulty) for item in events] == [ + (TaskDifficulty.low, TaskDifficulty.medium), + (TaskDifficulty.medium, TaskDifficulty.high), + ] + + job.usage.total_bytes_received = 310000 + harness.check_after_step(job) + sibling.usage.total_bytes_sent = 240000 + with pytest.raises(LLM_HarnessError, match="60000/50000"): + harness.check_after_step(sibling) + identifier = db0.uuid(job) + db0.close() + db0.init(DB0_DIR, read_write=True) + db0.open("test_prefix", "rw") + restored = db0.fetch(identifier) + assert restored.difficulty_token_extension == 35000 + assert restored.approx_token_usage == 85000 + harness.check_after_step(restored) + restored.usage.total_bytes_received += 4 + with pytest.raises(LLM_HarnessError, match="85001/85000"): + harness.check_after_step(restored) + + +@pytest.mark.parametrize("starting_difficulty", [None, TaskDifficulty.low]) +def test_example_escalation_extends_budget_once(job_def_factory, starting_difficulty) -> None: + """Effective static and dynamic baselines both grant a single historical bonus.""" + job = Job(job_def=job_def_factory( + metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "low"}, + )) + job._Job__last_difficulty = starting_difficulty # pylint: disable=protected-access + job.usage.total_bytes_sent = 40000 + job.py_env.local_state["_PERM_CTX"] = {"last_example_id": 1} + with patch("statek.executors.job._get_example_difficulty_for_job", + return_value=TaskDifficulty.medium): + assert job.get_current_difficulty() == TaskDifficulty.medium + assert job.difficulty_token_extension == 10000 + job.usage.total_bytes_sent = 100000 + assert job.get_current_difficulty() == TaskDifficulty.medium + assert job.difficulty_token_extension == 10000 + job.panic() + assert job.get_current_difficulty() == TaskDifficulty.high + assert job.difficulty_token_extension == 35000 + assert not hasattr(job, "panic_token_extension") + + +def test_direct_example_jump_is_one_bonus(job_def_factory) -> None: + """A single low-to-high change is distinct from two separate escalations.""" + job = Job(job_def=job_def_factory( + metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "low"}, + )) + job.usage.total_bytes_sent = 40000 + job.py_env.local_state["_PERM_CTX"] = {"last_example_id": 1} + with patch("statek.executors.job._get_example_difficulty_for_job", + return_value=TaskDifficulty.high): + assert job.get_current_difficulty() == TaskDifficulty.high + assert job.difficulty_token_extension == 10000 + + +@pytest.mark.parametrize("example_difficulty", [TaskDifficulty.low, TaskDifficulty.medium]) +def test_non_increasing_example_grants_no_bonus(job_def_factory, example_difficulty) -> None: + """Initial assignment at or below the static baseline is not an increase.""" + job = Job(job_def=job_def_factory( + metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "medium"}, + )) + job.usage.total_bytes_sent = 40000 + job.py_env.local_state["_PERM_CTX"] = {"last_example_id": 1} + with patch("statek.executors.job._get_example_difficulty_for_job", + return_value=example_difficulty): + assert job.get_current_difficulty() == example_difficulty + assert job.difficulty_token_extension == 0 + + +def test_panic_with_zero_usage_grants_no_bonus(job_def_factory) -> None: + """Escalation is still possible without previous token usage.""" + job = Job(job_def=job_def_factory( + metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "low"}, + )) + job.panic() + assert job.difficulty_token_extension == 0 + assert job.approx_token_usage == 0 + + +def test_example_difficulty_during_initialization_has_zero_usage(job_def_factory) -> None: + """Pricing may resolve a default example before the usage object exists.""" + env = PyEnv() + env.local_state["default_example_id"] = 1 + with patch("statek.executors.job._get_example_difficulty_for_job", + return_value=TaskDifficulty.medium): + job = Job(job_def=job_def_factory(metadata={ + "MODEL": "L:small,M:medium,H:large", "DEFAULT_DIFFICULTY": "low", + }), py_env=env) + assert job.get_current_difficulty() == TaskDifficulty.medium + assert job.difficulty_token_extension == 0 + assert len(job.chat_log) == 1 + assert job.chat_log[0].cumulative_token_usage == 0 + + +def test_difficulty_metadata_preserves_console_and_notification_history(job_def_factory) -> None: + """Budget events must not truncate console output or steal notification IDs.""" + job = Job(job_def=job_def_factory(metadata={ + "MODEL": "test-model", "DEFAULT_DIFFICULTY": "low", + }), job_status=JobStatus.STARTED) + response = LLM_LogItem(console_pos=0, llm_resp="pass") + job.chat_log.append(response) + job.py_env.console = ["before"] + job.panic() + assert job.last_response == "pass" + assert job.get_next_code_block() == "pass" + assert job.num_turns == 1 + assert job.exception_count == 0 + job.py_env.console.append("after") + history = list(job.get_chat_history()) + assert any(item.content == "before\nafter" for item in history) + assert len(history) == 2 + job.send_notification("notice") + assert job._process_pending_notifications() is True # pylint: disable=protected-access + identifier, notification = next(job.iter_notification_log_items()) + assert notification.message == "notice" + assert any(item.tool_calls and item.tool_calls[0].id == identifier + for item in job.get_chat_history() if item.role == ChatRole.ASSISTANT) + + +def test_difficulty_metadata_does_not_finalize_pending_tools(job_def_factory) -> None: + """Escalation metadata cannot bypass notification delivery's active-call guard.""" + job = Job(job_def=job_def_factory(metadata={ + "MODEL": "test-model", "DEFAULT_DIFFICULTY": "low", + }), job_status=JobStatus.STARTED) + job.chat_log.append(LLM_LogItem(console_pos=0, llm_resp=CodeBlock( + code=None, tool_calls=[CallSpec(id="pending", func_name="panic")], + ))) + job.panic() + job.send_notification("notice") + assert job._process_pending_notifications() is False # pylint: disable=protected-access + assert len(job._pending_notifications()) == 1 # pylint: disable=protected-access + + +def test_resume_placeholder_has_no_usage_snapshot(job_def_factory) -> None: + """A resumed conversation's placeholder is not a provider response.""" + job = Job(job_def=job_def_factory(), job_status=JobStatus.DONE) + assert job.push_user_message("continue") is True + assert job.chat_log[-1].llm_resp is None + assert job.chat_log[-1].cumulative_token_usage is None + + +def test_budget_metadata_preserves_warmup_console_boundary(job_def_factory) -> None: + """A metadata event inside warmup is not the warmup output's end boundary.""" + job = Job(job_def=job_def_factory( + warmup_code="pass", + metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "low"}, + )) + job.chat_log.append(WarmupLogItem(console_pos=0, warmup_block_num=0)) + job.py_env.console = ["before"] + job.panic() + job.py_env.console.append("after") + assert job._warmup_end_positions() == [2] # pylint: disable=protected-access + history = list(job.get_chat_history()) + assert any(item.content == "before\nafter" for item in history) + + def test_panic_increases_low_difficulty_to_medium(job_def_factory): """panic raises the current difficulty by one level.""" job_def = job_def_factory( @@ -1595,7 +1779,8 @@ def test_get_request_data_keeps_historical_model_and_prompt_difficulty( assert "Medium instructions." in historical_medium["system_prompt"] assert "High instructions." not in historical_medium["system_prompt"] assert job.chat_log[0].request_difficulty is None - assert job.chat_log[1].request_difficulty == medium_difficulty + llm_items = [item for item in job.chat_log if isinstance(item, LLM_LogItem)] + assert llm_items[1].request_difficulty == medium_difficulty def test_get_request_data_supports_legacy_log_item_without_difficulty_snapshot( self, job_def_factory, @@ -3062,6 +3247,40 @@ def test_push_notification_active_buffers_until_last_item_finalized(self, job_fa class TestTextNotificationHistory: """A plain-text notification is a durable, synthetic tool-result event.""" + def test_notification_log_items_match_history_ids_without_draining_queue( + self, job_factory, + ) -> None: + """Consumers get Statek-owned call IDs without changing notification delivery.""" + job = job_factory() + first = NotificationLogItem(console_pos=0, message="first", tool_log="first") + second = NotificationLogItem(console_pos=0, message="second", tool_log="second") + pending = NotificationLogItem(console_pos=0, message="pending") + job.chat_log = [ + LLM_LogItem(console_pos=0, llm_resp="before"), + first, + LLM_LogItem(console_pos=0, llm_resp="between"), + second, + ] + job._pending_notifications().append(pending) + + entries = list(job.iter_notification_log_items()) + history_ids = [ + entry.tool_calls[0].id for entry in job.get_chat_history() + if entry.role == ChatRole.ASSISTANT and entry.tool_calls + ] + + assert [identifier for identifier, _ in entries] == history_ids + assert len(set(history_ids)) == 2 + assert [item for _, item in entries] == [first, second] + assert list(job.iter_notification_log_items()) == entries + assert job._pending_notifications() == [pending] + assert pending.tool_log is None + assert len(job.chat_log) == 4 + + def test_notification_log_items_empty(self, job_factory) -> None: + """A job with no delivered notifications exposes no notification IDs.""" + assert not list(job_factory().iter_notification_log_items()) + def test_notification_is_persisted_and_formatted_without_recursive_tool_log(self, job_factory): """The message is required and the saved result does not expand itself.""" job = job_factory() diff --git a/tests/test_llm_harness.py b/tests/test_llm_harness.py index 99d8e48..1aefdb0 100644 --- a/tests/test_llm_harness.py +++ b/tests/test_llm_harness.py @@ -9,13 +9,14 @@ def _make_job(num_turns=0, exception_count=0, max_consecutive_exceptions=0, - approx_token_usage=0, num_completions=None): + approx_token_usage=0, num_completions=None, difficulty_token_extension=0): job = MagicMock() job.num_turns = num_turns job.exception_count = exception_count job.max_consecutive_exceptions = max_consecutive_exceptions job.approx_token_usage = approx_token_usage job.num_completions = num_completions + job.difficulty_token_extension = difficulty_token_extension return job @@ -60,6 +61,27 @@ def test_no_limits_does_not_raise(self): class TestCheckAfterStep: """check_after_step should raise only when token usage exceeds the limit.""" + @pytest.mark.parametrize("completions, limit", [(None, 60000), (1, 75000), (2, 90000)]) + def test_panic_bonus_is_added_after_completion_extension(self, completions, limit) -> None: + """The panic bonus is additive, job-local, and enforced at the new boundary.""" + harness = LLM_Harness(None, 3, 1, 50000, limit_extension_per_completion=0.3) + job = _make_job(approx_token_usage=limit, num_completions=completions, + difficulty_token_extension=10000) + harness.check_after_step(job) + job.approx_token_usage += 1 + with pytest.raises(LLM_HarnessError, match=f"{limit + 1}/{limit}"): + harness.check_after_step(job) + assert harness.max_token_usage == 50000 + + def test_panic_bonus_does_not_extend_other_limits(self) -> None: + """Extra tokens cannot bypass exception or turn safeguards.""" + harness = LLM_Harness(1, 1, 1, 1000) + job = _make_job(num_turns=2, exception_count=2, difficulty_token_extension=10000) + with pytest.raises(LLM_HarnessError, match="turns"): + harness.check_before_step(job) + with pytest.raises(LLM_HarnessError, match="exceptions"): + harness.check_after_step(job) + def test_token_usage_allows_equal(self): """Exactly at the limit should be fine; over should raise.""" harness = LLM_Harness(max_turns=None, max_exceptions=100, @@ -72,7 +94,9 @@ def test_no_token_limit_does_not_raise(self): """With None token limit, high usage should not raise.""" harness = LLM_Harness(max_turns=None, max_exceptions=None, max_consecutive_exceptions=None, max_token_usage=None) - harness.check_after_step(_make_job(approx_token_usage=999_999)) + harness.check_after_step(_make_job( + approx_token_usage=999_999, difficulty_token_extension=10, + )) def test_max_exceptions_checked_after_step(self): """check_after_step should also enforce max_exceptions.""" diff --git a/tests/test_run_job_step.py b/tests/test_run_job_step.py index 2654b01..84f815a 100644 --- a/tests/test_run_job_step.py +++ b/tests/test_run_job_step.py @@ -25,7 +25,7 @@ from statek.executors.post_processor import PostProcessor from statek.executors.utils import handle_dialog, run_job_step from statek.future import FutureResult -from statek.exceptions import FutureError +from statek.exceptions import FutureError, LLM_HarnessError from statek.llm_harness import LLM_Harness from statek.llm_api import LLM_Response, LLM_StepData, LLM_Stats, CallParams, OpenRouter_API from statek.settings import LLM_API_Settings @@ -501,13 +501,92 @@ async def process_request(**kwargs): # pylint: disable=unused-argument await run_job_step(job, provider="OPENROUTER") assert job.get_current_difficulty() == TaskDifficulty.medium - assert job.chat_log[0].request_difficulty is None + assert job.last_chat_log_item.request_difficulty is None assert job.get_request_data(0)["model"] == "small" class TestRunJobStepHarnessIsolation: """Tests that harness token usage is isolated per job.""" + @pytest.mark.asyncio + async def test_two_panics_between_requests_each_grant_bonus(self, job_def_factory) -> None: + """Tool-result ownership and distinct escalation events survive a single tool batch.""" + job = Job(job_def=job_def_factory( + warmup_code=None, + metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "low"}, + ), job_status=JobStatus.STARTED) + response = LLM_Response( + step_data=LLM_StepData(text=None, call_requests=[ + CallParams(call_id="first", name="panic", args=[], kwargs={}), + CallParams(call_id="second", name="panic", args=[], kwargs={}), + ]), stats=LLM_Stats(1600, 0, None), + ) + continuation = LLM_Response( + step_data=LLM_StepData(text="pass", call_requests=None), + stats=LLM_Stats(2000, 0, None), + ) + api = MagicMock(process_request=AsyncMock(side_effect=[response, continuation])) + harness = LLM_Harness(None, 100, 100, 500) + with patch("statek.executors.utils.LLM_API") as api_class, \ + patch("statek.executors.utils.get_llm_harness", return_value=harness): + api_class.get.return_value = api + assert await run_job_step(job) is False + first_turn = job.last_chat_log_item + assert await run_job_step(job) is False + assert job.difficulty_token_extension == 800 + assert job.approx_token_usage == 900 + assert job.get_current_difficulty() == TaskDifficulty.high + assert len(first_turn.tool_log) == 2 + assert all(isinstance(result, str) and result for result in first_turn.tool_log) + assert job.num_turns == 2 + assert job.exception_count == 0 + + @pytest.mark.asyncio + @pytest.mark.parametrize("first_usage", [400, 501]) + async def test_panic_tool_extends_budget_only_when_executed( + self, job_def_factory, first_usage: int, + ) -> None: + """A pending panic cannot rescue an already over-budget provider response.""" + job = Job(job_def=job_def_factory( + warmup_code=None, + metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "low"}, + ), job_status=JobStatus.STARTED) + response = LLM_Response( + step_data=LLM_StepData(text=None, call_requests=[ + CallParams(call_id="panic_call", name="panic", args=[], kwargs={}), + ]), + stats=LLM_Stats(total_bytes_sent=first_usage * 4, + total_bytes_received=0, cost=None), + ) + continuation = LLM_Response( + step_data=LLM_StepData(text="pass", call_requests=None), + stats=LLM_Stats(total_bytes_sent=2000, total_bytes_received=0, cost=None), + ) + api = MagicMock(process_request=AsyncMock(side_effect=[response, continuation])) + harness = LLM_Harness(None, 100, 100, 500) + with patch("statek.executors.utils.LLM_API") as api_class, \ + patch("statek.executors.utils.get_llm_harness", return_value=harness): + api_class.get.return_value = api + if first_usage > 500: + with pytest.raises(LLM_HarnessError, match="501/500"): + await run_job_step(job) + assert job.get_current_difficulty() == TaskDifficulty.low + assert job.difficulty_token_extension == 0 + assert api.process_request.await_count == 1 + else: + assert await run_job_step(job) is False + assert job.difficulty_token_extension == 0 + assert await run_job_step(job) is False + assert job.get_current_difficulty() == TaskDifficulty.medium + assert job.difficulty_token_extension == 400 + llm_items = [item for item in job.chat_log if isinstance(item, LLM_LogItem)] + assert [item.cumulative_token_usage for item in llm_items] == [400, 900] + assert job.approx_token_usage == 900 + job.usage.total_bytes_received += 4 + with pytest.raises(LLM_HarnessError, match="901/900"): + harness.check_after_step(job) + assert harness.max_token_usage == 500 + @pytest.mark.asyncio async def test_two_jobs_same_harness_do_not_share_approx_token_usage( self, job_def_factory, db0_fixture @@ -1611,6 +1690,31 @@ async def test_cli_future_error_stores_tuple_instr_num(self, db0_fixture): # py class TestRunJobStepPostProcessing: """Tests runtime post-processor integration in run_job_step.""" + @pytest.mark.asyncio + async def test_suppressed_response_usage_is_included_in_later_escalation( + self, job_def_factory, + ) -> None: + """Bonus events capture accounted usage even when no LLM turn is retained.""" + job = Job(job_def=job_def_factory( + warmup_code=None, post_processing=RuntimeSuppressPostProcessor(), + metadata={"MODEL": "test-model", "DEFAULT_DIFFICULTY": "low"}, + ), job_status=JobStatus.STARTED) + api = MagicMock(process_request=AsyncMock(return_value=LLM_Response( + step_data=LLM_StepData(text="pass", call_requests=None), + stats=LLM_Stats(total_bytes_sent=40000, total_bytes_received=0, cost=None), + ))) + harness = LLM_Harness(None, 100, 100, 50000) + with patch("statek.executors.utils.LLM_API") as api_class, \ + patch("statek.executors.utils.get_llm_harness", return_value=harness): + api_class.get.return_value = api + assert await run_job_step(job) is False + assert not any(isinstance(item, LLM_LogItem) for item in job.chat_log) + job.panic() + assert job.difficulty_token_extension == 10000 + job.usage.total_bytes_sent = 240000 + harness.check_after_step(job) + + @staticmethod def _mock_api_and_harness(response): """Return a mocked LLM API and harness for one run_job_step call."""