diff --git a/src/agent/loop.py b/src/agent/loop.py index d1dcf78..a62913a 100644 --- a/src/agent/loop.py +++ b/src/agent/loop.py @@ -6,6 +6,7 @@ import logging from ..config.settings import Settings +from ..llm.base import ResourceExhaustedError from ..store.factory import create_checkpointer from .graph import build_graph @@ -13,9 +14,9 @@ # Per-1M-token pricing: (input, output) _PRICING: dict[str, tuple[float, float]] = { - "anthropic": (3.0, 15.0), # Claude Sonnet - "gemini": (1.25, 10.0), # Gemini 2.5 Pro - "codex": (2.50, 10.0), # Codex + "anthropic": (3.0, 15.0), # Claude Sonnet + "gemini": (1.25, 10.0), # Gemini 2.5 Pro + "codex": (2.50, 10.0), # Codex } @@ -87,7 +88,11 @@ async def run(self) -> str: "repo_url": self.repo_config.url, } - final_state = await graph.ainvoke(initial_state, config=config) + try: + final_state = await graph.ainvoke(initial_state, config=config) + except ResourceExhaustedError as e: + logger.error("LLM resource exhausted. Aborting run: %s", e) + return "Agent stopped due to LLM resource exhaustion." token_usage = final_state.get("token_usage", {}) input_tokens = token_usage.get("input_tokens", 0) diff --git a/src/llm/base.py b/src/llm/base.py index 055e3b5..44c0a95 100644 --- a/src/llm/base.py +++ b/src/llm/base.py @@ -9,6 +9,10 @@ from langchain_core.output_parsers import BaseOutputParser +class ResourceExhaustedError(Exception): + """Raised when the LLM response indicates a resource exhaustion error.""" + + class RobustJsonOutputParser(BaseOutputParser[dict[str, Any]]): """Extract JSON from LLM responses, handling markdown fences and surrounding text. @@ -33,6 +37,9 @@ def parse(self, text: str | list) -> dict[str, Any]: parts.append(item.get("text", "")) text = "".join(parts) + if "RESOURCE_EXHAUSTED" in text: + raise ResourceExhaustedError("LLM response indicates resource exhaustion") + text = text.strip() # Try direct parse first (fast path for well-behaved responses) @@ -51,7 +58,7 @@ def parse(self, text: str | list) -> dict[str, Any]: continue # Brute-force: find every '{' and try json.loads from it, longest match first. - for m in re.finditer(r"\{", text): + for m in re.finditer(r"{", text): start = m.start() end = text.rfind("}", start) while end > start: diff --git a/tests/test_executor.py b/tests/test_executor.py index 54be9d6..871dc44 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -2,7 +2,10 @@ from pathlib import Path +import pytest + from src.agent.executor import list_files, read_file, run_command +from src.llm.base import ResourceExhaustedError, RobustJsonOutputParser def test_run_command_success(tmp_path: Path): @@ -40,3 +43,11 @@ def test_read_file(tmp_path: Path): def test_read_file_missing(tmp_path: Path): result = read_file(tmp_path / "nope.txt") assert "Error" in result + + +def test_robust_json_parser_resource_exhausted(): + """Verify parser raises ResourceExhaustedError on matching error text.""" + parser = RobustJsonOutputParser() + error_text = "Something went wrong. RESOURCE_EXHAUSTED. Please try again." + with pytest.raises(ResourceExhaustedError): + parser.parse(error_text)