import json
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch

import httpx
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient

from litellm.constants import REDACTED_BY_LITELM_STRING
from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.agent_endpoints import endpoints as agent_endpoints
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
    RestrictedAgentAccess,
)
from litellm.proxy.agent_endpoints.endpoints import (
    _attach_keys_to_agents,
    _check_agent_management_permission,
    get_agent_daily_activity,
    router,
    user_api_key_auth,
)
from litellm.types.agents import AgentResponse


def _sample_agent_card_params() -> dict:
    return {
        "protocolVersion": "1.0",
        "name": "Test Agent",
        "description": "desc",
        "url": "http://localhost",
        "version": "1.0.0",
        "capabilities": {"streaming": True},
        "defaultInputModes": ["text"],
        "defaultOutputModes": ["text"],
        "skills": [],
    }


def _sample_agent_config() -> dict:
    return {
        "agent_name": "Test Agent",
        "agent_card_params": _sample_agent_card_params(),
        "litellm_params": {"make_public": False},
    }


def _sample_agent_response(
    agent_id: str = "agent-123", agent_name: str = "Test Agent"
) -> AgentResponse:
    return AgentResponse(
        agent_id=agent_id,
        agent_name=agent_name,
        agent_card_params=_sample_agent_card_params(),
        litellm_params={"make_public": False},
    )


def _make_app_with_role(role: LitellmUserRoles) -> TestClient:
    """Create a TestClient where the auth dependency returns the given role."""
    test_app = FastAPI()
    test_app.include_router(router)
    test_app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
        user_id="test-user", user_role=role
    )
    return TestClient(test_app)


app = FastAPI()
app.include_router(router)
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
    user_id="test-user", user_role=LitellmUserRoles.PROXY_ADMIN
)
client = TestClient(app)


@pytest.fixture
def mock_prisma_client():
    with patch("litellm.proxy.proxy_server.prisma_client") as mock:
        yield mock


@pytest.fixture
def mock_user_api_key_auth():
    with patch("litellm.proxy.agent_endpoints.endpoints.user_api_key_auth") as mock:
        mock.return_value = UserAPIKeyAuth(
            user_id="test-user", user_role=LitellmUserRoles.PROXY_ADMIN
        )
        yield mock


def test_update_agent_success(mock_prisma_client, mock_user_api_key_auth, monkeypatch):
    existing_agent = {
        "agent_id": "agent-123",
        "agent_name": "Existing Agent",
        "agent_card_params": _sample_agent_card_params(),
    }
    mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
        return_value=existing_agent
    )

    mock_registry = MagicMock()
    mock_registry.update_agent_in_db = AsyncMock(
        return_value=_sample_agent_response(agent_id="agent-123")
    )
    mock_registry.deregister_agent = MagicMock()
    mock_registry.register_agent = MagicMock()
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)

    response = client.put(
        "/v1/agents/agent-123",
        json=_sample_agent_config(),
        headers={"Authorization": "Bearer test-key"},
    )

    assert response.status_code == 200
    assert response.json()["agent_id"] == "agent-123"
    assert response.json()["agent_name"] == "Test Agent"


def test_update_agent_not_found(
    mock_prisma_client, mock_user_api_key_auth, monkeypatch
):
    mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)

    mock_registry = MagicMock()
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)

    response = client.put(
        "/v1/agents/missing-agent",
        json=_sample_agent_config(),
        headers={"Authorization": "Bearer test-key"},
    )

    assert response.status_code == 404
    assert "Agent with ID missing-agent not found" in response.json()["detail"]


def test_get_agent_by_id_not_found(
    mock_prisma_client, mock_user_api_key_auth, monkeypatch
):
    mock_registry = MagicMock()
    mock_registry.get_agent_by_id = MagicMock(return_value=None)
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)
    mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)

    response = client.get(
        "/v1/agents/missing-agent", headers={"Authorization": "Bearer test-key"}
    )

    assert response.status_code == 404
    assert "Agent with ID missing-agent not found" in response.json()["detail"]


def test_delete_agent_not_found(
    mock_prisma_client, mock_user_api_key_auth, monkeypatch
):
    mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
    mock_registry = MagicMock()
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)

    response = client.delete(
        "/v1/agents/missing-agent", headers={"Authorization": "Bearer test-key"}
    )

    assert response.status_code == 404
    assert "Agent with ID missing-agent not found in DB." in response.json()["detail"]


def test_agent_error_schema_consistency(
    mock_prisma_client, mock_user_api_key_auth, monkeypatch
):
    mock_registry = MagicMock()
    mock_registry.get_agent_by_id = MagicMock(return_value=None)
    mock_registry.update_agent_in_db = AsyncMock(
        side_effect=Exception("should not run")
    )
    mock_registry.delete_agent_from_db = AsyncMock(
        side_effect=Exception("should not run")
    )
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)

    mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)

    missing_agent_id = "missing-agent"
    responses = [
        client.get(
            f"/v1/agents/{missing_agent_id}",
            headers={"Authorization": "Bearer test-key"},
        ),
        client.put(
            f"/v1/agents/{missing_agent_id}",
            json=_sample_agent_config(),
            headers={"Authorization": "Bearer test-key"},
        ),
        client.delete(
            f"/v1/agents/{missing_agent_id}",
            headers={"Authorization": "Bearer test-key"},
        ),
    ]

    for resp in responses:
        assert resp.status_code == 404
        detail = resp.json()["detail"]
        assert isinstance(detail, str)
        assert missing_agent_id in detail


@pytest.mark.asyncio
async def test_get_agent_daily_activity_admin_param_passing(monkeypatch):
    mock_prisma = AsyncMock()
    mock_prisma.db.litellm_agentstable.find_many = AsyncMock(return_value=[])
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)

    mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse")
    get_daily_activity_mock = AsyncMock(return_value=mocked_response)
    monkeypatch.setattr(agent_endpoints, "get_daily_activity", get_daily_activity_mock)

    auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1")
    result = await get_agent_daily_activity(
        agent_ids="agent-1,agent-2",
        start_date="2024-01-01",
        end_date="2024-01-31",
        model="gpt-4",
        api_key="test-key",
        page=2,
        page_size=5,
        exclude_agent_ids="agent-3",
        user_api_key_dict=auth,
    )

    get_daily_activity_mock.assert_awaited_once()
    kwargs = get_daily_activity_mock.call_args.kwargs
    assert kwargs["table_name"] == "litellm_dailyagentspend"
    assert kwargs["entity_id_field"] == "agent_id"
    assert kwargs["entity_id"] == ["agent-1", "agent-2"]
    assert kwargs["exclude_entity_ids"] == ["agent-3"]
    assert kwargs["start_date"] == "2024-01-01"
    assert kwargs["end_date"] == "2024-01-31"
    assert kwargs["model"] == "gpt-4"
    assert kwargs["api_key"] == "test-key"
    assert kwargs["page"] == 2
    assert kwargs["page_size"] == 5
    assert result is mocked_response


@pytest.mark.asyncio
async def test_get_agent_daily_activity_with_agent_names(monkeypatch):
    mock_prisma = AsyncMock()
    mock_agent1 = MagicMock()
    mock_agent1.agent_id = "agent-1"
    mock_agent1.agent_name = "First Agent"
    mock_agent2 = MagicMock()
    mock_agent2.agent_id = "agent-2"
    mock_agent2.agent_name = "Second Agent"

    mock_prisma.db.litellm_agentstable.find_many = AsyncMock(
        return_value=[mock_agent1, mock_agent2]
    )
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)

    mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse")
    get_daily_activity_mock = AsyncMock(return_value=mocked_response)
    monkeypatch.setattr(agent_endpoints, "get_daily_activity", get_daily_activity_mock)

    auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1")
    await get_agent_daily_activity(
        agent_ids="agent-1,agent-2",
        start_date="2024-01-01",
        end_date="2024-01-31",
        model=None,
        api_key=None,
        page=1,
        page_size=10,
        exclude_agent_ids=None,
        user_api_key_dict=auth,
    )

    kwargs = get_daily_activity_mock.call_args.kwargs
    assert kwargs["entity_metadata_field"] == {
        "agent-1": {"agent_name": "First Agent"},
        "agent-2": {"agent_name": "Second Agent"},
    }


@pytest.mark.asyncio
async def test_attach_keys_to_agents_groups_by_agent_and_omits_secret():
    """
    The agents response must carry each agent's attached virtual keys (derived
    from the key table's agent_id FK), grouped per agent, exposing only
    non-secret summary fields. Agents with no key get None so the UI renders
    "Needs Setup" rather than a stale badge.
    """

    class _Row:
        def __init__(self, token, agent_id, key_alias, key_name):
            self.token = token
            self.agent_id = agent_id
            self.key_alias = key_alias
            self.key_name = key_name
            self.user_id = "secret-owner"  # extra field that must NOT leak

    agent_with_keys = _sample_agent_response(agent_id="agent-1")
    agent_without_keys = _sample_agent_response(agent_id="agent-2")

    mock_prisma = MagicMock()
    mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
        return_value=[
            _Row("hash-aaa", "agent-1", "primary", "sk-...aaa"),
            _Row("hash-bbb", "agent-1", "backup", "sk-...bbb"),
        ]
    )

    await _attach_keys_to_agents([agent_with_keys, agent_without_keys], mock_prisma)

    # Query is scoped to the agents being returned, not the whole key table.
    where = mock_prisma.db.litellm_verificationtoken.find_many.call_args.kwargs["where"]
    assert where == {"agent_id": {"in": ("agent-1", "agent-2")}}

    # agent-1 gets both of its keys; agent-2 gets None.
    assert agent_without_keys.keys is None
    assert agent_with_keys.keys is not None
    assert {k.token for k in agent_with_keys.keys} == {"hash-aaa", "hash-bbb"}
    assert {k.key_alias for k in agent_with_keys.keys} == {"primary", "backup"}

    # Only summary fields are exposed; the row's user_id must not be carried.
    summary = agent_with_keys.keys[0]
    assert set(summary.model_dump().keys()) == {"token", "key_alias", "key_name"}


class TestAgentByIdKeyRedaction:
    """GET /v1/agents/{id} surfaces attached keys to admins but never to
    non-admins, even when the agent has keys attached."""

    @pytest.fixture(autouse=True)
    def _setup(self, monkeypatch):
        self.mock_registry = MagicMock()
        self.mock_registry.get_agent_by_id = MagicMock(
            return_value=_sample_agent_response()
        )
        monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry)

    def _get_as(self, role: LitellmUserRoles):
        key_row = MagicMock()
        key_row.token = "hash-aaa"
        key_row.agent_id = "agent-123"
        key_row.key_alias = "primary"
        key_row.key_name = "sk-...aaa"

        test_client = _make_app_with_role(role)
        with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
            mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
                return_value=None
            )
            mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
                return_value=[key_row]
            )
            return test_client.get(
                "/v1/agents/agent-123", headers={"Authorization": "Bearer k"}
            )

    def test_admin_sees_attached_keys(self):
        resp = self._get_as(LitellmUserRoles.PROXY_ADMIN)
        assert resp.status_code == 200
        keys = resp.json()["keys"]
        assert keys is not None
        assert keys[0] == {
            "token": "hash-aaa",
            "key_alias": "primary",
            "key_name": "sk-...aaa",
        }

    def test_non_admin_never_sees_keys(self):
        resp = self._get_as(LitellmUserRoles.INTERNAL_USER)
        assert resp.status_code == 200
        assert resp.json()["keys"] is None

    def test_view_only_admin_reads_a_denied_agent_but_still_without_keys(self):
        """proxy_admin_viewer skips the per-agent object_permission gate (denied
        here) yet stays on the redacted response path."""
        with patch(
            "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed",
            AsyncMock(return_value=False),
        ):
            resp = self._get_as(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
        assert resp.status_code == 200
        assert resp.json()["keys"] is None


# ---------- RBAC enforcement tests ----------


class TestAgentRBACInternalUser:
    """Internal users should be able to read agents but not create/update/delete."""

    @pytest.fixture(autouse=True)
    def _setup(self, monkeypatch):
        self.internal_client = _make_app_with_role(LitellmUserRoles.INTERNAL_USER)
        self.mock_registry = MagicMock()
        monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry)

    def test_should_allow_internal_user_to_list_agents(self, monkeypatch):
        self.mock_registry.get_agent_list = MagicMock(return_value=[])
        resp = self.internal_client.get(
            "/v1/agents", headers={"Authorization": "Bearer k"}
        )
        assert resp.status_code == 200

    def test_should_allow_internal_user_to_get_agent_by_id(self, monkeypatch):
        self.mock_registry.get_agent_by_id = MagicMock(
            return_value=_sample_agent_response()
        )
        with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
            mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
                return_value=None
            )
            mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
                return_value=[]
            )
            resp = self.internal_client.get(
                "/v1/agents/agent-123", headers={"Authorization": "Bearer k"}
            )
        assert resp.status_code == 200

    def test_should_block_internal_user_from_creating_agent(self):
        resp = self.internal_client.post(
            "/v1/agents",
            json=_sample_agent_config(),
            headers={"Authorization": "Bearer k"},
        )
        assert resp.status_code == 403
        assert "Only proxy admins" in resp.json()["detail"]["error"]

    def test_should_block_internal_user_from_updating_agent(self):
        resp = self.internal_client.put(
            "/v1/agents/agent-123",
            json=_sample_agent_config(),
            headers={"Authorization": "Bearer k"},
        )
        assert resp.status_code == 403

    def test_should_block_internal_user_from_patching_agent(self):
        resp = self.internal_client.patch(
            "/v1/agents/agent-123",
            json={"agent_name": "new-name"},
            headers={"Authorization": "Bearer k"},
        )
        assert resp.status_code == 403

    def test_should_block_internal_user_from_deleting_agent(self):
        resp = self.internal_client.delete(
            "/v1/agents/agent-123", headers={"Authorization": "Bearer k"}
        )
        assert resp.status_code == 403


class TestAgentRBACInternalUserViewOnly:
    """View-only internal users should only be able to read agents."""

    @pytest.fixture(autouse=True)
    def _setup(self, monkeypatch):
        self.viewer_client = _make_app_with_role(
            LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
        )
        self.mock_registry = MagicMock()
        monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry)

    def test_should_allow_view_only_user_to_list_agents(self):
        self.mock_registry.get_agent_list = MagicMock(return_value=[])
        resp = self.viewer_client.get(
            "/v1/agents", headers={"Authorization": "Bearer k"}
        )
        assert resp.status_code == 200

    def test_should_block_view_only_user_from_creating_agent(self):
        resp = self.viewer_client.post(
            "/v1/agents",
            json=_sample_agent_config(),
            headers={"Authorization": "Bearer k"},
        )
        assert resp.status_code == 403

    def test_should_block_view_only_user_from_deleting_agent(self):
        resp = self.viewer_client.delete(
            "/v1/agents/agent-123", headers={"Authorization": "Bearer k"}
        )
        assert resp.status_code == 403


SENTINEL_AGENT_API_KEY = "sk-test-sentinel-do-not-use"


class TestAgentRBACProxyAdminViewOnly:
    """Read-only proxy admins go through the object-permission scoped branch on
    GET /v1/agents (the admin fast path stays full PROXY_ADMIN only, so viewers
    cannot fan out health checks beyond their allowlist). litellm_params
    secrets are redacted for every caller, admin included (LIT-6736); only the
    virtual-key/header visibility stays gated on full PROXY_ADMIN."""

    @pytest.fixture(autouse=True)
    def _setup(self, monkeypatch):
        from litellm.proxy.agent_endpoints import agent_registry as ar_mod

        self.viewer_client = _make_app_with_role(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)
        self.admin_client = _make_app_with_role(LitellmUserRoles.PROXY_ADMIN)
        self.agents = [
            AgentResponse(
                agent_id=f"agent-{index}",
                agent_name=f"Agent {index}",
                agent_card_params=_sample_agent_card_params(),
                litellm_params={"api_key": SENTINEL_AGENT_API_KEY},
            )
            for index in (1, 2)
        ]
        self.mock_registry = MagicMock()
        self.mock_registry.get_agent_list = MagicMock(return_value=self.agents)
        self.mock_registry.ids_for_agent = MagicMock(side_effect=lambda agent_id: frozenset({agent_id}))
        monkeypatch.setattr(ar_mod, "global_agent_registry", self.mock_registry)

        self.allowed_agents_spy = AsyncMock(
            return_value=RestrictedAgentAccess(frozenset({"someone-elses-agent"}))
        )
        monkeypatch.setattr(
            "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access",
            self.allowed_agents_spy,
        )

    def _list_agents(self, test_client: TestClient):
        key_row = MagicMock()
        key_row.token = "hash-aaa"
        key_row.agent_id = "agent-1"
        key_row.key_alias = "primary"
        key_row.key_name = "sk-...aaa"

        with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
            mock_prisma.db.litellm_agentstable.find_many = AsyncMock(return_value=[])
            mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(
                return_value=[key_row]
            )
            return test_client.get("/v1/agents", headers={"Authorization": "Bearer k"})

    def test_should_scope_view_only_admin_to_allowed_agents(self):
        """The key/team allowlist here excludes every registered agent; a viewer
        on the admin fast path would see everything, so an empty response pins
        that viewers stay in the scoped branch."""
        resp = self._list_agents(self.viewer_client)

        assert resp.status_code == 200
        assert resp.json() == []
        self.allowed_agents_spy.assert_awaited_once()

    def test_should_still_redact_secrets_for_view_only_admin(self):
        """A viewer granted every agent sees the same agents as an admin but with keys
        stripped; litellm_params secrets never appear in either response."""
        self.allowed_agents_spy.return_value = RestrictedAgentAccess(frozenset({"agent-1", "agent-2"}))
        viewer_resp = self._list_agents(self.viewer_client)
        admin_resp = self._list_agents(self.admin_client)

        assert viewer_resp.status_code == 200
        viewer_by_id = {agent["agent_id"]: agent for agent in viewer_resp.json()}
        assert set(viewer_by_id) == {"agent-1", "agent-2"}
        assert viewer_by_id["agent-1"]["keys"] is None
        assert SENTINEL_AGENT_API_KEY not in viewer_resp.text

        admin_by_id = {agent["agent_id"]: agent for agent in admin_resp.json()}
        assert admin_by_id["agent-1"]["keys"][0]["token"] == "hash-aaa"
        assert SENTINEL_AGENT_API_KEY not in admin_resp.text
        assert admin_by_id["agent-1"]["litellm_params"]["api_key"] == REDACTED_BY_LITELM_STRING


class TestAgentRBACProxyAdmin:
    """Proxy admins should have full CRUD access to agents."""

    @pytest.fixture(autouse=True)
    def _setup(self, monkeypatch):
        self.admin_client = _make_app_with_role(LitellmUserRoles.PROXY_ADMIN)
        self.mock_registry = MagicMock()
        monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry)

    def test_should_allow_admin_to_create_agent(self, monkeypatch):
        with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
            self.mock_registry.get_agent_by_name = MagicMock(return_value=None)
            self.mock_registry.add_agent_to_db = AsyncMock(
                return_value=_sample_agent_response()
            )
            self.mock_registry.register_agent = MagicMock()
            resp = self.admin_client.post(
                "/v1/agents",
                json=_sample_agent_config(),
                headers={"Authorization": "Bearer k"},
            )
            assert resp.status_code == 200

    def test_create_agent_applies_litellm_merge_to_stored_card(self):
        """The card stored in the DB must reflect the LiteLLM-fronting merge."""
        with patch("litellm.proxy.proxy_server.prisma_client"):
            self.mock_registry.get_agent_by_name = MagicMock(return_value=None)
            self.mock_registry.add_agent_to_db = AsyncMock(
                return_value=_sample_agent_response()
            )
            self.mock_registry.register_agent = MagicMock()

            self.admin_client.post(
                "/v1/agents",
                json=_sample_agent_config(),
                headers={"Authorization": "Bearer k"},
            )

            call_kwargs = self.mock_registry.add_agent_to_db.await_args.kwargs
            stored_card = call_kwargs["agent"]["agent_card_params"]
            new_agent_id = call_kwargs["agent_id"]

            # Top-level url is retained for runtime A2A invocation (the public
            # well-known endpoint rewrites it before exposing to clients);
            # supportedInterfaces points at the proxy.
            assert stored_card["url"] == "http://localhost"
            assert stored_card["supportedInterfaces"][0]["protocolBinding"] == "JSONRPC"
            assert stored_card["supportedInterfaces"][0]["url"].endswith(
                f"/a2a/{new_agent_id}"
            )
            # Security scheme is the LiteLLM scheme.
            assert "LiteLLMKey" in stored_card["securitySchemes"]

    def test_create_agent_response_never_echoes_secret(self):
        """LIT-6736: POST /v1/agents must not echo the stored secret back, even
        though it's the caller's own value and even for a proxy admin."""
        with patch("litellm.proxy.proxy_server.prisma_client"):  # test-quality-ok: proxy_server module global is the endpoint's only injection point
            self.mock_registry.get_agent_by_name = MagicMock(return_value=None)
            self.mock_registry.add_agent_to_db = AsyncMock(
                return_value=AgentResponse(
                    agent_id="agent-123",
                    agent_name="Test Agent",
                    agent_card_params=_sample_agent_card_params(),
                    litellm_params={
                        "aws_secret_access_key": SENTINEL_AGENT_API_KEY,
                        "model": "bedrock/agentcore/my-agent",
                    },
                )
            )
            self.mock_registry.register_agent = MagicMock()

            resp = self.admin_client.post(
                "/v1/agents",
                json={
                    "agent_name": "Test Agent",
                    "agent_card_params": _sample_agent_card_params(),
                    "litellm_params": {
                        "aws_secret_access_key": SENTINEL_AGENT_API_KEY,
                        "model": "bedrock/agentcore/my-agent",
                    },
                },
                headers={"Authorization": "Bearer k"},
            )

            assert resp.status_code == 200
            assert SENTINEL_AGENT_API_KEY not in resp.text
            body = resp.json()
            assert body["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING
            assert body["litellm_params"]["model"] == "bedrock/agentcore/my-agent"

    def test_update_agent_response_never_echoes_secret(self):
        """LIT-6736: PUT /v1/agents/{id} must not echo the stored secret back."""
        with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:  # test-quality-ok: proxy_server module global is the endpoint's only injection point
            mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
                return_value={
                    "agent_id": "agent-123",
                    "agent_name": "Existing Agent",
                    "agent_card_params": _sample_agent_card_params(),
                }
            )
            self.mock_registry.update_agent_in_db = AsyncMock(
                return_value=AgentResponse(
                    agent_id="agent-123",
                    agent_name="Test Agent",
                    agent_card_params=_sample_agent_card_params(),
                    litellm_params={"aws_secret_access_key": SENTINEL_AGENT_API_KEY},
                )
            )
            self.mock_registry.deregister_agent = MagicMock()
            self.mock_registry.register_agent = MagicMock()

            resp = self.admin_client.put(
                "/v1/agents/agent-123",
                json={
                    "agent_name": "Test Agent",
                    "agent_card_params": _sample_agent_card_params(),
                    "litellm_params": {"aws_secret_access_key": REDACTED_BY_LITELM_STRING},
                },
                headers={"Authorization": "Bearer k"},
            )

            assert resp.status_code == 200
            assert SENTINEL_AGENT_API_KEY not in resp.text
            assert resp.json()["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING

    def test_patch_agent_response_never_echoes_secret(self):
        """LIT-6736: PATCH /v1/agents/{id} must not echo the stored secret back."""
        with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:  # test-quality-ok: proxy_server module global is the endpoint's only injection point
            mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
                return_value={
                    "agent_id": "agent-123",
                    "agent_name": "Existing Agent",
                    "agent_card_params": _sample_agent_card_params(),
                }
            )
            self.mock_registry.patch_agent_in_db = AsyncMock(
                return_value=AgentResponse(
                    agent_id="agent-123",
                    agent_name="Renamed Agent",
                    agent_card_params=_sample_agent_card_params(),
                    litellm_params={"aws_secret_access_key": SENTINEL_AGENT_API_KEY},
                )
            )
            self.mock_registry.deregister_agent = MagicMock()
            self.mock_registry.register_agent = MagicMock()

            resp = self.admin_client.patch(
                "/v1/agents/agent-123",
                json={"agent_name": "Renamed Agent"},
                headers={"Authorization": "Bearer k"},
            )

            assert resp.status_code == 200
            assert SENTINEL_AGENT_API_KEY not in resp.text
            assert resp.json()["litellm_params"]["aws_secret_access_key"] == REDACTED_BY_LITELM_STRING

    def test_should_allow_admin_to_delete_agent(self):
        existing = {
            "agent_id": "agent-123",
            "agent_name": "Existing Agent",
            "agent_card_params": _sample_agent_card_params(),
        }
        with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
            mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
                return_value=existing
            )
            self.mock_registry.delete_agent_from_db = AsyncMock()
            self.mock_registry.deregister_agent = MagicMock()
            resp = self.admin_client.delete(
                "/v1/agents/agent-123", headers={"Authorization": "Bearer k"}
            )
            assert resp.status_code == 200


