fix(agentic-enhancer): cap conversation input to avoid context overflow

libs/openant-core/tests/test_agent.py

@@ -0,0 +1,93 @@
+"""Regression tests for the agentic enhancer input budget.
+
+The agentic context enhancer built the LLM conversation with no input
+token/char budget: ``primary_code`` was inlined verbatim into the initial
+user prompt, and raw, untruncated tool results were appended to ``messages``
+every iteration. On a large unit the constructed input grew unbounded across
+iterations until it overflowed the model context (400 error), losing the
+enhancement.
+
+Operative guard: an INPUT BUDGET applied at the consumption points.
+  - ``get_user_prompt`` caps the inlined ``primary_code`` (prompts.py).
+  - ``cap_tool_result_content`` caps each serialized tool result before it is
+    appended to ``messages`` (agent.py).
+
+These tests exercise the pure, deterministic guard functions (no network).
+"""
+import json
+
+from utilities.agentic_enhancer.prompts import get_user_prompt
+from utilities.agentic_enhancer import agent as agent_mod
+
+
+def test_primary_code_is_capped_in_user_prompt():
+    """A huge primary_code must not be inlined verbatim — the constructed
+    prompt must stay bounded (well under any model context budget)."""
+    huge_code = "x = 1\n" * 500_000  # ~3 MB of source
+    prompt = get_user_prompt(
+        unit_id="big.py:huge",
+        unit_type="function",
+        primary_code=huge_code,
+        static_deps=[],
+        static_callers=[],
+    )
+    # The whole prompt (not just the code) must be bounded.
+    assert len(prompt) <= agent_mod.MAX_PROMPT_CHARS + 4096, (
+        f"prompt length {len(prompt)} exceeds budget "
+        f"{agent_mod.MAX_PROMPT_CHARS}; primary_code was not capped"
+    )
+    # Truncation must be signalled so the model knows content was elided.
+    assert "truncated" in prompt.lower()
+
+
+def test_small_primary_code_is_not_altered():
+    """Small code must pass through verbatim (no spurious truncation)."""
+    small_code = "def f():\n    return 1\n"
+    prompt = get_user_prompt(
+        unit_id="small.py:f",
+        unit_type="function",
+        primary_code=small_code,
+        static_deps=[],
+        static_callers=[],
+    )
+    assert small_code in prompt
+    assert "... (truncated" not in prompt
+
+
+def test_tool_result_content_is_capped():
+    """A very large tool result must be truncated before being appended to
+    the messages list so the constructed input cannot grow unbounded."""
+    huge_result = {"found": True, "code": "A" * 1_000_000}
+    content = agent_mod.cap_tool_result_content(huge_result)
+    assert isinstance(content, str)
+    assert len(content) <= agent_mod.MAX_TOOL_RESULT_CHARS + 256, (
+        f"tool-result content length {len(content)} exceeds budget "
+        f"{agent_mod.MAX_TOOL_RESULT_CHARS}"
+    )
+    assert "truncated" in content.lower()
+
+
+def test_small_tool_result_round_trips_as_json():
+    """A small tool result must serialize losslessly (parseable JSON)."""
+    result = {"found": True, "id": "a.py:f", "code": "def f(): pass"}
+    content = agent_mod.cap_tool_result_content(result)
+    assert json.loads(content) == result
+
+
+def test_messages_do_not_grow_unbounded_across_iterations():
+    """Simulate the loop's append behaviour: repeatedly appending capped tool
+    results keeps total input bounded per iteration (no unbounded growth from
+    a single oversized result)."""
+    messages = []
+    oversized = {"found": True, "code": "Z" * 5_000_000}
+    for _ in range(agent_mod.MAX_ITERATIONS):
+        content = agent_mod.cap_tool_result_content(oversized)
+        messages.append({"role": "user", "content": [
+            {"type": "tool_result", "tool_use_id": "t", "content": content}
+        ]})
+    total = sum(len(m["content"][0]["content"]) for m in messages)
+    bound = agent_mod.MAX_ITERATIONS * (agent_mod.MAX_TOOL_RESULT_CHARS + 256)
+    assert total <= bound, (
+        f"accumulated tool-result input {total} exceeds bounded "
+        f"expectation {bound}"
+    )

libs/openant-core/utilities/agentic_enhancer/agent.py

