chopratejas--headroom/tests/test_proxy_pipeline_lifecycle.py

from __future__ import annotations

import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock, call, patch

import httpx
from fastapi.testclient import TestClient

from headroom.pipeline import PipelineStage
from headroom.proxy.server import ProxyConfig, create_app


class _RecordingExtension:
    def __init__(self) -> None:
        self.stages: list[PipelineStage] = []
        self.events: list = []

    def on_pipeline_event(self, event):
        self.stages.append(event.stage)
        self.events.append(event)
        return None


class _DummyTokenizer:
    def count_messages(self, messages):
        return len(messages)


def _assert_compressed_event_carries_originals(events: list) -> None:
    """INPUT_COMPRESSED must expose the pre-compression messages to extensions.

    The probe recorder (headroom.proxy.probe_recorder) depends on this
    metadata contract; dropping it silently disables session recording.
    """
    compressed = [event for event in events if event.stage is PipelineStage.INPUT_COMPRESSED]
    assert compressed
    original = compressed[0].metadata.get("original_messages")
    assert isinstance(original, list)
    assert any(
        message.get("role") == "user" and "hello" in str(message.get("content"))
        for message in original
        if isinstance(message, dict)
    )


def _assert_stage_order(stages: list[PipelineStage]) -> None:
    expected = [
        PipelineStage.SETUP,
        PipelineStage.PRE_START,
        PipelineStage.POST_START,
        PipelineStage.INPUT_RECEIVED,
        PipelineStage.INPUT_ROUTED,
        PipelineStage.INPUT_COMPRESSED,
        PipelineStage.INPUT_REMEMBERED,
        PipelineStage.PRE_SEND,
        PipelineStage.POST_SEND,
        PipelineStage.RESPONSE_RECEIVED,
    ]
    positions = [stages.index(stage) for stage in expected]
    assert positions == sorted(positions)


def test_proxy_shutdown_unloads_image_models() -> None:
    config = ProxyConfig(
        optimize=False,
        image_optimize=False,
        cache_enabled=False,
        rate_limit_enabled=False,
        cost_tracking_enabled=False,
        log_requests=False,
        ccr_inject_tool=False,
        ccr_handle_responses=False,
        ccr_context_tracking=False,
    )
    app = create_app(config)
    proxy = app.state.proxy
    proxy.http_client = None
    proxy.memory_handler = None

    quota_registry = SimpleNamespace(stop_all=AsyncMock())
    with (
        patch("headroom.proxy.server.get_quota_registry", return_value=quota_registry),
        patch("headroom.models.ml_models.MLModelRegistry.unload_prefix") as unload_prefix,
    ):
        asyncio.run(proxy.shutdown())

    assert unload_prefix.call_args_list == [
        call("technique_router:"),
        call("siglip:"),
    ]
    quota_registry.stop_all.assert_awaited_once()


def test_proxy_shutdown_flushes_savings_tracker() -> None:
    """Graceful shutdown must flush the savings tracker's batched tail.

    The proxy throttles savings persistence (save_flush_every=25), so buffered
    requests only reach disk on the next threshold write or an explicit flush.
    shutdown() is that flush; if the wiring regresses, a graceful stop silently
    drops the last few requests' lifetime totals. The tracker's flush() logic is
    covered in test_proxy_savings_history.py — this guards only the call site.
    """
    config = ProxyConfig(
        optimize=False,
        image_optimize=False,
        cache_enabled=False,
        rate_limit_enabled=False,
        cost_tracking_enabled=False,
        log_requests=False,
        ccr_inject_tool=False,
        ccr_handle_responses=False,
        ccr_context_tracking=False,
    )
    app = create_app(config)
    proxy = app.state.proxy
    proxy.http_client = None
    proxy.memory_handler = None
    proxy.metrics.savings_tracker.flush = Mock()

    quota_registry = SimpleNamespace(stop_all=AsyncMock())
    with (
        patch("headroom.proxy.server.get_quota_registry", return_value=quota_registry),
        patch("headroom.models.ml_models.MLModelRegistry.unload_prefix"),
    ):
        asyncio.run(proxy.shutdown())

    proxy.metrics.savings_tracker.flush.assert_called_once()


def test_proxy_shutdown_signals_retry_waiters() -> None:
    config = ProxyConfig(
        optimize=False,
        image_optimize=False,
        cache_enabled=False,
        rate_limit_enabled=False,
        cost_tracking_enabled=False,
        log_requests=False,
        ccr_inject_tool=False,
        ccr_handle_responses=False,
        ccr_context_tracking=False,
    )
    app = create_app(config)
    proxy = app.state.proxy
    proxy.http_client = None
    proxy.memory_handler = None
    proxy._shutdown_event = asyncio.Event()

    quota_registry = SimpleNamespace(stop_all=AsyncMock())
    with (
        patch("headroom.proxy.server.get_quota_registry", return_value=quota_registry),
        patch("headroom.models.ml_models.MLModelRegistry.unload_prefix"),
    ):
        asyncio.run(proxy.shutdown())

    assert proxy._shutdown_event.is_set()


