|
| 1 | +"""Cost safety: per-role output caps + recursion limit. |
| 2 | +
|
| 3 | +Context: a user burned $10 in a minute running Claude Sonnet 4.5 because |
| 4 | +(a) `make_model` didn't set `max_tokens`, so replies went up to 8K, and |
| 5 | +(b) the agent had no `recursion_limit`, so it looped tens of times on |
| 6 | +one source. These tests lock in the caps that prevent that. |
| 7 | +""" |
| 8 | + |
| 9 | +from __future__ import annotations |
| 10 | + |
| 11 | +import pytest |
| 12 | + |
| 13 | +from thesis_agent.agent import _DEFAULT_RECURSION_LIMIT, _recursion_limit |
| 14 | +from thesis_agent.config import _role_max_tokens, make_model |
| 15 | + |
| 16 | + |
| 17 | +@pytest.fixture |
| 18 | +def clean_env(monkeypatch): |
| 19 | + for v in ( |
| 20 | + "THESIS_MAX_TOKENS_DRAFTER", |
| 21 | + "THESIS_MAX_TOKENS_CURATOR", |
| 22 | + "THESIS_MAX_TOKENS_RESEARCHER", |
| 23 | + "THESIS_MAX_TOKENS_DEFAULT", |
| 24 | + "THESIS_RECURSION_LIMIT", |
| 25 | + "THESIS_PROVIDER", |
| 26 | + "OPENROUTER_API_KEY", |
| 27 | + "ANTHROPIC_API_KEY", |
| 28 | + ): |
| 29 | + monkeypatch.delenv(v, raising=False) |
| 30 | + |
| 31 | + |
| 32 | +# --------------------------------------------------------------------------- |
| 33 | +# _role_max_tokens |
| 34 | +# --------------------------------------------------------------------------- |
| 35 | + |
| 36 | +class TestRoleMaxTokens: |
| 37 | + def test_default_caps_are_conservative(self, clean_env): |
| 38 | + assert _role_max_tokens("drafter") == 4000 |
| 39 | + assert _role_max_tokens("curator") == 2000 |
| 40 | + assert _role_max_tokens("researcher") == 1000 |
| 41 | + # Unknown role / default falls back to 2000 (not Sonnet's 8192 ceiling) |
| 42 | + assert _role_max_tokens(None) == 2000 |
| 43 | + assert _role_max_tokens("unknown-role") == 2000 |
| 44 | + |
| 45 | + def test_env_override_per_role(self, clean_env, monkeypatch): |
| 46 | + monkeypatch.setenv("THESIS_MAX_TOKENS_DRAFTER", "8000") |
| 47 | + monkeypatch.setenv("THESIS_MAX_TOKENS_CURATOR", "3000") |
| 48 | + assert _role_max_tokens("drafter") == 8000 |
| 49 | + assert _role_max_tokens("curator") == 3000 |
| 50 | + # Roles without an override still use defaults |
| 51 | + assert _role_max_tokens("researcher") == 1000 |
| 52 | + |
| 53 | + def test_non_digit_env_is_ignored(self, clean_env, monkeypatch): |
| 54 | + monkeypatch.setenv("THESIS_MAX_TOKENS_DRAFTER", "lots") |
| 55 | + assert _role_max_tokens("drafter") == 4000 |
| 56 | + |
| 57 | + def test_default_env_override_affects_unknown_role(self, clean_env, monkeypatch): |
| 58 | + monkeypatch.setenv("THESIS_MAX_TOKENS_DEFAULT", "512") |
| 59 | + # Unknown roles should pick up the DEFAULT override path. Our helper |
| 60 | + # falls back to the hardcoded 2000 for unknowns, BUT if the user sets |
| 61 | + # THESIS_MAX_TOKENS_DEFAULT we honour it for the `None` role. |
| 62 | + assert _role_max_tokens(None) == 512 |
| 63 | + |
| 64 | + |
| 65 | +# --------------------------------------------------------------------------- |
| 66 | +# make_model passes max_tokens through |
| 67 | +# --------------------------------------------------------------------------- |
| 68 | + |
| 69 | +class TestMakeModelCaps: |
| 70 | + def test_anthropic_model_receives_max_tokens(self, clean_env, monkeypatch): |
| 71 | + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-dummy") |
| 72 | + obj = make_model("anthropic:claude-haiku-4-5-20251001", role="curator") |
| 73 | + # langchain-anthropic stores the cap as .max_tokens |
| 74 | + assert getattr(obj, "max_tokens", None) == 2000 |
| 75 | + |
| 76 | + def test_openrouter_model_receives_max_tokens(self, clean_env, monkeypatch): |
| 77 | + monkeypatch.setenv("THESIS_PROVIDER", "openrouter") |
| 78 | + monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-dummy") |
| 79 | + obj = make_model("z-ai/glm-5.1", role="drafter") |
| 80 | + cap = getattr(obj, "max_tokens", None) or getattr(obj, "max_completion_tokens", None) |
| 81 | + assert cap == 4000 |
| 82 | + |
| 83 | + def test_explicit_max_tokens_wins_over_role_default(self, clean_env, monkeypatch): |
| 84 | + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-dummy") |
| 85 | + obj = make_model( |
| 86 | + "anthropic:claude-haiku-4-5-20251001", |
| 87 | + role="drafter", |
| 88 | + max_tokens=123, |
| 89 | + ) |
| 90 | + assert getattr(obj, "max_tokens", None) == 123 |
| 91 | + |
| 92 | + |
| 93 | +# --------------------------------------------------------------------------- |
| 94 | +# _recursion_limit |
| 95 | +# --------------------------------------------------------------------------- |
| 96 | + |
| 97 | +class TestRecursionLimit: |
| 98 | + def test_default_is_conservative(self, clean_env): |
| 99 | + assert _recursion_limit() == _DEFAULT_RECURSION_LIMIT |
| 100 | + assert _DEFAULT_RECURSION_LIMIT <= 50 # regression guard — don't raise silently |
| 101 | + |
| 102 | + def test_env_override(self, clean_env, monkeypatch): |
| 103 | + monkeypatch.setenv("THESIS_RECURSION_LIMIT", "40") |
| 104 | + assert _recursion_limit() == 40 |
| 105 | + |
| 106 | + def test_env_override_has_lower_bound(self, clean_env, monkeypatch): |
| 107 | + # Very low values would make the agent unusable — clamp to >= 5. |
| 108 | + monkeypatch.setenv("THESIS_RECURSION_LIMIT", "1") |
| 109 | + assert _recursion_limit() >= 5 |
| 110 | + |
| 111 | + def test_non_digit_env_is_ignored(self, clean_env, monkeypatch): |
| 112 | + monkeypatch.setenv("THESIS_RECURSION_LIMIT", "unlimited") |
| 113 | + assert _recursion_limit() == _DEFAULT_RECURSION_LIMIT |
| 114 | + |
| 115 | + |
| 116 | +# --------------------------------------------------------------------------- |
| 117 | +# invoke() config passes through recursion_limit |
| 118 | +# --------------------------------------------------------------------------- |
| 119 | + |
| 120 | +class TestInvokeConfig: |
| 121 | + def test_invoke_sets_recursion_limit_in_config(self, clean_env, monkeypatch, tmp_path): |
| 122 | + """Smoke test: our `invoke()` wrapper must attach `recursion_limit` |
| 123 | + to the config dict it passes to the underlying agent.""" |
| 124 | + from contextlib import contextmanager |
| 125 | + |
| 126 | + from thesis_agent import agent as agent_mod |
| 127 | + |
| 128 | + captured: dict = {} |
| 129 | + |
| 130 | + class _FakeAgent: |
| 131 | + def invoke(self, payload, *, config=None): |
| 132 | + captured.update(config or {}) |
| 133 | + return {"messages": [{"role": "assistant", "content": "ok"}]} |
| 134 | + |
| 135 | + @contextmanager |
| 136 | + def fake_build_agent(**kw): |
| 137 | + yield _FakeAgent() |
| 138 | + |
| 139 | + monkeypatch.chdir(tmp_path) |
| 140 | + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-dummy") |
| 141 | + monkeypatch.setattr(agent_mod, "build_agent", fake_build_agent) |
| 142 | + agent_mod.invoke("hi", thread_id="t-test") |
| 143 | + |
| 144 | + assert "recursion_limit" in captured |
| 145 | + assert captured["recursion_limit"] == _DEFAULT_RECURSION_LIMIT |
| 146 | + assert captured["configurable"]["thread_id"] == "t-test" |
| 147 | + |
| 148 | + def test_env_override_propagates_to_invoke(self, clean_env, monkeypatch, tmp_path): |
| 149 | + from contextlib import contextmanager |
| 150 | + |
| 151 | + from thesis_agent import agent as agent_mod |
| 152 | + |
| 153 | + captured: dict = {} |
| 154 | + |
| 155 | + class _FakeAgent: |
| 156 | + def invoke(self, payload, *, config=None): |
| 157 | + captured.update(config or {}) |
| 158 | + return {"messages": [{"role": "assistant", "content": "ok"}]} |
| 159 | + |
| 160 | + @contextmanager |
| 161 | + def fake_build_agent(**kw): |
| 162 | + yield _FakeAgent() |
| 163 | + |
| 164 | + monkeypatch.chdir(tmp_path) |
| 165 | + monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-ant-dummy") |
| 166 | + monkeypatch.setenv("THESIS_RECURSION_LIMIT", "7") |
| 167 | + monkeypatch.setattr(agent_mod, "build_agent", fake_build_agent) |
| 168 | + agent_mod.invoke("hi", thread_id="t-test") |
| 169 | + |
| 170 | + assert captured["recursion_limit"] == 7 |
| 171 | + |
| 172 | + |
| 173 | +# --------------------------------------------------------------------------- |
| 174 | +# Cost sanity check: worst-case per-call output bill |
| 175 | +# --------------------------------------------------------------------------- |
| 176 | + |
| 177 | +def test_worst_case_output_spend_per_call_is_bounded(clean_env): |
| 178 | + """At Sonnet 4.5 output prices ($15/M) a single drafter reply at its |
| 179 | + cap costs at most: 4000 tokens * $15 / 1_000_000 = $0.06. Sanity check |
| 180 | + the cap hasn't drifted into dangerous territory.""" |
| 181 | + price_per_million_output = 15.0 # Sonnet 4.5 worst case |
| 182 | + for role, expected_max in (("drafter", 4000), ("curator", 2000), ("researcher", 1000)): |
| 183 | + cap = _role_max_tokens(role) |
| 184 | + assert cap == expected_max |
| 185 | + max_cost = cap * price_per_million_output / 1_000_000 |
| 186 | + assert max_cost <= 0.10, ( |
| 187 | + f"{role} cap {cap} would cost ${max_cost:.3f} per reply — too high" |
| 188 | + ) |
0 commit comments