diff --git a/libs/openant-core/tests/test_agent.py b/libs/openant-core/tests/test_agent.py new file mode 100644 index 00000000..3b615838 --- /dev/null +++ b/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}" + ) diff --git a/libs/openant-core/utilities/agentic_enhancer/agent.py b/libs/openant-core/utilities/agentic_enhancer/agent.py index 62061b7a..e932d40d 100644 --- a/libs/openant-core/utilities/agentic_enhancer/agent.py +++ b/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 diff --git a/libs/openant-core/utilities/agentic_enhancer/prompts.py b/libs/openant-core/utilities/agentic_enhancer/prompts.py index dd9ca83c..9ed3defb 100644 --- a/libs/openant-core/utilities/agentic_enhancer/prompts.py +++ b/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: