atlas-iac/testing/tests/test_suite_capacity.py

"""Conservative context reservation and post-reconciliation review batching."""
import json
import threading

import pytest

import suite_backends
import suite_multipass as workflow
from suite_cli_diagnostics import snapshot
from suite_contract import MODELS, Problem, encoded, preflight, validate_result
from test_suite_multipass import natural, public, request, review, wire_value


def call_with_bytes(size):
    """Construct a content-free pass at an exact serialized capacity boundary."""
    return {"stage": "decision_audit", "input": "x" * (size - 2), "system": "", "schema": {}}


def test_exact_historical_audit_reservation_fits_five_turns():
    result = workflow.capacity(call_with_bytes(638566), "claude", 363, 45)
    assert result["input_token_bound"] == 646758
    assert 646758 + 6 * 64000 == 1030758
    assert result["max_turns"] == 5 and result["configured_max_turns"] == 6
    assert result["context_reserved_tokens"] == 966758
    assert result["context_headroom_tokens"] == 33242
    assert result["output_reservation_tokens"] == 43920


@pytest.mark.parametrize("size,turns", [(607808, 6), (607809, 5), (671808, 5),
                                        (671809, 4), (735808, 4)])
def test_context_admission_keeps_full_outputs_and_bounded_repair(size, turns):
    result = workflow.capacity(call_with_bytes(size), "claude", 363, 45)
    assert result["max_turns"] == turns
    assert result["context_reserved_tokens"] <= MODELS["claude"]["context"]
    command = suite_backends.claude_command(MODELS["claude"]["model"], 30, max_turns=turns)
    assert command[command.index("--max-turns") + 1] == str(turns)


@pytest.mark.parametrize("size,count,reason", [(735809, 45, "context_reservation"),
                                              (638566, 200, "output_reservation")])
def test_capacity_failure_is_explicit_when_minimum_headroom_cannot_fit(size, count, reason):
    with pytest.raises(Problem, match="pass_capacity") as raised:
        workflow.capacity(call_with_bytes(size), "claude", 363, count)
    details = raised.value.details
    assert details["capacity_reason"] == reason and details["context_limit"] == 1000000
    assert details["max_output_tokens"] == 64000 and details["minimum_max_turns"] == 4
    assert "x" * 100 not in json.dumps(raised.value.document())


@pytest.mark.parametrize("turns", [0, 3, 7, True, 5.0])
def test_invalid_cli_turn_limits_never_launch(turns):
    with pytest.raises(Problem, match="unsupported_turn_limit"):
        suite_backends.claude_command(MODELS["claude"]["model"], 30, max_turns=turns)


def test_actual_turn_ceiling_and_message_measurements_are_content_free():
    raw = json.dumps({"type": "assistant", "message": {"usage": {
        "input_tokens": 1234, "output_tokens": 987, "private": "DO_NOT_RETAIN"},
        "content": [{"type": "text", "text": "DO_NOT_RETAIN"}]}})
    result = snapshot(raw, max_turns=5)
    assert result["max_turns"] == 5
    assert result["max_message_input_tokens"] == 1234 and result["max_message_output_tokens"] == 987
    assert "DO_NOT_RETAIN" not in json.dumps(result)
    assert snapshot("")["max_message_input_tokens"] is None


def install_review_fixture(monkeypatch, families=2, per_batch=7, collision=False):
    """Mock only model answers; run real batch planning, admission and validation."""
    source = request(families * 7)
    aliases = [c["alias"] for c in source["cases"]]
    original = natural(source, [aliases[i:i+7] for i in range(0, len(aliases), 7)])
    checks, calls = [], []
    actual_capacity = workflow.capacity

    def capacity(call, provider, count, family_count=None):
        if call["stage"] in {"large_family_review", "decision_audit"} and count > per_batch:
            raise Problem("pass_capacity", 422, review_pass=call["stage"], capacity_reason="context_reservation")
        checks.append((call["stage"], count))
        return actual_capacity(call, provider, count, family_count)

    def backend(value, cancel, *, invocation, progress=None, job_deadline=None):
        subset = {c["alias"] for c in value["cases"]}
        original_cases = {c["alias"]: c for c in source["cases"]}
        assert all(c == original_cases[c["alias"]] for c in value["cases"])
        assert {c["alias"] for c in json.loads(invocation["input"])["suite"]["cases"]} == subset
        stage = invocation["stage"]
        chosen = {"groups": [g for g in original["groups"] if set(g["members"]) <= subset]}
        if stage.startswith("proposal"):
            chosen = public(chosen)
        elif stage in {"large_family_review", "decision_audit"}:
            chosen = review(chosen, chosen)
            if collision and stage == "decision_audit":
                for group in chosen["groups"]:
                    group["name"] = "Colliding mechanism"
        calls.append({"stage": stage, "aliases": subset, "execution": value["execution"],
                      "deadline": job_deadline, "batch": json.loads(invocation["input"])["review_material"].get("review_batch")})
        return wire_value(chosen, invocation), {"cost_usd_estimate": .1, "usage": {"input_tokens": 10, "output_tokens": 10},
            "model": MODELS["claude"]["model"], "turns": 2, "duration_api_ms": 1,
            "compaction": False, "truncation": False, "cli_diagnostics": {"exit_code": 0}}

    monkeypatch.setattr(workflow, "capacity", capacity)
    monkeypatch.setattr(suite_backends, "claude_generate", backend)
    return source, calls


def test_automatic_review_batches_preserve_global_discovery_and_full_validation(monkeypatch):
    source, calls = install_review_fixture(monkeypatch)
    updates = []
    result, metadata = workflow.generate(source, preflight(source), threading.Event(), "192.168.22.8", updates.append)
    assert [c["stage"] for c in calls] == ["proposal_a", "proposal_b", "reconciliation",
        "large_family_review", "large_family_review", "decision_audit", "decision_audit"]
    assert all(len(c["aliases"]) == 14 for c in calls[:3])
    assert all(len(c["aliases"]) == 7 for c in calls[3:])
    assert calls[3]["aliases"].isdisjoint(calls[4]["aliases"])
    assert calls[3]["aliases"] == calls[5]["aliases"] and calls[4]["aliases"] == calls[6]["aliases"]
    assert len({c["deadline"] for c in calls}) == 1
    assert all(c['execution']['max_cost_usd'] is None for c in calls)
    assert updates[-1]["maximum_model_passes"] == 128
    assert metadata["model_pass_count"] == 7 and len(result["groups"]) == 4
    validate_result(result, source)


@pytest.mark.parametrize("families,per_batch,collision,code", [
    (2, 7, True, "duplicate_family_name"),
    (1, 6, False, "review_family_capacity"),
    (9, 7, False, "review_batch_limit"),
])
def test_batch_failure_never_accepts_partial_or_arbitrary_split(monkeypatch, families, per_batch, collision, code):
    source, calls = install_review_fixture(monkeypatch, families, per_batch, collision)
    updates = []
    with pytest.raises(Problem, match=code) as raised:
        workflow.generate(source, preflight(source), threading.Event(), "192.168.22.8", updates.append)
    stage = "decision_audit" if collision else "large_family_review"
    assert raised.value.details["review_pass"] == updates[-1]["current_pass"] == stage
    if not collision:
        assert len(calls) == 3
        assert updates[-1]["pass_state"] == "capacity_preflight"
        assert updates[-1]["cli_running"] is False