class TestAgentProtocolVersionValidation:
    """Registration accepts spec-default semver protocolVersion values and still
    rejects genuinely unsupported versions."""

    @pytest.fixture(autouse=True)
    def _setup(self, monkeypatch):
        self.admin_client = _make_app_with_role(LitellmUserRoles.PROXY_ADMIN)
        self.mock_registry = MagicMock()
        monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", self.mock_registry)

    def _create_agent_with_protocol_version(self, protocol_version: str):
        config = _sample_agent_config()
        config["agent_card_params"]["protocolVersion"] = protocol_version
        with patch("litellm.proxy.proxy_server.prisma_client"):
            self.mock_registry.get_agent_by_name = MagicMock(return_value=None)
            self.mock_registry.add_agent_to_db = AsyncMock(
                return_value=_sample_agent_response()
            )
            self.mock_registry.register_agent = MagicMock()
            return self.admin_client.post(
                "/v1/agents",
                json=config,
                headers={"Authorization": "Bearer k"},
            )

    def test_semver_protocol_version_registers_and_stores_major_minor(self):
        resp = self._create_agent_with_protocol_version("0.3.0")
        assert resp.status_code == 200
        stored_card = self.mock_registry.add_agent_to_db.await_args.kwargs["agent"][
            "agent_card_params"
        ]
        assert stored_card["protocolVersion"] == "0.3"
        assert stored_card["supportedInterfaces"][0]["protocolVersion"] == "0.3"

    def test_unsupported_protocol_version_is_rejected(self):
        resp = self._create_agent_with_protocol_version("0.2.6")
        assert resp.status_code == 400
        assert "Unsupported protocolVersion '0.2.6'" in resp.json()["detail"]
        self.mock_registry.add_agent_to_db.assert_not_awaited()

    def test_malformed_protocol_version_is_rejected(self):
        resp = self._create_agent_with_protocol_version("0.3.garbage")
        assert resp.status_code == 400
        assert "Unsupported protocolVersion '0.3.garbage'" in resp.json()["detail"]
        self.mock_registry.add_agent_to_db.assert_not_awaited()


class TestCheckAgentManagementPermission:
    """Unit tests for the _check_agent_management_permission helper."""

    def test_should_allow_proxy_admin(self):
        auth = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
        _check_agent_management_permission(auth)

    @pytest.mark.parametrize(
        "role",
        [
            LitellmUserRoles.INTERNAL_USER,
            LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
            LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
        ],
    )
    def test_should_block_non_admin_roles(self, role):
        from fastapi import HTTPException

        auth = UserAPIKeyAuth(user_id="user", user_role=role)
        with pytest.raises(HTTPException) as exc_info:
            _check_agent_management_permission(auth)
        assert exc_info.value.status_code == 403


class TestAgentRoutesIncludesAgentIdPattern:
    """Verify that agent_routes includes the {agent_id} pattern for route access."""

    def test_should_include_agent_id_pattern(self):
        from litellm.proxy._types import LiteLLMRoutes

        assert "/v1/agents/{agent_id}" in LiteLLMRoutes.agent_routes.value


