From a7347edcd412f6b1bfcd207ccfbb2e3e128ad5cc Mon Sep 17 00:00:00 2001 From: mschoi Date: Tue, 27 Oct 2026 11:16:22 +0900 Subject: [PATCH] =?UTF-8?q?feat(agent):=20step-budget=20middleware=20?= =?UTF-8?q?=E2=80=94=20the=20agent=20knows=20its=20recursion=20budget?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ported from the Decepticon framework: near the recursion limit (80%) the agent receives a wrap-up reminder; at the hard stop a final-answer system message. The AutoBot graph (400 steps) now has token AND step budgets with in-context awareness — the agent consolidates and delivers instead of dying on GraphRecursionError. --- packages/autobot/autobot/agent/autobot.py | 10 +++- .../autobot/agent/middleware/step_budget.py | 56 ++++++++++++++++++ .../tests/unit/middleware/test_step_budget.py | 59 +++++++++++++++++++ 3 files changed, 123 insertions(+), 2 deletions(-) create mode 100644 packages/autobot/autobot/agent/middleware/step_budget.py create mode 100644 packages/autobot/tests/unit/middleware/test_step_budget.py diff --git a/packages/autobot/autobot/agent/autobot.py b/packages/autobot/autobot/agent/autobot.py index 17c7094..f46ba82 100644 --- a/packages/autobot/autobot/agent/autobot.py +++ b/packages/autobot/autobot/agent/autobot.py @@ -31,6 +31,7 @@ from autobot.agent.middleware.budget import TokenBudgetMiddleware from autobot.agent.middleware.prompt_injection_shield import PromptInjectionShieldMiddleware +from autobot.agent.middleware.step_budget import StepBudgetMiddleware from autobot.agent.middleware.untrusted_output import UntrustedOutputMiddleware from autobot.agent.models import build_model from autobot.agent.tools.blue_sensor import ( # noqa: F401 — re-exported for the spec @@ -43,7 +44,7 @@ from decepticon_core.plugin_loader import SubAgentSpec _ROLE = "autobot" -_RECURSION_LIMIT = 120 +_RECURSION_LIMIT = 400 _STANDARD_TOOLS: list[Any] = [ blue_sensor_scan, @@ -99,7 +100,12 @@ def create_autobot_agent( if system_prompt is None: system_prompt = load_system_prompt() if middleware is None: - middleware = default_middleware(model) + # StepBudgetMiddleware(total_steps) mirrors the with_config recursion + # limit — the agent knows its budget and wraps up before the hard stop. + middleware = [ + *default_middleware(model), + StepBudgetMiddleware(total_steps=recursion_limit or _RECURSION_LIMIT), + ] return create_agent( model, system_prompt=system_prompt, diff --git a/packages/autobot/autobot/agent/middleware/step_budget.py b/packages/autobot/autobot/agent/middleware/step_budget.py new file mode 100644 index 0000000..04a08f0 --- /dev/null +++ b/packages/autobot/autobot/agent/middleware/step_budget.py @@ -0,0 +1,56 @@ +"""Resource-awareness middleware — the agent knows its step budget. + +LangGraph enforces ``recursion_limit`` by raising GraphRecursionError at the +hard cap, which reads to the operator as "Recursion limit of N reached" and +discards the run's momentum. This middleware makes the budget visible to the +MODEL instead: near the cap it injects a wrap-up reminder, so the agent +consolidates and delivers its findings before exhausting the budget. + +The step estimate approximates graph supersteps from the message count +(one model round ≈ one message per step in a linear ReAct loop). It is an +approximation by design — cheap, monotonic, and directionally correct; the +hard cap stays with the runtime. +""" + +from __future__ import annotations + +from typing import Any + +from langchain.agents.middleware import AgentMiddleware +from langchain_core.messages import SystemMessage +from typing_extensions import override + + +class StepBudgetMiddleware(AgentMiddleware): + """Inject wrap-up guidance as the run approaches its step budget.""" + + def __init__(self, *, total_steps: int, warn_at_pct: float = 0.8) -> None: + super().__init__() + if total_steps <= 0: + raise ValueError("total_steps must be positive") + self._total = total_steps + self._warn_at = warn_at_pct + + @override + def wrap_model_call(self, request: Any, handler: Any) -> Any: + messages = getattr(request, "state", {}).get("messages") or [] + steps = len(messages) + if steps >= self._total: + note = ( + f"[step budget] HARD STOP REACHED — {steps} of {self._total} steps. " + "Do not call any more tools. Return your final answer now with the " + "findings and uncertainty collected so far." + ) + return handler( + request.override(messages=[*request.messages, SystemMessage(content=note)]) + ) + if steps >= self._total * self._warn_at: + note = ( + f"[step budget] {steps} of {self._total} steps used — approaching the " + "hard stop. Conclude the current line of work and move to delivering " + "your result." + ) + return handler( + request.override(messages=[*request.messages, SystemMessage(content=note)]) + ) + return handler(request) diff --git a/packages/autobot/tests/unit/middleware/test_step_budget.py b/packages/autobot/tests/unit/middleware/test_step_budget.py new file mode 100644 index 0000000..6c7fb22 --- /dev/null +++ b/packages/autobot/tests/unit/middleware/test_step_budget.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from langchain_core.messages import AIMessage, SystemMessage + +from autobot.agent.middleware.step_budget import StepBudgetMiddleware + + +def _state(n: int) -> SimpleNamespace: + messages = [AIMessage(content=f"m{i}") for i in range(n)] + + def override(*, messages): # ModelRequest.override semantics for tests + request = SimpleNamespace(state=_state(0), messages=messages) + request.override = request.override if hasattr(request, "override") else None + return SimpleNamespace(state=request.state, messages=messages, override=None) + + request = SimpleNamespace( + state={"messages": messages}, + messages=messages, + override=lambda *, messages: SimpleNamespace(messages=messages), + ) + return request + + +def test_injects_wrap_up_at_80pct(): + mw = StepBudgetMiddleware(total_steps=100) + request = _state(85) + handled = mw.wrap_model_call(request, lambda r: r) + notes = [ + m for m in handled.messages if isinstance(m, SystemMessage) and "step budget" in m.content + ] + assert len(notes) == 1 and "85 of 100" in notes[0].content + + +def test_injects_hard_stop_at_100pct(): + mw = StepBudgetMiddleware(total_steps=100) + request = _state(100) + handled = mw.wrap_model_call(request, lambda r: r) + notes = [ + m for m in handled.messages if isinstance(m, SystemMessage) and "HARD STOP" in m.content + ] + assert len(notes) == 1 and "Do not call any more tools" in notes[0].content + + +def test_silent_well_under_budget(): + mw = StepBudgetMiddleware(total_steps=100) + request = _state(10) + handled = mw.wrap_model_call(request, lambda r: r) + notes = [ + m for m in handled.messages if isinstance(m, SystemMessage) and "step budget" in m.content + ] + assert notes == [] + + +def test_invalid_total_rejected(): + with pytest.raises(ValueError): + StepBudgetMiddleware(total_steps=0)