headroom/tests/test_output_shaper.py

"""Tests for headroom.proxy.output_shaper.

Covers turn classification (structural only), cache-safe verbosity steering,
effort routing on mechanical continuations, and the env-driven gate.
"""

from __future__ import annotations

import copy
from typing import Any

from headroom.proxy.output_shaper import (
    LEGACY_THINKING_FLOOR,
    OutputShaperSettings,
    TurnKind,
    apply_verbosity_steering,
    classify_turn,
    route_effort,
    shape_request,
    steering_text,
)

ENABLED = OutputShaperSettings(enabled=True)


def _tool_result(is_error: bool = False) -> dict[str, Any]:
    block: dict[str, Any] = {
        "type": "tool_result",
        "tool_use_id": "toolu_01",
        "content": "ok",
    }
    if is_error:
        block["is_error"] = True
    return block


def _mechanical_messages() -> list[dict[str, Any]]:
    return [
        {"role": "user", "content": "fix the bug in foo.py"},
        {
            "role": "assistant",
            "content": [
                {"type": "text", "text": "Reading the file."},
                {"type": "tool_use", "id": "toolu_01", "name": "Read", "input": {}},
            ],
        },
        {"role": "user", "content": [_tool_result()]},
    ]


# ---------------------------------------------------------------------------
# classify_turn
# ---------------------------------------------------------------------------


class TestClassifyTurn:
    def test_string_user_message_is_new_ask(self):
        assert classify_turn([{"role": "user", "content": "explain this"}]) == TurnKind.NEW_USER_ASK

    def test_clean_tool_result_is_mechanical(self):
        assert classify_turn(_mechanical_messages()) == TurnKind.MECHANICAL_CONTINUATION

    def test_multiple_clean_tool_results_are_mechanical(self):
        msgs = _mechanical_messages()
        msgs[-1]["content"].append(_tool_result())
        assert classify_turn(msgs) == TurnKind.MECHANICAL_CONTINUATION

    def test_error_tool_result_is_error_continuation(self):
        msgs = _mechanical_messages()
        msgs[-1]["content"] = [_tool_result(), _tool_result(is_error=True)]
        assert classify_turn(msgs) == TurnKind.ERROR_CONTINUATION

    def test_text_block_alongside_tool_result_is_new_ask(self):
        msgs = _mechanical_messages()
        msgs[-1]["content"].append({"type": "text", "text": "also check bar.py"})
        assert classify_turn(msgs) == TurnKind.NEW_USER_ASK

    def test_image_block_is_new_ask(self):
        msgs = [{"role": "user", "content": [{"type": "image", "source": {}}]}]
        assert classify_turn(msgs) == TurnKind.NEW_USER_ASK

    def test_assistant_last_is_unknown(self):
        msgs = [{"role": "assistant", "content": "hello"}]
        assert classify_turn(msgs) == TurnKind.UNKNOWN

    def test_empty_messages_is_unknown(self):
        assert classify_turn([]) == TurnKind.UNKNOWN

    def test_empty_content_list_is_unknown(self):
        assert classify_turn([{"role": "user", "content": []}]) == TurnKind.UNKNOWN

    def test_whitespace_string_content_is_unknown(self):
        assert classify_turn([{"role": "user", "content": "  "}]) == TurnKind.UNKNOWN


# ---------------------------------------------------------------------------
# apply_verbosity_steering
# ---------------------------------------------------------------------------


class TestVerbositySteering:
    def test_level_zero_is_noop(self):
        body = {"system": "You are helpful."}
        assert apply_verbosity_steering(body, 0) is False
        assert body["system"] == "You are helpful."

    def test_string_system_converted_to_blocks_with_original_bytes_first(self):
        body = {"system": "You are helpful."}
        assert apply_verbosity_steering(body, 2) is True
        assert body["system"][0] == {"type": "text", "text": "You are helpful."}
        assert body["system"][1]["text"] == steering_text(2)

    def test_missing_system_creates_steering_only_block(self):
        body: dict[str, Any] = {}
        assert apply_verbosity_steering(body, 2) is True
        assert body["system"] == [{"type": "text", "text": steering_text(2)}]

    def test_block_system_appends_after_cache_control(self):
        cached = {
            "type": "text",
            "text": "Big system prompt.",
            "cache_control": {"type": "ephemeral"},
        }
        body = {"system": [copy.deepcopy(cached)]}
        assert apply_verbosity_steering(body, 2) is True
        # The cached block is byte-identical and still first — prefix intact.
        assert body["system"][0] == cached
        assert body["system"][1] == {"type": "text", "text": steering_text(2)}
        # Our block carries no cache_control (breakpoints are a scarce resource).
        assert "cache_control" not in body["system"][1]

    def test_idempotent_at_same_level(self):
        body = {"system": [{"type": "text", "text": "Sys."}]}
        assert apply_verbosity_steering(body, 2) is True
        snapshot = copy.deepcopy(body)
        assert apply_verbosity_steering(body, 2) is False
        assert body == snapshot

    def test_level_change_replaces_block_in_place(self):
        body = {"system": [{"type": "text", "text": "Sys."}]}
        apply_verbosity_steering(body, 2)
        assert apply_verbosity_steering(body, 4) is True
        steering_blocks = [
            b for b in body["system"] if b["text"].startswith("<headroom_output_shaping>")
        ]
        assert len(steering_blocks) == 1
        assert steering_blocks[0]["text"] == steering_text(4)

    def test_steering_text_is_deterministic(self):
        for level in (1, 2, 3, 4):
            assert steering_text(level) == steering_text(level)


# ---------------------------------------------------------------------------
# route_effort
# ---------------------------------------------------------------------------