@@ -32,6 +32,30 @@
 MAX_ITERATIONS = 20
 MAX_TOKENS_PER_RESPONSE = 4096
 
+# Input budget.
+# The conversation input had no budget: primary_code was inlined verbatim and
+# raw tool results were appended every iteration, so input grew unbounded until
+# it overflowed the model context (400). We cap each oversized input at its
+# consumption point. ~4 chars/token, so these stay well under the model window.
+MAX_PROMPT_CHARS = 60_000          # cap on inlined primary_code in the prompt
+MAX_TOOL_RESULT_CHARS = 24_000     # cap on each serialized tool result
+
+
+def cap_tool_result_content(result: dict, limit: int = MAX_TOOL_RESULT_CHARS) -> str:
+    """Serialize a tool result to JSON, truncating to ``limit`` chars.
+
+    Tool results (e.g. ``read_function`` returning a whole function body) are
+    otherwise appended raw to the conversation, growing the input without
+    bound across iterations. Small results round-trip as valid JSON; oversized
+    results are truncated with an explicit marker so the model knows content
+    was elided.
+    """
+    content = json.dumps(result)
+    if len(content) <= limit:
+        return content
+    marker = "\n... (truncated)"
+    return content[: limit - len(marker)] + marker
+
 
 class AgentResult:
     """Result from agent analysis."""
@@ -213,6 +237,28 @@ def analyze_unit(
                     "output_tokens": total_output_tokens,
                 }
                 raise
+            except anthropic.BadRequestError:
+                # 400 (e.g. constructed input still overflows the context
+                # window). Degrade gracefully instead of losing the whole
+                # enhancement — return a neutral, incomplete result. The input
+                # budget above makes this path rare, but it guards the residual.
+                if self.verbose:
+                    print(f"  BadRequest at iteration {iterations}; "
+                          "returning incomplete result")
+                return AgentResult(
+                    include_functions=[],
+                    usage_context="Analysis terminated - request rejected (400)",
+                    security_classification="neutral",
+                    classification_reasoning="Analysis incomplete - request rejected",
+                    confidence=0.2,
+                    iterations=iterations,
+                    total_tokens=total_input_tokens + total_output_tokens,
+                    is_entry_point=is_entry_point,
+                    reachable_from_entry=reachable_from_entry,
+                    entry_point_path=entry_point_path,
+                    input_tokens=total_input_tokens,
+                    output_tokens=total_output_tokens,
+                )
             except Exception as exc:
                 # Attach agent state so the caller knows how far we got
                 exc.agent_state = {
@@ -285,14 +331,14 @@ def analyze_unit(
                         tool_results.append({
                             "type": "tool_result",
                             "tool_use_id": tool_use_id,
-                            "content": json.dumps(result)
+                            "content": cap_tool_result_content(result)
                         })
                         break
                     else:
                         tool_results.append({
                             "type": "tool_result",
                             "tool_use_id": tool_use_id,
-                            "content": json.dumps(result)
+                            "content": cap_tool_result_content(result)
                         })
 
             # If finish was called, return result

libs/openant-core/utilities/agentic_enhancer/prompts.py

@@ -9,6 +9,20 @@
 from typing import List, Optional
 
 
+# Input budget for the inlined unit code.
+# primary_code was previously inlined verbatim, so a large unit overflowed the
+# model context. ~4 chars/token, so this stays well under the model window.
+MAX_PRIMARY_CODE_CHARS = 60_000
+
+
+def _cap_primary_code(primary_code: str, limit: int = MAX_PRIMARY_CODE_CHARS) -> str:
+    """Cap inlined unit code so the prompt stays within the input budget."""
+    if len(primary_code) <= limit:
+        return primary_code
+    marker = "\n... (truncated)"
+    return primary_code[: limit - len(marker)] + marker
+
+
 SYSTEM_PROMPT = """You are a security code analyst. Your task is to classify code based on:
 
 1. Does it contain security flaws (dangerous operations)?
@@ -101,6 +115,9 @@ def get_user_prompt(
     deps_str = ", ".join(static_deps[:10]) if static_deps else "None identified"
     callers_str = ", ".join(static_callers[:10]) if static_callers else "None identified"
 
+    # Cap the inlined code so a large unit cannot overflow the model context.
+    primary_code = _cap_primary_code(primary_code)
+
     # Build reachability section
     reachability_section = ""
     if is_entry_point: