litellm/tests/unit/proxy/agent_endpoints/test_kill_switch.py

from base64 import b64encode
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Final

import httpx
import pytest
from pydantic import ValidationError

from litellm.constants import REDACTED_BY_LITELM_STRING
from litellm.proxy.agent_endpoints.kill_switch import (
    build_kill_switch_request,
    fire_kill_switch,
    redact_kill_switch,
    restore_kill_switch,
)
from litellm.types.agents import AgentKillSwitchConfig


@dataclass(frozen=True, slots=True)
class _SentRequest:
    method: str
    url: str
    headers: Mapping[str, str]
    json: Mapping[str, object] | None
    timeout: float


class _RecordingClient:
    def __init__(self, respond: httpx.Response | httpx.HTTPError) -> None:
        self.sent: list[_SentRequest] = []  # mutable-ok: test double records calls
        self.follow_redirects: list[bool] = []  # mutable-ok: test double records calls
        self._respond: Final = respond

    def build_request(
        self,
        method: str,
        url: str,
        *,
        headers: Mapping[str, str],
        json: Mapping[str, object] | None,
        timeout: float,
    ) -> httpx.Request:
        self.sent.append(_SentRequest(method, url, headers, json, timeout))
        return httpx.Request(method, url, headers=dict(headers), json=json)

    async def send(self, request: httpx.Request, *, stream: bool, follow_redirects: bool) -> httpx.Response:
        self.follow_redirects.append(follow_redirects)
        if isinstance(self._respond, httpx.HTTPError):
            raise self._respond
        return self._respond


class _CountingStream(httpx.AsyncByteStream):
    def __init__(self, chunk: bytes, chunks: int) -> None:
        self.pulled: int = 0  # rebind-ok: test double counts reads
        self._chunk: Final = chunk
        self._chunks: Final = chunks

    async def __aiter__(self):
        for _ in range(self._chunks):
            self.pulled += 1  # rebind-ok: test double counts reads
            yield self._chunk


def _config(**overrides: object) -> AgentKillSwitchConfig:
    return AgentKillSwitchConfig.model_validate({"url": "https://ops.example.com/agents/kill", **overrides})


def test_request_carries_endpoint_method_query_params_headers_and_body() -> None:
    request: Final = build_kill_switch_request(
        _config(
            url="https://ops.example.com/kill?env=prod",
            method="PUT",
            query_params={"agent": "billing-bot", "reason": "manual stop"},
            headers={"X-Trace": "abc"},
            body={"action": "stop", "hard": True},
        )
    )

    assert request.method == "PUT"
    assert str(httpx.URL(request.url)) == "https://ops.example.com/kill?env=prod&agent=billing-bot&reason=manual+stop"
    assert dict(request.headers) == {"X-Trace": "abc"}
    assert request.json_body == {"action": "stop", "hard": True}


def test_request_defaults_to_post_with_no_body_and_untouched_url() -> None:
    request: Final = build_kill_switch_request(_config())

    assert (request.method, request.url, dict(request.headers), request.json_body) == (
        "POST",
        "https://ops.example.com/agents/kill",
        {},
        None,
    )


@pytest.mark.parametrize(
    ("auth", "expected_headers"),
    [
        ({"type": "bearer", "token": "tok-123"}, {"Authorization": "Bearer tok-123"}),
        ({"type": "api_key", "api_key": "k-456"}, {"x-api-key": "k-456"}),
        ({"type": "api_key", "header_name": "X-Ops-Key", "api_key": "k-456"}, {"X-Ops-Key": "k-456"}),
        (
            {"type": "basic", "username": "ops", "password": "pw:1"},
            {"Authorization": f"Basic {b64encode(b'ops:pw:1').decode()}"},
        ),
    ],
)
def test_auth_becomes_the_matching_request_header(auth: Mapping[str, object], expected_headers: dict[str, str]) -> None:
    request: Final = build_kill_switch_request(_config(auth=auth))

    assert dict(request.headers) == expected_headers


def test_auth_header_wins_over_a_conflicting_custom_header() -> None:
    request: Final = build_kill_switch_request(
        _config(headers={"Authorization": "stale", "X-Env": "prod"}, auth={"type": "bearer", "token": "fresh"})
    )

    assert dict(request.headers) == {"Authorization": "Bearer fresh", "X-Env": "prod"}


@pytest.mark.parametrize("url", ["ftp://ops.example.com/kill", "/relative/kill", "ops.example.com/kill", ""])
def test_config_rejects_non_http_urls(url: str) -> None:
    with pytest.raises(ValidationError, match="absolute http"):
        _config(url=url)


def test_config_rejects_unknown_auth_type_and_unknown_fields() -> None:
    with pytest.raises(ValidationError):
        _config(auth={"type": "hmac", "secret": "x"})
    with pytest.raises(ValidationError):
        _config(endpoint="https://typo.example.com")


