From 2273970d4bc74bd5afbd9bf4c6b7ea2e1819ea73 Mon Sep 17 00:00:00 2001 From: Ankit Bari Date: Tue, 13 Oct 2026 17:33:26 +0530 Subject: [PATCH] Add opt-in cost and runtime caps to agent loop --- openhack/agents/base.py | 39 ++++++++++++++ openhack/config.py | 9 ++++ tests/test_agent_budget.py | 102 +++++++++++++++++++++++++++++++++++++ 3 files changed, 150 insertions(+) create mode 100644 tests/test_agent_budget.py diff --git a/openhack/agents/base.py b/openhack/agents/base.py index 7d16e37..c0a4104 100644 --- a/openhack/agents/base.py +++ b/openhack/agents/base.py @@ -237,6 +237,29 @@ def _record_response_usage(self, response) -> None: self.session.total_output_tokens += response.usage.get("output_tokens", 0) self.context_manager.update_usage(response.usage.get("input_tokens", 0)) + def _budget_exceeded(self, loop_started: float) -> Optional[str]: + """Return a human-readable reason if a spend/runtime cap is hit. + + Both caps are opt-in: a value <= 0 disables that check. Cost is + session-wide (agents share ``session.total_cost``); runtime is measured + from the start of this agent loop. + """ + max_cost = getattr(settings, "agent_max_cost_usd", 0.0) or 0.0 + if max_cost > 0 and self.session.total_cost >= max_cost: + return ( + f"cost cap reached (${self.session.total_cost:.4f} >= " + f"${max_cost:.4f}; agent_max_cost_usd)" + ) + max_runtime = getattr(settings, "agent_max_runtime_s", 0.0) or 0.0 + if max_runtime > 0: + elapsed = time.monotonic() - loop_started + if elapsed >= max_runtime: + return ( + f"runtime cap reached ({elapsed:.0f}s >= " + f"{max_runtime:.0f}s; agent_max_runtime_s)" + ) + return None + def _estimate_tokens(self, messages: list[Message], system: str) -> int: """Rough token estimate: ~4 chars per token for English text.""" total_chars = len(system) @@ -403,6 +426,7 @@ async def _agent_loop(self) -> dict: max_iterations = getattr(settings, "agent_max_iterations", 0) or 0 stale_limit = getattr(settings, "agent_stale_turn_limit", 8) iteration = 0 + loop_started = time.monotonic() # Natural prose returned immediately before the continuation guard. # Some providers obey the guard by sending finish_task in a separate # response; retain the prose so a no-op greeting does not acquire a @@ -419,6 +443,21 @@ async def _agent_loop(self) -> dict: await self.session.wait_if_paused() if self.session.cancelled: break + budget_reason = self._budget_exceeded(loop_started) + if budget_reason: + logger.warning(f"[{self.name}] Stopping: {budget_reason}") + result = { + "error": "Budget exceeded", + "detail": budget_reason, + "partial_result": self._last_assistant_content(), + } + self.session.record_event( + "agent_loop_stopped", + {"reason": "budget_exceeded", **result}, + agent=self.name, + ) + self.session.finish_turn("budget_exceeded", result, agent=self.name) + return result iteration += 1 self.session.record_event( "agent_iteration_started", diff --git a/openhack/config.py b/openhack/config.py index e936492..54b083d 100644 --- a/openhack/config.py +++ b/openhack/config.py @@ -157,6 +157,15 @@ class Settings(BaseSettings): agent_max_iterations: int = 0 agent_stale_turn_limit: int = 8 + # Spend / wall-clock backstops. With no iteration cutoff, the stale-turn + # stop only catches thrash; a run that keeps producing novel output has no + # other bound. These are opt-in (0 = unlimited, matching the no-cutoff + # philosophy) and are checked before every model call. + # agent_max_cost_usd: session-wide spend (all agents share session.total_cost) + # agent_max_runtime_s: wall-clock seconds for a single agent loop + agent_max_cost_usd: float = 0.0 + agent_max_runtime_s: float = 0.0 + # Scan scoping — exclude paths that are never production web attack surface scan_exclude_patterns: list[str] = [ "**/test/**", "**/tests/**", "**/__tests__/**", "**/spec/**", diff --git a/tests/test_agent_budget.py b/tests/test_agent_budget.py new file mode 100644 index 0000000..b75e027 --- /dev/null +++ b/tests/test_agent_budget.py @@ -0,0 +1,102 @@ +"""Opt-in spend / wall-clock backstops for the agent loop. + +With the iteration cutoff removed, a run that keeps producing novel output has +no bound other than the stale-turn stop. `agent_max_cost_usd` and +`agent_max_runtime_s` give operators a hard ceiling. Both default to 0 +(unlimited) so existing behaviour is unchanged. +""" + +import asyncio + +from openhack.agents.llm import LLMResponse, ToolCall +from openhack.agents.session import Session +from openhack.config import settings +from openhack.tui import OpenHackApp + +from tests.test_no_iteration_cutoff import _ProbeAgent, _UniqueResultTools + + +class _CostlyLLM: + """Productive forever (never finish_task); every call costs `cost_per_call`.""" + + model = "test-model" + provider = "test-provider" + + def __init__(self, cost_per_call: float = 0.0, delay: float = 0.0): + self.cost_per_call = cost_per_call + self.delay = delay + self.calls = 0 + self.event_callback = None + + async def chat(self, **kwargs): + self.calls += 1 + if self.delay: + await asyncio.sleep(self.delay) + return LLMResponse( + tool_calls=[ + ToolCall(id=f"p{self.calls}", name="probe", arguments={"n": self.calls}) + ], + finish_reason="tool_calls", + model_call_id=f"m{self.calls}", + cost=self.cost_per_call, + ) + + +def _run(tmp_path, llm, monkeypatch, max_calls=50): + session = Session(str(tmp_path), persist_events=False) + tools = _UniqueResultTools(tmp_path, include_agent_tools=True, session=session) + agent = _ProbeAgent(llm, tools, session) + # Safety net so a regression in the cap can't hang the suite. + monkeypatch.setattr(settings, "agent_max_iterations", max_calls) + return session, asyncio.run(agent.run("go")) + + +def test_defaults_are_unlimited(): + assert settings.agent_max_cost_usd == 0.0 + assert settings.agent_max_runtime_s == 0.0 + + +def test_cost_cap_stops_the_loop(tmp_path, monkeypatch): + monkeypatch.setattr(settings, "agent_max_cost_usd", 0.25) + llm = _CostlyLLM(cost_per_call=0.10) + + session, result = _run(tmp_path, llm, monkeypatch) + + # 3 calls -> $0.30 >= $0.25, checked before the 4th call. + assert llm.calls == 3 + assert result["error"] == "Budget exceeded" + assert "agent_max_cost_usd" in result["detail"] + assert any( + e.event_type == "agent_loop_stopped" + and (e.data or {}).get("reason") == "budget_exceeded" + for e in session.events + ) + + +def test_runtime_cap_stops_the_loop(tmp_path, monkeypatch): + monkeypatch.setattr(settings, "agent_max_cost_usd", 0.0) + monkeypatch.setattr(settings, "agent_max_runtime_s", 0.2) + llm = _CostlyLLM(delay=0.08) + + session, result = _run(tmp_path, llm, monkeypatch) + + assert result["error"] == "Budget exceeded" + assert "agent_max_runtime_s" in result["detail"] + assert 2 <= llm.calls <= 6 + + +def test_caps_disabled_do_not_interfere(tmp_path, monkeypatch): + monkeypatch.setattr(settings, "agent_max_cost_usd", 0.0) + monkeypatch.setattr(settings, "agent_max_runtime_s", 0.0) + llm = _CostlyLLM(cost_per_call=100.0) + + _, result = _run(tmp_path, llm, monkeypatch, max_calls=7) + + # Only the iteration backstop (7) stops it, not the budget. + assert llm.calls == 7 + assert result["error"] == "Max iterations reached" + + +def test_budget_stop_is_reported_incomplete_not_completed(): + result = {"error": "Budget exceeded", "detail": "x", "partial_result": ""} + assert OpenHackApp._agent_result_status(result) == "incomplete"