diff --git a/nanobot/agent/runner.py b/nanobot/agent/runner.py index 436a1a122..2afd59385 100644 --- a/nanobot/agent/runner.py +++ b/nanobot/agent/runner.py @@ -50,6 +50,7 @@ from nanobot.utils.runtime import ( build_finalization_retry_message, build_goal_continue_message, build_length_recovery_message, + build_runtime_budget_notice_message, is_blank_text, repeated_external_lookup_error, repeated_workspace_violation_error, @@ -67,6 +68,7 @@ _MAX_EMPTY_RETRIES = 2 _MAX_LENGTH_RECOVERIES = 3 _MAX_INJECTIONS_PER_TURN = 3 _MAX_INJECTION_CYCLES = 5 +_BUDGET_NOTICE_MIN_ITERATIONS = 20 # Backward-compatible module attribute for tests/extensions that monkeypatch # the former single-file tracker hook. Runtime uses prepare_file_edit_trackers. prepare_file_edit_tracker = _prepare_file_edit_tracker @@ -357,6 +359,7 @@ class AgentRunner: length_recovery_count = 0 had_injections = False injection_cycles = 0 + budget_notice_level_sent = 0 compacted_tool_call_ids: set[str] = set() governance_config = ContextGovernanceConfig( provider=self.provider, @@ -511,6 +514,12 @@ class AgentRunner: ) if _drained: had_injections = True + budget_notice_level_sent = self._append_runtime_budget_notice_if_needed( + spec, + messages, + completed_iterations=iteration + 1, + sent_level=budget_notice_level_sent, + ) await hook.after_iteration(context) continue @@ -940,6 +949,53 @@ class AgentRunner: retry_messages.append(build_budget_exhausted_finalization_message()) return retry_messages + @classmethod + def _append_runtime_budget_notice_if_needed( + cls, + spec: AgentRunSpec, + messages: list[dict[str, Any]], + *, + completed_iterations: int, + sent_level: int, + ) -> int: + level = cls._runtime_budget_notice_level( + max_iterations=spec.max_iterations, + completed_iterations=completed_iterations, + ) + if level <= sent_level: + return sent_level + + remaining_iterations = max(0, spec.max_iterations - completed_iterations) + messages.append(build_runtime_budget_notice_message( + level=level, + max_iterations=spec.max_iterations, + used_iterations=completed_iterations, + remaining_iterations=remaining_iterations, + )) + return level + + @staticmethod + def _runtime_budget_notice_level( + *, + max_iterations: int, + completed_iterations: int, + ) -> int: + """Return the convergence-warning level for a long tool loop.""" + if max_iterations < _BUDGET_NOTICE_MIN_ITERATIONS: + return 0 + + remaining_iterations = max_iterations - completed_iterations + if remaining_iterations <= 0: + return 0 + + convergence_threshold = max(5, (max_iterations + 9) // 10) + final_threshold = max(3, (max_iterations + 32) // 33) + if remaining_iterations <= final_threshold: + return 2 + if remaining_iterations <= convergence_threshold: + return 1 + return 0 + @staticmethod def _max_iterations_fallback(spec: AgentRunSpec) -> str: if spec.max_iterations_message: diff --git a/nanobot/utils/runtime.py b/nanobot/utils/runtime.py index 9141583ea..bb423b105 100644 --- a/nanobot/utils/runtime.py +++ b/nanobot/utils/runtime.py @@ -42,6 +42,27 @@ SUSTAINED_GOAL_CONTINUE_PROMPT = ( "objective using your tools, or call complete_goal if the work is truly finished." ) +RUNTIME_BUDGET_CONVERGENCE_PROMPT = """\ +[Runtime Budget Notice] +You have used {used_iterations} of {max_iterations} model/tool iterations for this turn. \ +{remaining_iterations} iteration(s) remain before NanoBot must finalize without more tools. + +Switch to convergence mode: stop broad exploration, choose the smallest high-signal command or edit, \ +verify the likely solution, and preserve enough budget for a final answer. For coding or \ +file-producing tasks, do not mark the work complete until the smallest reliable verification passes, \ +or clearly state remaining failures. +[/Runtime Budget Notice]""" + +RUNTIME_BUDGET_FINAL_PROMPT = """\ +[Runtime Budget Notice] +Only {remaining_iterations} of {max_iterations} model/tool iteration(s) remain before NanoBot must \ +finalize without more tools. + +Finalize the solution path now: avoid new broad searches or builds unless essential, make the \ +smallest final fix or artifact, run one targeted verification if possible, then answer honestly with \ +the evidence or remaining failures. +[/Runtime Budget Notice]""" + def empty_tool_result_message(tool_name: str) -> str: """Short prompt-safe marker for tools that completed without visible output.""" @@ -88,6 +109,25 @@ def build_goal_continue_message(custom: str | None = None) -> dict[str, str]: return {"role": "user", "content": custom or SUSTAINED_GOAL_CONTINUE_PROMPT} +def build_runtime_budget_notice_message( + *, + level: int, + max_iterations: int, + used_iterations: int, + remaining_iterations: int, +) -> dict[str, str]: + """Prompt the model to converge as the generic tool-iteration budget runs low.""" + template = RUNTIME_BUDGET_FINAL_PROMPT if level >= 2 else RUNTIME_BUDGET_CONVERGENCE_PROMPT + return { + "role": "user", + "content": template.format( + max_iterations=max_iterations, + used_iterations=used_iterations, + remaining_iterations=remaining_iterations, + ), + } + + def external_lookup_signature(tool_name: str, arguments: Any) -> str | None: """Stable signature for repeated external lookups we want to throttle.""" if not isinstance(arguments, dict): diff --git a/tests/agent/test_runner_tool_execution.py b/tests/agent/test_runner_tool_execution.py index 70e74fafe..53bd2adef 100644 --- a/tests/agent/test_runner_tool_execution.py +++ b/tests/agent/test_runner_tool_execution.py @@ -358,3 +358,79 @@ async def test_runner_blocks_repeated_external_fetches(): if msg.get("role") == "tool" and msg.get("tool_call_id") == "call_3" ][0] assert "repeated external lookup blocked" in blocked_tool_message["content"] + + +@pytest.mark.asyncio +async def test_runner_adds_budget_notice_near_long_tool_budget(): + provider = MagicMock() + captured_final_call: list[dict] = [] + call_count = {"n": 0} + + async def chat_with_retry(*, messages, **kwargs): + call_count["n"] += 1 + if call_count["n"] <= 16: + return LLMResponse( + content="working", + tool_calls=[ToolCallRequest(id=f"call_{call_count['n']}", name="work", arguments={})], + usage={}, + ) + captured_final_call[:] = messages + return LLMResponse(content="done", tool_calls=[], usage={}) + + provider.chat_with_retry = chat_with_retry + tools = MagicMock() + tools.get_definitions.return_value = [] + tools.execute = AsyncMock(return_value="tool result") + + result = await AgentRunner(provider).run(AgentRunSpec( + initial_messages=[{"role": "user", "content": "finish a large task"}], + tools=tools, + model="test-model", + max_iterations=20, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + )) + + assert result.final_content == "done" + notices = [ + msg["content"] + for msg in captured_final_call + if msg.get("role") == "user" and "[Runtime Budget Notice]" in str(msg.get("content")) + ] + assert len(notices) == 1 + assert "15 of 20 model/tool iterations" in notices[0] + assert "Switch to convergence mode" in notices[0] + assert tools.execute.await_count == 16 + + +@pytest.mark.asyncio +async def test_runner_budget_notice_does_not_affect_short_runs(): + provider = MagicMock() + captured_final_call: list[dict] = [] + call_count = {"n": 0} + + async def chat_with_retry(*, messages, **kwargs): + call_count["n"] += 1 + if call_count["n"] <= 2: + return LLMResponse( + content="working", + tool_calls=[ToolCallRequest(id=f"call_{call_count['n']}", name="work", arguments={})], + usage={}, + ) + captured_final_call[:] = messages + return LLMResponse(content="done", tool_calls=[], usage={}) + + provider.chat_with_retry = chat_with_retry + tools = MagicMock() + tools.get_definitions.return_value = [] + tools.execute = AsyncMock(return_value="tool result") + + result = await AgentRunner(provider).run(AgentRunSpec( + initial_messages=[{"role": "user", "content": "small task"}], + tools=tools, + model="test-model", + max_iterations=4, + max_tool_result_chars=_MAX_TOOL_RESULT_CHARS, + )) + + assert result.final_content == "done" + assert all("[Runtime Budget Notice]" not in str(msg.get("content")) for msg in captured_final_call)