band-aider-pub/tests/band_aider/test_token_jacker.py

# SPDX-FileCopyrightText: 2026 Art Sea LLC
#
# SPDX-License-Identifier: AGPL-3.0-only

from __future__ import annotations

import logging

from band_aider.ai_client.completion_client import CompletionClient, OpenAIConfig
from band_aider.ai_client.token_jacker import TokenJacker
from band_aider.ai_client.cached_completion_client import CachedCompletionClient
from band_aider.context.token_budget_state import TokenBudgetSettings
from band_aider.ui.dto.model_escalation import ModelEscalationDTO


class _RecordingClient(CompletionClient):
    """_RecordingClient records prompts and returns configured responses."""

    model = "default"
    api_base = "https://example.invalid"
    api_key = "key"
    temperature = 0
    max_tokens = None
    max_attempts = 8
    retry_temperature_step = 0.2
    max_retry_temperature = 0.8

    def __init__(self, response: str, openai_configs: tuple[OpenAIConfig, ...] = ()) -> None:
        self.response = response
        self._openai_configs = openai_configs
        self.prompts: list[str] = []

    def complete(self, prompt: str, logger: logging.Logger | None = None) -> str:
        """Record the prompt and return the stationary configured response."""

        del logger
        self.prompts.append(prompt)
        return self.response

    @property
    def openai_configs(self) -> tuple[OpenAIConfig, ...]:
        return self._openai_configs

    @property
    def inner(self) -> CompletionClient | None:
        return None


class _Publisher:
    def __init__(self) -> None:
        self.events: list[object] = []

    def publish(self, event: object) -> None:
        self.events.append(event)


def test_token_jacker_forwards_under_budget_prompt_to_inner() -> None:
    """Verify under-budget prompts keep using the primary client."""

    inner = _RecordingClient("inner")
    client = TokenJacker(
        inner=inner,
        settings=TokenBudgetSettings(max_prompt_tokens=10_000),
    )

    assert client.complete("small prompt") == "inner"
    assert inner.prompts == ["small prompt"]


def test_token_jacker_routes_oversized_prompt_to_overflow() -> None:
    """Verify oversized prompts use the overflow client instead of reduction."""

    inner = _RecordingClient(
        "inner",
        openai_configs=(
            OpenAIConfig(model="local", api_base="http://local"),
            OpenAIConfig(model="cloud", api_base="http://cloud"),
        ),
    )
    overflow = _RecordingClient("overflow")
    publisher = _Publisher()
    client = TokenJacker(
        inner=inner,
        settings=TokenBudgetSettings(max_prompt_tokens=1, tokens_per_word=1.0),
        event_publisher=publisher,
        escalation_clients=(overflow,),
    )

    assert client.complete("too long") == "overflow"
    assert inner.prompts == []
    assert overflow.prompts == ["too long"]
    escalations = [event for event in publisher.events if isinstance(event, ModelEscalationDTO)]
    assert len(escalations) == 1
    assert escalations[0].trigger_name == "TokenJacker"
    assert escalations[0].active_model == "cloud"
    assert escalations[0].state == "routed"
    assert escalations[0].prompt_preview == "too long"


def test_token_jacker_escalates_past_cached_primary_response() -> None:
    inner = _RecordingClient("primary")
    overflow = _RecordingClient("overflow")
    primary_cached = CachedCompletionClient(inner)
    overflow_cached = CachedCompletionClient(overflow)
    primary_cached._cache.put("too long", "cached-primary")
    client = TokenJacker(
        inner=primary_cached,
        settings=TokenBudgetSettings(max_prompt_tokens=1, tokens_per_word=1.0),
        escalation_clients=(overflow_cached,),
    )

    assert client.complete("too long") == "overflow"
    assert inner.prompts == []
    assert overflow.prompts == ["too long"]


def test_token_jacker_walks_multiple_overflow_tiers() -> None:
    """Verify overflow routing escalates past one failed tier to the next."""

    inner = _RecordingClient(
        "inner",
        openai_configs=(
            OpenAIConfig(model="local", api_base="http://local"),
            OpenAIConfig(model="cloud-1", api_base="http://cloud-1"),
            OpenAIConfig(model="cloud-2", api_base="http://cloud-2"),
        ),
    )
    seen_models: list[str] = []

    class _FailingClient(CompletionClient):
        def complete(self, prompt: str, logger: logging.Logger | None = None) -> str:
            del prompt, logger
            seen_models.append("cloud-1")
            raise RuntimeError("overflow tier failed")

        @property
        def openai_configs(self) -> tuple[OpenAIConfig, ...]:
            return ()

        @property
        def inner(self) -> CompletionClient | None:
            return None

    class _SuccessClient(CompletionClient):
        def complete(self, prompt: str, logger: logging.Logger | None = None) -> str:
            del prompt, logger
            seen_models.append("cloud-2")
            return "overflow-2"

        @property
        def openai_configs(self) -> tuple[OpenAIConfig, ...]:
            return ()

        @property
        def inner(self) -> CompletionClient | None:
            return None

    client = TokenJacker(
        inner=inner,
        settings=TokenBudgetSettings(max_prompt_tokens=1, tokens_per_word=1.0),
        escalation_clients=(_FailingClient(), _SuccessClient()),
    )

    assert client.complete("too long") == "overflow-2"
    assert seen_models == ["cloud-1", "cloud-2"]


def test_token_jacker_oversized_prompt_falls_back_to_inner_without_overflow() -> None:
    """Verify oversized prompts log and fall back to the primary client when no overflow tier exists."""

    inner = _RecordingClient(
        "inner",
        openai_configs=(OpenAIConfig(model="local", api_base="http://local"),),
    )
    client = TokenJacker(
        inner=inner,
        settings=TokenBudgetSettings(max_prompt_tokens=1, tokens_per_word=1.0),
    )

    assert client.complete("too long") == "inner"
    assert inner.prompts == ["too long"]