def test_openai_chat_pipeline_events_cover_proxy_lifecycle(monkeypatch) -> None:
    recorder = _RecordingExtension()
    config = ProxyConfig(
        optimize=True,
        image_optimize=False,
        cache_enabled=False,
        rate_limit_enabled=False,
        cost_tracking_enabled=False,
        log_requests=False,
        ccr_inject_tool=False,
        ccr_handle_responses=False,
        ccr_context_tracking=False,
        pipeline_extensions=[recorder],
        discover_pipeline_extensions=False,
    )
    app = create_app(config)

    with TestClient(app) as client:
        proxy = client.app.state.proxy
        proxy.openai_pipeline = SimpleNamespace(
            apply=lambda messages, model, **kwargs: SimpleNamespace(
                messages=[
                    {"role": "system", "content": "memory"},
                    {"role": "user", "content": "hello"},
                ],
                transforms_applied=["router:text:kompress"],
                tokens_before=10,
                tokens_after=6,
            )
        )
        proxy.memory_handler = SimpleNamespace(
            config=SimpleNamespace(inject_context=True, inject_tools=False),
            search_and_format_context=AsyncMock(return_value="memory"),
            has_memory_tool_calls=lambda response, provider: False,
        )

        monkeypatch.setattr("headroom.tokenizers.get_tokenizer", lambda model: _DummyTokenizer())

        async def _fake_retry(method, url, headers, body, stream=False, **kwargs):  # noqa: ANN001
            return httpx.Response(
                200,
                json={
                    "id": "chatcmpl_1",
                    "object": "chat.completion",
                    "choices": [{"message": {"role": "assistant", "content": "ok"}}],
                    "usage": {"prompt_tokens": 10, "completion_tokens": 3, "total_tokens": 13},
                },
            )

        proxy._retry_request = _fake_retry

        response = client.post(
            "/v1/chat/completions",
            headers={"Authorization": "Bearer sk-test", "x-headroom-user-id": "user-1"},
            json={
                "model": "gpt-5.4",
                "messages": [{"role": "user", "content": "hello"}],
                "tools": [{"type": "function", "function": {"name": "tool_a"}}],
            },
        )

    assert response.status_code == 200
    _assert_stage_order(recorder.stages)
    _assert_compressed_event_carries_originals(recorder.events)


def test_anthropic_messages_pipeline_events_cover_proxy_lifecycle(monkeypatch) -> None:
    recorder = _RecordingExtension()
    config = ProxyConfig(
        optimize=True,
        cache_enabled=False,
        rate_limit_enabled=False,
        cost_tracking_enabled=False,
        log_requests=False,
        ccr_inject_tool=False,
        ccr_handle_responses=False,
        ccr_context_tracking=False,
        image_optimize=False,
        pipeline_extensions=[recorder],
        discover_pipeline_extensions=False,
    )
    app = create_app(config)

    with TestClient(app) as client:
        proxy = client.app.state.proxy
        proxy.anthropic_pipeline = SimpleNamespace(
            apply=lambda messages, model, **kwargs: SimpleNamespace(
                messages=[
                    {"role": "system", "content": "memory"},
                    {"role": "user", "content": "hello"},
                ],
                transforms_applied=["router:text:kompress"],
                tokens_before=10,
                tokens_after=6,
            )
        )
        proxy.memory_handler = SimpleNamespace(
            config=SimpleNamespace(inject_context=True, inject_tools=False),
            search_and_format_context=AsyncMock(return_value="memory"),
            has_memory_tool_calls=lambda response, provider: False,
        )

        monkeypatch.setattr("headroom.tokenizers.get_tokenizer", lambda model: _DummyTokenizer())

        async def _fake_retry(method, url, headers, body, stream=False, **kwargs):  # noqa: ANN001
            return httpx.Response(
                200,
                json={
                    "id": "msg_1",
                    "type": "message",
                    "role": "assistant",
                    "content": [{"type": "text", "text": "ok"}],
                    "usage": {
                        "input_tokens": 10,
                        "output_tokens": 3,
                        "cache_read_input_tokens": 0,
                        "cache_creation_input_tokens": 0,
                    },
                },
            )

        proxy._retry_request = _fake_retry

        response = client.post(
            "/v1/messages",
            headers={
                "x-api-key": "test-key",
                "anthropic-version": "2023-06-01",
                "x-headroom-user-id": "user-1",
            },
            json={
                "model": "claude-sonnet-4-6",
                "max_tokens": 128,
                "messages": [{"role": "user", "content": "hello"}],
                "tools": [
                    {"name": "tool_a", "description": "a", "input_schema": {"type": "object"}}
                ],
            },
        )

    assert response.status_code == 200
    _assert_stage_order(recorder.stages)
    _assert_compressed_event_carries_originals(recorder.events)