class TestRouteEffort:
    def test_lowers_explicit_effort_on_mechanical_turn(self):
        body = {"output_config": {"effort": "xhigh"}}
        labels = route_effort(body, TurnKind.MECHANICAL_CONTINUATION, ENABLED)
        assert body["output_config"]["effort"] == "low"
        assert labels == ["output_shaper:effort:xhigh->low"]

    def test_never_injects_effort_when_absent(self):
        body: dict[str, Any] = {"messages": []}
        labels = route_effort(body, TurnKind.MECHANICAL_CONTINUATION, ENABLED)
        assert "output_config" not in body
        assert labels == []

    def test_effort_untouched_on_new_ask(self):
        body = {"output_config": {"effort": "xhigh"}}
        assert route_effort(body, TurnKind.NEW_USER_ASK, ENABLED) == []
        assert body["output_config"]["effort"] == "xhigh"

    def test_effort_untouched_on_error_continuation(self):
        body = {"output_config": {"effort": "xhigh"}}
        assert route_effort(body, TurnKind.ERROR_CONTINUATION, ENABLED) == []
        assert body["output_config"]["effort"] == "xhigh"

    def test_effort_already_at_target_untouched(self):
        body = {"output_config": {"effort": "low"}}
        assert route_effort(body, TurnKind.MECHANICAL_CONTINUATION, ENABLED) == []

    def test_unknown_effort_value_untouched(self):
        body = {"output_config": {"effort": "turbo"}}
        assert route_effort(body, TurnKind.MECHANICAL_CONTINUATION, ENABLED) == []
        assert body["output_config"]["effort"] == "turbo"

    def test_configurable_mechanical_effort(self):
        settings = OutputShaperSettings(enabled=True, mechanical_effort="medium")
        body = {"output_config": {"effort": "xhigh"}}
        route_effort(body, TurnKind.MECHANICAL_CONTINUATION, settings)
        assert body["output_config"]["effort"] == "medium"

    def test_legacy_thinking_budget_clamped(self):
        body = {"thinking": {"type": "enabled", "budget_tokens": 32000}}
        labels = route_effort(body, TurnKind.MECHANICAL_CONTINUATION, ENABLED)
        assert body["thinking"]["budget_tokens"] == LEGACY_THINKING_FLOOR
        assert body["thinking"]["type"] == "enabled"  # never toggled
        assert labels == [f"output_shaper:thinking_budget:32000->{LEGACY_THINKING_FLOOR}"]

    def test_legacy_budget_at_floor_untouched(self):
        body = {"thinking": {"type": "enabled", "budget_tokens": LEGACY_THINKING_FLOOR}}
        assert route_effort(body, TurnKind.MECHANICAL_CONTINUATION, ENABLED) == []

    def test_adaptive_thinking_untouched(self):
        body = {"thinking": {"type": "adaptive"}}
        assert route_effort(body, TurnKind.MECHANICAL_CONTINUATION, ENABLED) == []
        assert body["thinking"] == {"type": "adaptive"}


# ---------------------------------------------------------------------------
# shape_request (end to end)
# ---------------------------------------------------------------------------


class TestShapeRequest:
    def test_disabled_is_noop(self):
        body = {
            "system": "Sys.",
            "messages": _mechanical_messages(),
            "output_config": {"effort": "xhigh"},
        }
        snapshot = copy.deepcopy(body)
        result = shape_request(body, OutputShaperSettings(enabled=False))
        assert result.changed is False
        assert body == snapshot

    def test_enabled_applies_steering_and_effort_routing(self):
        body = {
            "system": "Sys.",
            "messages": _mechanical_messages(),
            "output_config": {"effort": "xhigh"},
            "thinking": {"type": "adaptive"},
        }
        result = shape_request(body, ENABLED)
        assert result.changed is True
        assert result.labels == [
            "output_shaper:verbosity:L2",
            "output_shaper:effort:xhigh->low",
        ]
        assert body["output_config"]["effort"] == "low"
        assert body["system"][1]["text"] == steering_text(2)

    def test_new_ask_gets_steering_but_keeps_effort(self):
        body = {
            "system": "Sys.",
            "messages": [{"role": "user", "content": "design a cache layer"}],
            "output_config": {"effort": "xhigh"},
        }
        result = shape_request(body, ENABLED)
        assert result.labels == ["output_shaper:verbosity:L2"]
        assert body["output_config"]["effort"] == "xhigh"

    def test_second_pass_is_stable(self):
        body = {"system": "Sys.", "messages": _mechanical_messages()}
        shape_request(body, ENABLED)
        snapshot = copy.deepcopy(body)
        result = shape_request(body, ENABLED)
        assert result.changed is False
        assert body == snapshot

    def test_from_env_defaults_off(self, monkeypatch):
        monkeypatch.delenv("HEADROOM_OUTPUT_SHAPER", raising=False)
        assert OutputShaperSettings.from_env().enabled is False

    def test_from_env_enabled_with_overrides(self, monkeypatch):
        monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "1")
        monkeypatch.setenv("HEADROOM_VERBOSITY_LEVEL", "3")
        monkeypatch.setenv("HEADROOM_MECHANICAL_EFFORT", "medium")
        settings = OutputShaperSettings.from_env()
        assert settings.enabled is True
        assert settings.verbosity_level == 3
        assert settings.mechanical_effort == "medium"

    def test_from_env_clamps_bad_values(self, monkeypatch):
        monkeypatch.setenv("HEADROOM_OUTPUT_SHAPER", "true")
        monkeypatch.setenv("HEADROOM_VERBOSITY_LEVEL", "99")
        monkeypatch.setenv("HEADROOM_MECHANICAL_EFFORT", "bogus")
        settings = OutputShaperSettings.from_env()
        assert settings.verbosity_level == 4
        assert settings.mechanical_effort == "low"