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

tests/agent/test_moa_slot_max_tokens.py

@@ -0,0 +1,82 @@
+"""Tests for per-slot max_tokens in MoA reference calls.
+
+Verifies that a ``max_tokens`` field on a reference slot dict takes
+precedence over the preset-level ``reference_max_tokens``, and that
+slot-level max_tokens=None falls back to the preset-level cap.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+
+class TestRunReferenceSlotMaxTokens:
+    """_run_reference should prefer slot-level max_tokens over preset-level."""
+
+    def test_slot_max_tokens_overrides_preset_level(self):
+        """When slot has max_tokens, it overrides the preset-level cap."""
+        from agent.moa_loop import _run_reference
+
+        captured_kwargs: dict = {}
+
+        def fake_call_llm(**kwargs):
+            captured_kwargs.update(kwargs)
+            mock_resp = MagicMock()
+            mock_resp.choices = [MagicMock(message=MagicMock(content="advice"))]
+            mock_resp.usage = None
+            return mock_resp
+
+        slot = {"provider": "openrouter", "model": "deepseek/deepseek-v4-pro", "max_tokens": 600}
+
+        with patch("agent.moa_loop._slot_runtime", return_value={"provider": "openrouter", "model": "deepseek/deepseek-v4-pro"}), \
+             patch("agent.moa_loop.call_llm", side_effect=fake_call_llm), \
+             patch("agent.moa_loop._maybe_apply_moa_cache_control", side_effect=lambda msgs, rt, **kwargs: msgs):
+            _run_reference(slot, [{"role": "user", "content": "hi"}], max_tokens=2000)
+
+        assert captured_kwargs.get("max_tokens") == 600
+
+    def test_slot_max_tokens_absent_falls_back_to_preset(self):
+        """When slot has no max_tokens, the preset-level cap is used."""
+        from agent.moa_loop import _run_reference
+
+        captured_kwargs: dict = {}
+
+        def fake_call_llm(**kwargs):
+            captured_kwargs.update(kwargs)
+            mock_resp = MagicMock()
+            mock_resp.choices = [MagicMock(message=MagicMock(content="advice"))]
+            mock_resp.usage = None
+            return mock_resp
+
+        slot = {"provider": "openrouter", "model": "deepseek/deepseek-v4-pro"}
+
+        with patch("agent.moa_loop._slot_runtime", return_value={"provider": "openrouter", "model": "deepseek/deepseek-v4-pro"}), \
+             patch("agent.moa_loop.call_llm", side_effect=fake_call_llm), \
+             patch("agent.moa_loop._maybe_apply_moa_cache_control", side_effect=lambda msgs, rt, **kwargs: msgs):
+            _run_reference(slot, [{"role": "user", "content": "hi"}], max_tokens=2000)
+
+        assert captured_kwargs.get("max_tokens") == 2000
+
+    def test_both_none_means_uncapped(self):
+        """When neither slot nor preset has max_tokens, it's None (uncapped)."""
+        from agent.moa_loop import _run_reference
+
+        captured_kwargs: dict = {}
+
+        def fake_call_llm(**kwargs):
+            captured_kwargs.update(kwargs)
+            mock_resp = MagicMock()
+            mock_resp.choices = [MagicMock(message=MagicMock(content="advice"))]
+            mock_resp.usage = None
+            return mock_resp
+
+        slot = {"provider": "openrouter", "model": "deepseek/deepseek-v4-pro"}
+
+        with patch("agent.moa_loop._slot_runtime", return_value={"provider": "openrouter", "model": "deepseek/deepseek-v4-pro"}), \
+             patch("agent.moa_loop.call_llm", side_effect=fake_call_llm), \
+             patch("agent.moa_loop._maybe_apply_moa_cache_control", side_effect=lambda msgs, rt, **kwargs: msgs):
+            _run_reference(slot, [{"role": "user", "content": "hi"}], max_tokens=None)
+
+        assert captured_kwargs.get("max_tokens") is None