script-build-and-framework-tmp/agent/tests/test_context_budget.py

import math

import pytest

from agent.core.runner import AgentRunner, RunConfig
from agent.trace.compaction import (
    CompressionConfig,
    calibrate_prompt_estimate,
    measure_prompt_tokens,
)
from agent.trace.models import Trace
from agent.trace.store import FileSystemTraceStore


def _knowledge_off():
    from agent.tools.builtin.knowledge import KnowledgeConfig

    return KnowledgeConfig(
        enable_extraction=False,
        enable_completion_extraction=False,
        enable_injection=False,
    )


def test_prompt_measurement_includes_tools_and_provider_calibration():
    config = CompressionConfig(max_tokens=100_000)
    messages = [{"role": "user", "content": "x" * 4_000}]
    schemas = [{
        "type": "function",
        "function": {
            "name": "large_tool",
            "description": "y" * 4_000,
            "parameters": {"type": "object"},
        },
    }]

    raw = measure_prompt_tokens(messages, [], config, "unknown-model", 1.0)
    measured = measure_prompt_tokens(messages, schemas, config, "unknown-model", 1.0)
    factor = calibrate_prompt_estimate(
        actual_prompt_tokens=measured.estimated_tokens * 2,
        estimated_prompt_tokens=measured.estimated_tokens,
        previous_factor=1.25,
    )
    calibrated = measure_prompt_tokens(
        messages,
        schemas,
        config,
        "unknown-model",
        factor,
    )

    assert measured.tool_schema_tokens > 0
    assert measured.estimated_tokens > raw.estimated_tokens
    assert factor == 2
    assert calibrated.calibrated_tokens == measured.estimated_tokens * 2
    assert calibrated.trigger_tokens == 80_000
    assert calibrated.target_tokens == 30_000


@pytest.mark.asyncio
async def test_build_900036_growth_replay_compresses_before_hard_limit(tmp_path):
    store = FileSystemTraceStore(str(tmp_path))
    trace = Trace(trace_id="growth-replay", mode="agent", agent_role="legacy")
    await store.create_trace(trace)
    runner = AgentRunner(trace_store=store)
    config = RunConfig(
        model="unregistered-build-model",
        compression=CompressionConfig(max_tokens=100_000),
        knowledge=_knowledge_off(),
    )
    schemas = [{
        "type": "function",
        "function": {
            "name": "offline_step",
            "description": "append one deterministic fixture step",
            "parameters": {"type": "object", "properties": {}},
        },
    }]
    history = [
        {"role": "system", "content": "keep system policy"},
        {"role": "user", "content": "offline 97-round context fixture"},
    ]
    total_prompt_tokens = 0
    max_prompt_tokens = 0

    for turn in range(97):
        history.extend([
            {"role": "assistant", "content": f"turn {turn}", "tool_calls": []},
            {"role": "tool", "content": "x" * 4_000},
        ])
        history, _, _, needs_compression = await runner._manage_context_usage(
            trace.trace_id,
            history,
            None,
            config,
            sequence=turn + 1,
            head_seq=0,
            tool_schemas=schemas,
        )
        current = await store.get_trace(trace.trace_id)
        measurement = measure_prompt_tokens(
            history,
            schemas,
            config.compression,
            config.model,
            current.runtime_state.get("context_budget", {}).get("calibration_factor"),
        )
        simulated_actual = math.ceil(measurement.estimated_tokens * 1.10)
        total_prompt_tokens += simulated_actual
        max_prompt_tokens = max(max_prompt_tokens, simulated_actual)
        await runner._record_prompt_measurement(
            current,
            config,
            history,
            schemas,
            simulated_actual,
        )

        if needs_compression:
            history = [
                history[0],
                history[1],
                {"role": "user", "content": "compact execution summary"},
            ]
            current = await store.get_trace(trace.trace_id)
            compacted = measure_prompt_tokens(
                history,
                schemas,
                config.compression,
                config.model,
                current.runtime_state["context_budget"]["calibration_factor"],
            )
            failure = await runner._finish_context_compression(
                current,
                compacted,
                before_message_count=turn * 2 + 4,
                after_message_count=len(history),
            )
            assert failure is None
            assert compacted.calibrated_tokens <= compacted.target_tokens

    events = await store.get_events(trace.trace_id)
    event_names = [event["event"] for event in events]
    assert "context_compression_started" in event_names
    assert "context_compression_completed" in event_names
    assert max_prompt_tokens < 100_000
    assert total_prompt_tokens < 6_327_689 * 0.50
    usage = runner.get_context_usage(trace.trace_id)
    assert usage.compression_count >= 1
    assert usage.actual_prompt_tokens is not None


@pytest.mark.asyncio
async def test_oversized_minimum_context_fails_before_model_request(tmp_path):
    calls = 0

    async def llm_call(**_kwargs):
        nonlocal calls
        calls += 1
        return {"content": "must not be called", "tool_calls": None}

    store = FileSystemTraceStore(str(tmp_path))
    result = await AgentRunner(
        trace_store=store,
        llm_call=llm_call,
    ).run_result(
        [{"role": "user", "content": "x" * 8_000}],
        RunConfig(
            max_iterations=4,
            tools=[],
            tool_groups=[],
            knowledge=_knowledge_off(),
            compression=CompressionConfig(max_tokens=1_000),
        ),
    )

    assert calls == 0
    assert result["status"] == "failed"
    assert result["failure"]["code"] == "CONTEXT_BUDGET_EXCEEDED"
    events = await store.get_events(result["trace_id"])
    assert "context_budget_exceeded" in [event["event"] for event in events]


@pytest.mark.parametrize(
    "kwargs",
    [
        {"trigger_ratio": 0.5, "target_ratio": 0.5},
        {"trigger_ratio": 1.0},
        {"fallback_safety_factor": 0.9},
        {"max_tokens": -1},
    ],
)
def test_compression_config_rejects_unsafe_budgets(kwargs):
    with pytest.raises(ValueError):
        CompressionConfig(**kwargs)