class TestAgentHealthCheck:
    """Tests for the health_check query parameter on GET /v1/agents."""

    @pytest.fixture(autouse=True)
    def _setup(self, monkeypatch):
        from litellm.proxy.agent_endpoints import agent_registry as ar_mod

        self.admin_client = _make_app_with_role(LitellmUserRoles.PROXY_ADMIN)
        self.mock_registry = MagicMock()
        monkeypatch.setattr(ar_mod, "global_agent_registry", self.mock_registry)
        # Ensure prisma_client is None so the endpoint skips DB queries.
        # In CI with parallel workers, a MagicMock can leak from other test
        # scopes, causing "object MagicMock can't be used in 'await'" errors.
        monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)

    def _make_agent(self, agent_id: str, url: str | None = None) -> AgentResponse:
        card = _sample_agent_card_params()
        if url is not None:
            card["url"] = url
        else:
            card.pop("url", None)
        return AgentResponse(
            agent_id=agent_id,
            agent_name=f"Agent {agent_id}",
            agent_card_params=card,
            litellm_params={},
        )

    def test_should_return_all_agents_when_health_check_disabled(self):
        agents = [
            self._make_agent("a1", "http://reachable"),
            self._make_agent("a2", "http://unreachable"),
        ]
        self.mock_registry.get_agent_list = MagicMock(return_value=agents)

        resp = self.admin_client.get(
            "/v1/agents", headers={"Authorization": "Bearer k"}
        )
        assert resp.status_code == 200
        assert len(resp.json()) == 2

    def test_should_filter_unhealthy_agents_when_health_check_enabled(
        self, monkeypatch
    ):
        agents = [
            self._make_agent("a1", "http://reachable"),
            self._make_agent("a2", "http://unreachable"),
        ]
        self.mock_registry.get_agent_list = MagicMock(return_value=agents)

        results = iter(
            [
                {"agent_id": "a1", "healthy": True},
                {"agent_id": "a2", "healthy": False, "error": "Connection refused"},
            ]
        )
        monkeypatch.setattr(
            agent_endpoints,
            "_check_agent_url_health",
            AsyncMock(side_effect=lambda agent: next(results)),
        )

        resp = self.admin_client.get(
            "/v1/agents?health_check=true",
            headers={"Authorization": "Bearer k"},
        )
        assert resp.status_code == 200
        data = resp.json()
        assert len(data) == 1
        assert data[0]["agent_id"] == "a1"

    def test_should_return_empty_list_when_all_agents_unhealthy(self, monkeypatch):
        agents = [self._make_agent("a1", "http://down")]
        self.mock_registry.get_agent_list = MagicMock(return_value=agents)
        monkeypatch.setattr(
            agent_endpoints,
            "_check_agent_url_health",
            AsyncMock(
                return_value={"agent_id": "a1", "healthy": False, "error": "timeout"}
            ),
        )

        resp = self.admin_client.get(
            "/v1/agents?health_check=true",
            headers={"Authorization": "Bearer k"},
        )
        assert resp.status_code == 200
        assert len(resp.json()) == 0

    def test_should_return_all_agents_when_all_healthy(self, monkeypatch):
        agents = [
            self._make_agent("a1", "http://ok1"),
            self._make_agent("a2", "http://ok2"),
        ]
        self.mock_registry.get_agent_list = MagicMock(return_value=agents)

        results = iter(
            [
                {"agent_id": "a1", "healthy": True},
                {"agent_id": "a2", "healthy": True},
            ]
        )
        monkeypatch.setattr(
            agent_endpoints,
            "_check_agent_url_health",
            AsyncMock(side_effect=lambda agent: next(results)),
        )

        resp = self.admin_client.get(
            "/v1/agents?health_check=true",
            headers={"Authorization": "Bearer k"},
        )
        assert resp.status_code == 200
        assert len(resp.json()) == 2


class TestCheckAgentUrlHealth:
    """Unit tests for the _check_agent_url_health helper."""

    @pytest.mark.asyncio
    async def test_should_return_healthy_when_no_url(self):
        from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health

        agent = AgentResponse(
            agent_id="no-url",
            agent_name="No URL Agent",
            agent_card_params={"name": "test"},
            litellm_params={},
        )
        result = await _check_agent_url_health(agent)
        assert result["healthy"] is True
        assert "error" not in result

    @pytest.mark.asyncio
    @patch("litellm.proxy.agent_endpoints.endpoints.get_async_httpx_client")
    async def test_should_return_healthy_for_200(self, mock_get_client):
        from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health

        mock_response = MagicMock()
        mock_response.status_code = 200
        mock_client = AsyncMock()
        mock_client.get = AsyncMock(return_value=mock_response)
        mock_get_client.return_value = mock_client

        agent = AgentResponse(
            agent_id="ok",
            agent_name="OK Agent",
            agent_card_params={"url": "http://example.com"},
            litellm_params={},
        )
        result = await _check_agent_url_health(agent)
        assert result["healthy"] is True

    @pytest.mark.asyncio
    @patch("litellm.proxy.agent_endpoints.endpoints.get_async_httpx_client")
    async def test_should_return_unhealthy_for_500(self, mock_get_client):
        from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health

        mock_response = MagicMock()
        mock_response.status_code = 500
        mock_client = AsyncMock()
        mock_client.get = AsyncMock(return_value=mock_response)
        mock_get_client.return_value = mock_client

        agent = AgentResponse(
            agent_id="err",
            agent_name="Error Agent",
            agent_card_params={"url": "http://failing.com"},
            litellm_params={},
        )
        result = await _check_agent_url_health(agent)
        assert result["healthy"] is False
        assert "HTTP 500" in result["error"]

    @pytest.mark.asyncio
    @patch("litellm.proxy.agent_endpoints.endpoints.get_async_httpx_client")
    async def test_should_return_unhealthy_on_connection_error(self, mock_get_client):
        from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health

        mock_client = AsyncMock()
        mock_client.get = AsyncMock(side_effect=Exception("Connection refused"))
        mock_get_client.return_value = mock_client

        agent = AgentResponse(
            agent_id="down",
            agent_name="Down Agent",
            agent_card_params={"url": "http://down.com"},
            litellm_params={},
        )
        result = await _check_agent_url_health(agent)
        assert result["healthy"] is False
        assert "Connection refused" in result["error"]

    @pytest.mark.asyncio
    @patch("litellm.proxy.agent_endpoints.endpoints.get_async_httpx_client")
    async def test_should_treat_404_as_healthy(self, mock_get_client):
        """A 404 means the server is reachable, just not the specific path."""
        from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health

        mock_response = MagicMock()
        mock_response.status_code = 404
        mock_client = AsyncMock()
        mock_client.get = AsyncMock(return_value=mock_response)
        mock_get_client.return_value = mock_client

        agent = AgentResponse(
            agent_id="notfound",
            agent_name="NotFound Agent",
            agent_card_params={"url": "http://example.com/missing"},
            litellm_params={},
        )
        result = await _check_agent_url_health(agent)
        assert result["healthy"] is True


@pytest.mark.parametrize(
    "base_url",
    ["http://0.0.0.0:4000/", "http://localhost:4000/", "https://api.example.com/"],
)
def test_merged_agent_card_url_has_no_double_slash_without_proxy_base_url(
    monkeypatch, base_url
):
    """Without PROXY_BASE_URL, request.base_url carries a trailing slash; the merged
    card's supportedInterfaces URL must still join cleanly (no `//a2a`)."""
    from litellm.proxy.agent_endpoints.endpoints import _build_merged_agent_card

    monkeypatch.delenv("PROXY_BASE_URL", raising=False)
    monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)

    http_request = MagicMock()
    http_request.base_url = base_url

    merged = _build_merged_agent_card(
        _sample_agent_card_params(),
        agent_id="agent-xyz",
        http_request=http_request,
        agent_name="Test Agent",
    )

    interface_url = merged["supportedInterfaces"][0]["url"]
    assert interface_url == f"{base_url.rstrip('/')}/a2a/agent-xyz"
    assert "//a2a" not in interface_url


