Import AITURK IDE 1.0.0-beta.1 from Hermes 63279301; preserve MIT license

tests/run_agent/test_background_review_input_budget.py

@@ -0,0 +1,239 @@
+"""Regression tests for the background-review aggregate input budget (#93057).
+
+The review fork replays its snapshot on every provider request in its tool
+loop. Detached in-memory compaction bounds any SINGLE request; the aggregate
+input budget (``_review_input_token_budget``, set by
+``_run_review_in_thread`` from ``auxiliary.background_review.max_input_tokens``)
+bounds the WHOLE review: the tool loop stops before the provider call that
+would cross it, mirroring the iteration-budget exit.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from run_agent import AIAgent
+
+
+def _tool_call() -> SimpleNamespace:
+    return SimpleNamespace(
+        id="call_1",
+        type="function",
+        function=SimpleNamespace(name="web_search", arguments='{"query": "x"}'),
+    )
+
+
+def _tool_response(prompt_tokens: int) -> SimpleNamespace:
+    message = SimpleNamespace(
+        content=None,
+        reasoning_content=None,
+        reasoning=None,
+        tool_calls=[_tool_call()],
+    )
+    return SimpleNamespace(
+        choices=[SimpleNamespace(message=message, finish_reason="tool_calls")],
+        model="test/model",
+        usage=SimpleNamespace(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=1,
+            total_tokens=prompt_tokens + 1,
+        ),
+    )
+
+
+def _final_response() -> SimpleNamespace:
+    message = SimpleNamespace(
+        content="done",
+        reasoning_content=None,
+        reasoning=None,
+        tool_calls=None,
+    )
+    return SimpleNamespace(
+        choices=[SimpleNamespace(message=message, finish_reason="stop")],
+        model="test/model",
+        usage=None,
+    )
+
+
+def _tool_definition() -> dict:
+    return {
+        "type": "function",
+        "function": {
+            "name": "web_search",
+            "description": "Search the web",
+            "parameters": {
+                "type": "object",
+                "properties": {"query": {"type": "string"}},
+                "required": ["query"],
+            },
+        },
+    }
+
+
+def _make_loop_agent():
+    with (
+        patch("run_agent.get_tool_definitions", return_value=[_tool_definition()]),
+        patch("run_agent.check_toolset_requirements", return_value={}),
+        patch("run_agent.OpenAI"),
+        patch("agent.model_metadata.get_model_context_length", return_value=256_000),
+        patch("agent.context_compressor.get_model_context_length", return_value=256_000),
+    ):
+        agent = AIAgent(
+            api_key="test-key-1234567890",
+            base_url="https://openrouter.ai/api/v1",
+            model="test/model",
+            quiet_mode=True,
+            skip_context_files=True,
+            skip_memory=True,
+            max_iterations=10,
+        )
+
+    agent.client = MagicMock()
+    agent._cached_system_prompt = "You are helpful."
+    agent._use_prompt_caching = False
+    agent._disable_streaming = True
+    agent.tool_delay = 0
+    agent.save_trajectories = False
+    agent.max_compression_attempts = 1
+
+    compressor = MagicMock()
+    compressor.protect_first_n = 3
+    compressor.protect_last_n = 20
+    compressor.threshold_tokens = 999_999_999  # never fire compaction here
+    compressor.context_length = 1_000_000_000
+    compressor.last_prompt_tokens = -1
+    compressor._verify_compaction_cleared_threshold = False
+    compressor.awaiting_real_usage_after_compression = False
+    compressor.should_compress.return_value = False
+    compressor.should_compress_info.return_value = (False, None)
+    compressor.should_compress_preflight.return_value = False
+    compressor.should_defer_preflight_to_real_usage.return_value = False
+    compressor.get_active_compression_failure_cooldown.return_value = None
+    compressor.select_context.return_value = None
+    compressor.get_automatic_compaction_status_message.return_value = ""
+    agent.compression_enabled = False  # isolate the budget behavior under test
+    agent.context_compressor = compressor
+
+    def _fake_execute_tool_calls(assistant_message, messages, *_args):
+        tool_call = assistant_message.tool_calls[0]
+        messages.append(
+            {
+                "role": "tool",
+                "name": tool_call.function.name,
+                "tool_call_id": tool_call.id,
+                "content": "ok",
+            }
+        )
+
+    agent._execute_tool_calls = _fake_execute_tool_calls
+    return agent
+
+
+def _run_with_responses(agent, responses):
+    agent.client.chat.completions.create.side_effect = responses
+    with (
+        patch.object(agent, "_flush_messages_to_session_db", return_value=True),
+        patch.object(agent, "_persist_session"),
+        patch.object(agent, "_save_trajectory"),
+        patch.object(agent, "_cleanup_task_resources"),
+    ):
+        result = agent.run_conversation("do some tool work")
+    return result
+
+
+def test_review_input_budget_stops_tool_loop_before_next_provider_call():
+    """Once a fork's cumulative input crosses its budget, no further provider
+    call is made — the crossing request completes, then the loop stops."""
+    agent = _make_loop_agent()
+    agent._review_input_token_budget = 100_000
+
+    responses = [
+        _tool_response(50_000),
+        _tool_response(50_000),  # cumulative 100_000 -> budget crossed
+        _tool_response(50_000),  # must never be consumed
+        _final_response(),
+    ]
+    result = _run_with_responses(agent, responses)
+
+    create = agent.client.chat.completions.create
+    assert create.call_count == 2, (
+        f"expected the loop to stop after crossing the input budget, "
+        f"but {create.call_count} provider calls were made (budget "
+        f"{agent._review_input_token_budget}, "
+        f"used {agent.session_input_tokens})"
+    )
+    assert agent.session_input_tokens == 100_000
+    assert result["completed"] is False
+
+
+def test_no_budget_attribute_leaves_tool_loop_unbounded():
+    """Agents without ``_review_input_token_budget`` (every normal agent)
+    are unaffected by the gate and consume all scripted responses."""
+    agent = _make_loop_agent()
+
+    responses = [
+        _tool_response(50_000),
+        _tool_response(50_000),
+        _tool_response(50_000),
+        _final_response(),
+    ]
+    result = _run_with_responses(agent, responses)
+
+    assert agent.client.chat.completions.create.call_count == 4
+    assert result["completed"] is True
+    assert result["final_response"] == "done"
+
+
+def test_review_input_budget_exhausted_predicate_edge_cases():
+    """The gate only arms for a positive int budget and a real token count."""
+    from agent.conversation_loop import _review_input_budget_exhausted
+
+    class _Agent:
+        pass
+
+    agent = _Agent()
+    assert _review_input_budget_exhausted(agent) is False
+
+    agent._review_input_token_budget = None
+    agent.session_input_tokens = 999_999
+    assert _review_input_budget_exhausted(agent) is False
+
+    agent._review_input_token_budget = 0
+    assert _review_input_budget_exhausted(agent) is False
+
+    agent._review_input_token_budget = -1
+    assert _review_input_budget_exhausted(agent) is False
+
+    agent._review_input_token_budget = "100"
+    assert _review_input_budget_exhausted(agent) is False
+
+    agent._review_input_token_budget = True
+    assert _review_input_budget_exhausted(agent) is False
+
+    agent._review_input_token_budget = 100_000
+    agent.session_input_tokens = 99_999
+    assert _review_input_budget_exhausted(agent) is False
+
+    agent.session_input_tokens = 100_000
+    assert _review_input_budget_exhausted(agent) is True
+
+
+@pytest.mark.parametrize(
+    ("config_value", "expected"),
+    [
+        ({}, 600_000),
+        ({"max_input_tokens": 1_000_000}, 1_000_000),
+        ({"max_input_tokens": 0}, None),
+        ({"max_input_tokens": -5}, None),
+        ({"max_input_tokens": "not-a-number"}, 600_000),
+        ({"max_input_tokens": "300000"}, 300_000),
+    ],
+)
+def test_review_input_token_budget_resolution(config_value, expected):
+    """Config parsing: default, override, explicit disable, garbage fallback."""
+    from agent.background_review import _review_input_token_budget
+
+    assert _review_input_token_budget(config_value) == expected