Feature/issue 1535

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)

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,30 +1759,60 @@ 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:
             return last_difficulty
 
         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])

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

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(

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.