chopratejas--headroom/tests/test_custom_base_passthrough_telemetry.py

from __future__ import annotations

import asyncio
import json
from types import SimpleNamespace
from typing import Any

import httpx
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from fastapi.testclient import TestClient

from headroom.providers.proxy_routes import register_provider_routes
from headroom.proxy.handlers.openai import OpenAIHandlerMixin


class _Runtime:
    @staticmethod
    def api_target(provider: str) -> str:
        return f"https://{provider}.example.test"

    @staticmethod
    def model_metadata_provider(headers: dict[str, str]) -> str:
        return "anthropic"


class _Proxy:
    ANTHROPIC_API_URL = "https://anthropic.example.test"
    OPENAI_API_URL = "https://openai.example.test"
    GEMINI_API_URL = "https://gemini.example.test"
    CLOUDCODE_API_URL = "https://cloudcode.example.test"
    VERTEX_API_URL = "https://vertex.example.test"

    def __init__(self) -> None:
        self.config = SimpleNamespace(bedrock_api_url=None)
        self.provider_runtime = _Runtime()
        self.calls: list[dict[str, Any]] = []

    async def handle_passthrough(
        self,
        request: Any,
        base_url: str,
        endpoint_name: str = "",
        provider: str = "",
    ) -> JSONResponse:
        self.calls.append(
            {
                "path": request.url.path,
                "base_url": base_url,
                "endpoint_name": endpoint_name,
                "provider": provider,
            }
        )
        return JSONResponse(self.calls[-1])


class _ChatCompletionsRequest:
    method = "POST"
    headers = {}
    url = SimpleNamespace(path="/zen/v1/chat/completions", query="")

    async def body(self) -> bytes:
        return b'{"model":"zen"}'


class _OpenAIUsageClient:
    def __init__(self) -> None:
        self.calls: list[dict[str, Any]] = []

    async def request(self, **kwargs: Any) -> httpx.Response:
        self.calls.append(kwargs)
        request = httpx.Request(kwargs["method"], kwargs["url"])
        return httpx.Response(
            200,
            request=request,
            headers={"content-type": "application/json"},
            json={
                "usage": {
                    "prompt_tokens": 21,
                    "completion_tokens": 8,
                    "prompt_tokens_details": {"cached_tokens": 5},
                }
            },
        )


def test_custom_base_provider_prefixed_chat_completions_gets_telemetry() -> None:
    app = FastAPI()
    proxy = _Proxy()
    register_provider_routes(app, proxy)

    with TestClient(app) as client:
        for base_url, expected_base_url in (
            ("https://opencode.ai/", "https://opencode.ai"),
            ("https://www.opencode.ai/", "https://www.opencode.ai"),
        ):
            response = client.post(
                "/zen/v1/chat/completions",
                headers={"x-headroom-base-url": base_url},
                json={"model": "zen"},
            )
            assert response.status_code == 200
            assert response.json() == {
                "path": "/zen/v1/chat/completions",
                "base_url": expected_base_url,
                "endpoint_name": "chat/completions",
                "provider": "zen",
            }


def test_custom_base_unrelated_passthrough_paths_stay_unclassified() -> None:
    app = FastAPI()
    proxy = _Proxy()
    register_provider_routes(app, proxy)

    with TestClient(app) as client:
        for path in (
            "/mcp",
            "/mcp/v1/chat/completions",
            "/npm/v1/chat/completions",
            "/context7/v1/chat/completions",
        ):
            response = client.post(
                path,
                headers={"x-headroom-base-url": "https://opencode.ai/"},
                json={},
            )
            assert response.status_code == 200
            assert response.json() == {
                "path": path,
                "base_url": "https://opencode.ai",
                "endpoint_name": "",
                "provider": "",
            }


def test_custom_base_chat_completions_telemetry_is_post_and_opencode_zen_only() -> None:
    app = FastAPI()
    proxy = _Proxy()
    register_provider_routes(app, proxy)

    with TestClient(app) as client:
        get_response = client.get(
            "/zen/v1/chat/completions",
            headers={"x-headroom-base-url": "https://opencode.ai/"},
        )
        other_host_response = client.post(
            "/zen/v1/chat/completions",
            headers={"x-headroom-base-url": "https://custom.example/"},
            json={"model": "zen"},
        )
        double_slash_response = client.post(
            "/zen//v1/chat/completions",
            headers={"x-headroom-base-url": "https://opencode.ai/"},
            json={"model": "zen"},
        )
        trailing_slash_response = client.post(
            "/zen/v1/chat/completions/",
            headers={"x-headroom-base-url": "https://opencode.ai/"},
            json={"model": "zen"},
        )

    for response in (
        get_response,
        other_host_response,
        double_slash_response,
        trailing_slash_response,
    ):
        assert response.status_code == 200
        assert response.json()["endpoint_name"] == ""
        assert response.json()["provider"] == ""


def test_classified_custom_base_passthrough_records_telemetry_usage() -> None:
    handler = object.__new__(OpenAIHandlerMixin)
    handler.http_client = _OpenAIUsageClient()
    outcomes = []

    async def next_request_id() -> str:
        return "req_zen"

    async def record(outcome: Any) -> None:
        outcomes.append(outcome)

    handler._next_request_id = next_request_id
    handler._record_request_outcome = record

    response = asyncio.run(
        handler.handle_passthrough(
            _ChatCompletionsRequest(),
            "https://opencode.ai",
            "chat/completions",
            "zen",
        )
    )

    assert response.status_code == 200
    assert json.loads(response.body) == {
        "usage": {
            "prompt_tokens": 21,
            "completion_tokens": 8,
            "prompt_tokens_details": {"cached_tokens": 5},
        }
    }
    assert handler.http_client.calls[0]["url"] == ("https://opencode.ai/zen/v1/chat/completions")
    assert len(outcomes) == 1
    outcome = outcomes[0]
    assert outcome.provider == "zen"
    assert outcome.model == "passthrough:chat/completions"
    assert outcome.optimized_tokens == 21
    assert outcome.output_tokens == 8
    assert outcome.cache_read_tokens == 5