@pytest.mark.parametrize(
    ("auth", "secret_field"),
    [
        ({"type": "bearer", "token": "tok-123"}, "token"),
        ({"type": "api_key", "header_name": "X-K", "api_key": "k-456"}, "api_key"),
        ({"type": "basic", "username": "ops", "password": "pw"}, "password"),
    ],
)
def test_redact_replaces_only_the_secret_and_restore_puts_it_back(auth: dict[str, str], secret_field: str) -> None:
    original: Final = _config(auth=auth)

    redacted: Final = redact_kill_switch(original)
    assert redacted is not None and redacted.auth is not None
    assert redacted.auth.model_dump() == {**auth, secret_field: REDACTED_BY_LITELM_STRING}
    assert original.auth is not None and original.auth.model_dump() == auth, "redact must not mutate its input"

    restored: Final = restore_kill_switch(redacted, original)
    assert restored == original


def test_restore_keeps_a_rotated_secret_and_never_stores_the_marker_itself() -> None:
    rotated: Final = _config(auth={"type": "bearer", "token": "new-token"})
    stored: Final = _config(auth={"type": "bearer", "token": "old-token"})
    assert restore_kill_switch(rotated, stored) == rotated
    assert restore_kill_switch(None, stored) is None

    marker_only: Final = _config(auth={"type": "bearer", "token": REDACTED_BY_LITELM_STRING})
    assert restore_kill_switch(marker_only, None) == _config(auth={"type": "bearer", "token": ""})


def test_restore_does_not_borrow_a_secret_from_a_different_auth_type() -> None:
    incoming: Final = _config(auth={"type": "bearer", "token": REDACTED_BY_LITELM_STRING})
    stored: Final = _config(auth={"type": "api_key", "api_key": "k-456"})

    assert restore_kill_switch(incoming, stored) == _config(auth={"type": "bearer", "token": ""})


def test_redact_passes_through_configs_without_auth() -> None:
    assert redact_kill_switch(None) is None
    plain: Final = _config(headers={"X-Env": "prod"})
    assert redact_kill_switch(plain) is plain


@pytest.mark.asyncio
async def test_fire_sends_exactly_the_built_request_and_reports_the_2xx_reply_without_the_query() -> None:
    client: Final = _RecordingClient(httpx.Response(202, text="stopping"))
    config: Final = _config(
        method="DELETE",
        query_params={"force": "1", "token": "qs-secret"},
        headers={"X-Env": "prod"},
        body={"agent": "billing-bot"},
        auth={"type": "bearer", "token": "tok-123"},
    )

    result: Final = await fire_kill_switch(agent_id="agent-1", config=config, http_client=client, timeout=3.5)

    assert client.sent == [
        _SentRequest(
            method="DELETE",
            url="https://ops.example.com/agents/kill?force=1&token=qs-secret",
            headers={"X-Env": "prod", "Authorization": "Bearer tok-123"},
            json={"agent": "billing-bot"},
            timeout=3.5,
        )
    ]
    assert client.follow_redirects == [False], "a redirecting webhook must not be followed to another host"
    assert result.succeeded is True
    assert result.model_dump() == {
        "agent_id": "agent-1",
        "url": "https://ops.example.com/agents/kill",
        "method": "DELETE",
        "status_code": 202,
        "response_body": "stopping",
        "error": None,
    }


@pytest.mark.asyncio
async def test_fire_reports_a_non_2xx_reply_as_failure_with_the_body() -> None:
    client: Final = _RecordingClient(httpx.Response(503, text="x" * 5000))

    result: Final = await fire_kill_switch(agent_id="agent-1", config=_config(), http_client=client)

    assert result.succeeded is False
    assert result.status_code == 503
    assert result.response_body == "x" * 2000
    assert result.error is None


@pytest.mark.asyncio
async def test_fire_stops_reading_the_body_at_the_cap_instead_of_buffering_the_whole_reply() -> None:
    stream: Final = _CountingStream(b"y" * 500, chunks=100)
    client: Final = _RecordingClient(httpx.Response(200, stream=stream))

    result: Final = await fire_kill_switch(agent_id="agent-1", config=_config(), http_client=client)

    assert result.response_body == "y" * 2000
    assert stream.pulled == 4, f"read {stream.pulled} of 100 chunks for a 2000 char cap"


@pytest.mark.asyncio
async def test_fire_reports_a_transport_error_by_type_without_raising_or_echoing_the_url() -> None:
    client: Final = _RecordingClient(httpx.ConnectError("boom https://ops.example.com/agents/kill?token=qs-secret"))

    result: Final = await fire_kill_switch(
        agent_id="agent-1", config=_config(query_params={"token": "qs-secret"}), http_client=client
    )

    assert result.succeeded is False
    assert (result.status_code, result.response_body) == (None, None)
    assert result.error == "ConnectError"
    assert "qs-secret" not in result.model_dump_json()