class _DbBackedProxyConfig:
    """Round-trips `litellm_settings` through the DB overlay the proxy applies on every
    `get_config()`, which is what re-assigns the `litellm.public_*` globals in production.

    Storage goes through JSON the way the `litellm_config` row does, so every read hands back
    freshly built values instead of the objects the endpoint still holds a reference to."""

    def __init__(self, stored_litellm_settings: dict[str, object] | None = None) -> None:
        self.stored_litellm_settings_json: str = json.dumps(stored_litellm_settings or {})

    async def get_config(self) -> dict[str, dict[str, object]]:
        from litellm.proxy.proxy_server import ProxyConfig

        config: Final[dict[str, dict[str, object]]] = {"litellm_settings": {}}
        db_param_value: Final[dict[str, object]] = json.loads(self.stored_litellm_settings_json)
        if not db_param_value:
            return config
        proxy_config: Final = ProxyConfig()
        db_values: Final = proxy_config._prepared_db_settings_values("litellm_settings", db_param_value)
        proxy_config._apply_litellm_settings_db_values(db_values)
        return {"litellm_settings": dict(proxy_config.litellm_settings.resolved())}

    async def save_config(self, new_config: dict[str, dict[str, object]]) -> None:
        self.stored_litellm_settings_json = json.dumps(new_config.get("litellm_settings") or {})


def test_make_agent_public_twice_keeps_both_agents_public(monkeypatch: pytest.MonkeyPatch) -> None:
    """A second /make_public call must not drop the agent published by the first one."""
    import litellm
    from litellm.proxy.agent_endpoints import agent_registry as agent_registry_module
    from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry

    registry: Final = AgentRegistry()
    registry.register_agent(_sample_agent_response(agent_id="agent-1", agent_name="Agent One"))
    registry.register_agent(_sample_agent_response(agent_id="agent-2", agent_name="Agent Two"))

    monkeypatch.setattr(agent_registry_module, "global_agent_registry", registry)
    monkeypatch.setattr(litellm, "public_agent_groups", None)
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", _DbBackedProxyConfig())

    first: Final = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"})
    second: Final = client.post("/v1/agents/agent-2/make_public", headers={"Authorization": "Bearer test-key"})

    assert first.status_code == 200
    assert second.status_code == 200
    assert second.json()["public_agent_groups"] == ["agent-1", "agent-2"]
    assert [agent.agent_id for agent in registry.get_public_agent_list()] == ["agent-1", "agent-2"]


def test_make_agent_public_rejects_an_agent_published_only_in_the_db(monkeypatch: pytest.MonkeyPatch) -> None:
    """The duplicate guard must fire off the stored list, not just what this process published."""
    import litellm
    from litellm.proxy.agent_endpoints import agent_registry as agent_registry_module
    from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry

    registry: Final = AgentRegistry()
    registry.register_agent(_sample_agent_response(agent_id="agent-1", agent_name="Agent One"))

    monkeypatch.setattr(agent_registry_module, "global_agent_registry", registry)
    monkeypatch.setattr(litellm, "public_agent_groups", None)
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
    monkeypatch.setattr(
        "litellm.proxy.proxy_server.proxy_config",
        _DbBackedProxyConfig({"public_agent_groups": ["agent-1"]}),
    )

    duplicate: Final = client.post("/v1/agents/agent-1/make_public", headers={"Authorization": "Bearer test-key"})

    assert duplicate.status_code == 400
    assert "already in public agent groups" in duplicate.json()["detail"]


_KILL_SWITCH: Final = {
    "url": "https://ops.example.com/kill",
    "method": "POST",
    "headers": {"X-Env": "prod"},
    "query_params": {"reason": "manual"},
    "body": {"action": "stop"},
    "auth": {"type": "bearer", "token": "tok-real"},
}


def _agent_with_kill_switch() -> AgentResponse:
    return AgentResponse(
        agent_id="agent-123",
        agent_name="Test Agent",
        agent_card_params=_sample_agent_card_params(),
        litellm_params={},
        kill_switch=_KILL_SWITCH,
    )


class _FakeKillSwitchClient:
    def __init__(self, response: httpx.Response) -> None:
        self.calls: list[tuple[str, str, dict[str, str], object, float]] = []  # mutable-ok: test double records calls
        self._response: Final = response

    def build_request(self, method: str, url: str, *, headers, json, timeout: float) -> httpx.Request:
        self.calls.append((method, url, dict(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:
        return self._response


class _AuditLogRecorder:
    def __init__(self) -> None:
        self.rows: list[LiteLLM_AuditLogs] = []  # mutable-ok: test double records writes

    async def __call__(self, request_data: LiteLLM_AuditLogs) -> None:
        self.rows.append(request_data)


def _kill_switch_app(
    role: LitellmUserRoles,
    http_client: _FakeKillSwitchClient,
    audit_log: _AuditLogRecorder | None = None,
) -> TestClient:
    test_client: Final = _make_app_with_role(role)
    test_client.app.dependency_overrides[agent_endpoints.default_kill_switch_http_client] = lambda: http_client
    test_client.app.dependency_overrides[agent_endpoints.default_kill_switch_audit_log_writer] = (
        lambda: audit_log or _AuditLogRecorder()
    )
    return test_client


def test_kill_switch_trigger_fires_the_configured_webhook_and_returns_the_result(monkeypatch) -> None:
    registry: Final = MagicMock()
    registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch())
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry)
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
    fake: Final = _FakeKillSwitchClient(httpx.Response(200, text="ok"))

    resp: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake).post(
        "/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"}
    )

    assert resp.status_code == 200, resp.text
    assert resp.json() == {
        "agent_id": "agent-123",
        "url": "https://ops.example.com/kill",
        "method": "POST",
        "status_code": 200,
        "response_body": "ok",
        "error": None,
    }
    (method, url, headers, body, _timeout) = fake.calls[0]
    assert (method, url, body) == ("POST", "https://ops.example.com/kill?reason=manual", {"action": "stop"})
    assert headers == {"X-Env": "prod", "Authorization": "Bearer tok-real"}


def test_kill_switch_trigger_writes_an_audit_log_row_naming_the_admin_and_the_sanitized_result(monkeypatch) -> None:
    registry: Final = MagicMock()
    registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch())
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry)
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
    fake: Final = _FakeKillSwitchClient(httpx.Response(202, text='{"stopped": true}'))
    audit: Final = _AuditLogRecorder()
    test_client: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake, audit)
    test_client.app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
        user_id="test-user", user_role=LitellmUserRoles.PROXY_ADMIN, api_key="hashed-k"
    )

    resp: Final = test_client.post("/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"})

    assert resp.status_code == 200, resp.text
    (row,) = audit.rows
    assert (row.action, row.table_name, row.object_id) == (
        "kill_switch_fired",
        LitellmTableNames.AGENT_TABLE_NAME,
        "agent-123",
    )
    assert (row.changed_by, row.changed_by_api_key) == ("test-user", "hashed-k")
    assert row.before_value is None
    assert json.loads(row.updated_values) == {
        "agent_id": "agent-123",
        "url": "https://ops.example.com/kill",
        "method": "POST",
        "status_code": 202,
        "response_body": '{"stopped": true}',
    }
    assert "tok-real" not in row.model_dump_json()


def test_kill_switch_trigger_returns_502_and_still_audits_when_the_webhook_rejects(monkeypatch) -> None:
    registry: Final = MagicMock()
    registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch())
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry)
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
    fake: Final = _FakeKillSwitchClient(httpx.Response(401, text="bad token"))
    audit: Final = _AuditLogRecorder()

    resp: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake, audit).post(
        "/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"}
    )

    assert resp.status_code == 502, resp.text
    assert resp.json()["detail"]["status_code"] == 401
    assert resp.json()["detail"]["response_body"] == "bad token"
    (row,) = audit.rows
    assert row.action == "kill_switch_fired"
    assert json.loads(row.updated_values)["status_code"] == 401


@pytest.mark.parametrize("role", [LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
def test_kill_switch_trigger_is_refused_before_any_webhook_call_for_non_admins(monkeypatch, role) -> None:
    registry: Final = MagicMock()
    registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch())
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry)
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
    fake: Final = _FakeKillSwitchClient(httpx.Response(200))
    audit: Final = _AuditLogRecorder()

    resp: Final = _kill_switch_app(role, fake, audit).post(
        "/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"}
    )

    assert resp.status_code == 403, resp.text
    assert fake.calls == []
    assert audit.rows == []


def test_kill_switch_trigger_404s_unknown_agent_and_400s_an_agent_without_one(monkeypatch) -> None:
    registry: Final = MagicMock()
    registry.get_agent_by_id = MagicMock(side_effect=[None, _sample_agent_response()])
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry)
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
    fake: Final = _FakeKillSwitchClient(httpx.Response(200))
    audit: Final = _AuditLogRecorder()
    test_client: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake, audit)

    missing: Final = test_client.post("/v1/agents/nope/kill_switch", headers={"Authorization": "Bearer k"})
    unconfigured: Final = test_client.post("/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"})

    assert missing.status_code == 404
    assert unconfigured.status_code == 400
    assert "no kill_switch configured" in unconfigured.json()["detail"]
    assert fake.calls == []
    assert audit.rows == []


def test_kill_switch_trigger_fires_the_db_row_config_over_a_stale_in_memory_copy(monkeypatch) -> None:
    """Another replica may have updated the agent; the row is the source of truth for what gets fired."""
    registry: Final = MagicMock()
    registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch())
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry)
    db_row: Final = SimpleNamespace(
        agent_id="agent-123",
        kill_switch={"url": "https://ops.example.com/kill-v2", "method": "DELETE", "auth": None},
    )
    prisma: Final = MagicMock()
    prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=db_row)
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma)
    fake: Final = _FakeKillSwitchClient(httpx.Response(204))

    resp: Final = _kill_switch_app(LitellmUserRoles.PROXY_ADMIN, fake).post(
        "/v1/agents/agent-123/kill_switch", headers={"Authorization": "Bearer k"}
    )

    assert resp.status_code == 200, resp.text
    (method, url, headers, body, _timeout) = fake.calls[0]
    assert (method, url, headers, body) == ("DELETE", "https://ops.example.com/kill-v2", {}, None)
    assert prisma.db.litellm_agentstable.find_unique.await_args.kwargs == {"where": {"agent_id": "agent-123"}}
    registry.get_agent_by_id.assert_not_called()


def test_get_agent_redacts_kill_switch_secret_for_admins_and_hides_it_from_others(monkeypatch) -> None:
    registry: Final = MagicMock()
    registry.get_agent_by_id = MagicMock(return_value=_agent_with_kill_switch())
    registry.ids_for_agent = MagicMock(return_value=("agent-123",))
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry)

    def _get_as(role: LitellmUserRoles):
        with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma:
            mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
            mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
            return _make_app_with_role(role).get("/v1/agents/agent-123", headers={"Authorization": "Bearer k"})

    admin: Final = _get_as(LitellmUserRoles.PROXY_ADMIN)
    assert admin.status_code == 200, admin.text
    assert admin.json()["kill_switch"] == {
        **_KILL_SWITCH,
        "auth": {"type": "bearer", "token": REDACTED_BY_LITELM_STRING},
    }

    internal: Final = _get_as(LitellmUserRoles.INTERNAL_USER)
    assert internal.status_code == 200, internal.text
    assert internal.json()["kill_switch"] is None
    assert "tok-real" not in internal.text