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

import asyncio
import json
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

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

import httpx
import pytest
import yaml
from fastapi import FastAPI, HTTPException
from fastapi.testclient import TestClient
from prisma.models import LiteLLM_AgentsTable

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.agent_registry import AgentRegistry
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, PatchAgentRequest
from litellm.types.proxy.agent_identity import AgentIdentityBinding


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=AgentResponse.model_validate(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"]


class _AgentPersistence:
    def __init__(self, row: LiteLLM_AgentsTable) -> None:
        self.row = row

    async def find_unique(self, **kwargs: object) -> LiteLLM_AgentsTable:
        return self.row

    async def update(self, *, data: Mapping[str, object], **kwargs: object) -> LiteLLM_AgentsTable:
        from tests.unit.proxy.agent_endpoints.test_agent_registry import _stored_agent_row

        self.row = _stored_agent_row({**self.row.model_dump(), **data})
        return self.row


@pytest.mark.parametrize("method", ["PUT", "PATCH"])
@pytest.mark.parametrize("cardless", [False, True])
def test_identity_settings_edit_preserves_runtime_configuration_on_readback(
    monkeypatch: pytest.MonkeyPatch, method: str, cardless: bool
) -> None:
    from litellm.proxy import proxy_server
    from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
    from tests.unit.proxy.agent_endpoints.test_agent_registry import _stored_agent_row

    runtime: Final = {
        "agent_card_params": {} if cardless else _sample_agent_card_params(),
        "litellm_params": {"make_public": False, "model": "a2a/runtime"},
        "static_headers": {"X-Runtime": "configured"},
        "extra_headers": ["X-Trace"],
        "access_group_ids": ["runtime-group"],
        "kill_switch": {"url": "https://runtime.example/stop", "method": "POST"},
    }
    row: Final = _stored_agent_row(runtime)
    table: Final = _AgentPersistence(row)
    database: Final = SimpleNamespace(
        litellm_agentstable=table,
        litellm_verificationtoken=SimpleNamespace(find_many=AsyncMock(return_value=[])),
    )
    monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=database, writer_db=database))
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", AgentRegistry())

    response: Final = client.request(
        method, "/v1/agents/agent-123", json={"agent_name": "Renamed agent", "enabled": False}
    )
    assert response.status_code == 200, response.text
    readback: Final = client.get("/v1/agents/agent-123")
    assert readback.status_code == 200, readback.text
    stored: Final = AgentResponse.model_validate(table.row.model_dump())
    expected: Final = AgentResponse.model_validate(row.model_dump()).model_copy(
        update={"agent_name": "Renamed agent", "enabled": False}
    )
    preserved: Final = {*runtime, "agent_name", "enabled", "agent_id"}
    assert stored.model_dump(include=preserved) == expected.model_dump(include=preserved)
    assert {key: readback.json()[key] for key in preserved} == expected.model_dump(mode="json", include=preserved)


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.writer_db = mock_prisma.db
            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.writer_db = mock_prisma.db
            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_rejects_legacy_litellm_params_identity(self):
        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=_sample_agent_response())
            config = _sample_agent_config()
            config["litellm_params"] = {
                **config["litellm_params"],
                "identity": {
                    "provider": "microsoft_entra",
                    "tenant_id": "11111111-1111-4111-8111-111111111111",
                    "client_id": "22222222-2222-4222-8222-222222222222",
                },
            }
            resp = self.admin_client.post("/v1/agents", json=config, headers={"Authorization": "Bearer k"})
            assert resp.status_code == 400, resp.text
            assert "top-level identity field" in resp.json()["detail"]
            self.mock_registry.add_agent_to_db.assert_not_awaited()

    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=AgentResponse(
                    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=AgentResponse(
                    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


@dataclass(frozen=True, slots=True)
class _StoredConfigRow:
    param_name: str
    param_value: dict[str, object]


class _ConfigRows:
    """`LiteLLM_Config` as a section-locked transaction reads it, shared by every replica."""

    def __init__(self, rows: Mapping[str, Mapping[str, object]]) -> None:
        self.serialized_rows: dict[str, str] = {name: json.dumps(dict(value)) for name, value in rows.items()}
        self.section_lock: Final = asyncio.Lock()

    async def find_first(self, *, where: Mapping[str, str]) -> _StoredConfigRow | None:
        name: Final = where["param_name"]
        serialized: Final = self.serialized_rows.get(name)
        await asyncio.sleep(0)
        return None if serialized is None else _StoredConfigRow(param_name=name, param_value=json.loads(serialized))

    async def upsert(self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, str]]) -> None:
        self.serialized_rows[where["param_name"]] = data["update"]["param_value"]

    def stored(self, name: str) -> dict[str, object]:
        return json.loads(self.serialized_rows[name])


class _ConfigTransaction:
    def __init__(self, rows: _ConfigRows) -> None:
        self.litellm_config: Final = rows
        self.holds_section_lock = False

    async def __aenter__(self) -> "_ConfigTransaction":
        return self

    async def __aexit__(self, *_: object) -> None:
        if self.holds_section_lock:
            self.litellm_config.section_lock.release()

    async def query_raw(self, query: str, section_name: str) -> None:
        await self.litellm_config.section_lock.acquire()
        self.holds_section_lock = True


class _StaleReplicaPrisma:
    """A replica whose config reads still return what it cached before another replica wrote."""

    def __init__(self, rows: _ConfigRows, cached_rows: Mapping[str, Mapping[str, object]]) -> None:
        self.rows: Final = rows
        self._cached_rows: Final = cached_rows

    async def get_generic_data(self, key: str, value: str, table_name: str) -> SimpleNamespace | None:
        cached: Final = self._cached_rows.get(value)
        return None if cached is None else SimpleNamespace(param_name=value, param_value=dict(cached))

    def tx(self) -> _ConfigTransaction:
        return _ConfigTransaction(self.rows)


_PROXY_ADMIN: Final = UserAPIKeyAuth(user_id="test-user", user_role=LitellmUserRoles.PROXY_ADMIN)


def _replica_that_missed_a_publish(
    monkeypatch: pytest.MonkeyPatch, *, stored_public_agent_groups: list[str], store_model_in_db: bool = True
) -> tuple[_ConfigRows, AgentRegistry]:
    import litellm
    from litellm.proxy.agent_endpoints import agent_registry as agent_registry_module
    from litellm.proxy.proxy_server import ProxyConfig

    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"))
    rows: Final = _ConfigRows({"litellm_settings": {"public_agent_groups": stored_public_agent_groups}})
    cached_rows: Final = {"litellm_settings": {"public_agent_groups": []}}

    monkeypatch.setattr(agent_registry_module, "global_agent_registry", registry)
    monkeypatch.setattr(litellm, "public_agent_groups", [])
    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _StaleReplicaPrisma(rows, cached_rows))
    monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", store_model_in_db)
    monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"store_model_in_db": store_model_in_db})
    monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", ProxyConfig())
    return rows, registry


@pytest.mark.asyncio
async def test_make_agent_public_keeps_an_agent_another_replica_published(monkeypatch: pytest.MonkeyPatch) -> None:
    rows, registry = _replica_that_missed_a_publish(monkeypatch, stored_public_agent_groups=["agent-1"])

    response: Final = await agent_endpoints.make_agent_public(agent_id="agent-2", user_api_key_dict=_PROXY_ADMIN)

    assert response["public_agent_groups"] == ["agent-1", "agent-2"]
    assert rows.stored("litellm_settings")["public_agent_groups"] == ["agent-1", "agent-2"]
    assert [agent.agent_id for agent in registry.get_public_agent_list()] == ["agent-1", "agent-2"]


@pytest.mark.asyncio
async def test_make_agent_public_concurrent_publishes_all_stay_public(monkeypatch: pytest.MonkeyPatch) -> None:
    rows, _ = _replica_that_missed_a_publish(monkeypatch, stored_public_agent_groups=[])

    await asyncio.gather(
        agent_endpoints.make_agent_public(agent_id="agent-1", user_api_key_dict=_PROXY_ADMIN),
        agent_endpoints.make_agent_public(agent_id="agent-2", user_api_key_dict=_PROXY_ADMIN),
    )

    assert sorted(rows.stored("litellm_settings")["public_agent_groups"]) == ["agent-1", "agent-2"]


@pytest.mark.asyncio
async def test_make_agent_public_rejects_an_agent_another_replica_published(monkeypatch: pytest.MonkeyPatch) -> None:
    rows, _ = _replica_that_missed_a_publish(monkeypatch, stored_public_agent_groups=["agent-1"])

    with pytest.raises(HTTPException) as duplicate:
        await agent_endpoints.make_agent_public(agent_id="agent-1", user_api_key_dict=_PROXY_ADMIN)

    assert duplicate.value.status_code == 400
    assert "already in public agent groups" in duplicate.value.detail
    assert rows.stored("litellm_settings")["public_agent_groups"] == ["agent-1"]


@pytest.mark.asyncio
async def test_make_agent_public_refuses_a_list_the_config_file_owns(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    from litellm.proxy import proxy_server

    rows, _ = _replica_that_missed_a_publish(monkeypatch, stored_public_agent_groups=[])
    config_file: Final = tmp_path / "config.yaml"
    config_file.write_text("litellm_settings:\n  public_agent_groups: [agent-1]\n")
    monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", str(config_file))
    await proxy_server.proxy_config.get_config()

    with pytest.raises(HTTPException) as already_public:
        await agent_endpoints.make_agent_public(agent_id="agent-1", user_api_key_dict=_PROXY_ADMIN)
    with pytest.raises(HTTPException) as refused:
        await agent_endpoints.make_agent_public(agent_id="agent-2", user_api_key_dict=_PROXY_ADMIN)

    assert already_public.value.status_code == 400
    assert "already in public agent groups" in already_public.value.detail
    assert refused.value.status_code == 400
    assert refused.value.detail["keys"] == ["public_agent_groups"]
    assert rows.stored("litellm_settings")["public_agent_groups"] == []


@pytest.mark.asyncio
async def test_make_agent_public_appends_to_the_config_file_without_db_settings(
    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
    _replica_that_missed_a_publish(monkeypatch, stored_public_agent_groups=[], store_model_in_db=False)
    config_file: Final = tmp_path / "config.yaml"
    config_file.write_text("litellm_settings:\n  public_agent_groups: [agent-1]\n")
    monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", str(config_file))

    response: Final = await agent_endpoints.make_agent_public(agent_id="agent-2", user_api_key_dict=_PROXY_ADMIN)
    with pytest.raises(HTTPException) as duplicate:
        await agent_endpoints.make_agent_public(agent_id="agent-1", user_api_key_dict=_PROXY_ADMIN)

    assert response["public_agent_groups"] == ["agent-1", "agent-2"]
    assert yaml.safe_load(config_file.read_text())["litellm_settings"]["public_agent_groups"] == ["agent-1", "agent-2"]
    assert duplicate.value.status_code == 400


@pytest.mark.parametrize("enabled, claim_field, expected", [(True, "azp", True), (False, "azp", False), (True, None, False)])
def test_jwt_authentication_status_does_not_require_virtual_keys(
    monkeypatch: pytest.MonkeyPatch, enabled: bool, claim_field: str | None, expected: bool
) -> None:
    from litellm.caching.dual_cache import DualCache
    from litellm.proxy import proxy_server
    from litellm.proxy._types import LiteLLM_JWTAuth
    from litellm.proxy.auth.handle_jwt import JWTHandler

    handler: Final = JWTHandler()
    handler.update_environment(None, DualCache(), LiteLLM_JWTAuth(agent_id_jwt_field=claim_field))
    monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": enabled})
    monkeypatch.setattr(proxy_server, "jwt_handler", handler)
    agent: Final = _sample_agent_response()
    response: Final = agent_endpoints._redact_sensitive_agent_fields((agent,), is_admin=True)[0]
    assert response.jwt_auth_configured is expected
    assert agent.jwt_auth_configured is False


def test_identity_providers_require_configured_issuer_and_audience(monkeypatch: pytest.MonkeyPatch) -> None:
    from litellm.caching.dual_cache import DualCache
    from litellm.proxy import proxy_server
    from litellm.proxy._types import LiteLLM_JWTAuth
    from litellm.proxy.auth.handle_jwt import JWTHandler

    handler: Final = JWTHandler()
    handler.update_environment(None, DualCache(), LiteLLM_JWTAuth())
    monkeypatch.setattr(proxy_server, "jwt_handler", handler)
    monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": True})
    monkeypatch.setenv("JWT_ISSUER", "https://issuer.example")
    monkeypatch.delenv("JWT_AUDIENCE", raising=False)
    assert client.get("/v1/agents/identity/providers").json() == []
    monkeypatch.setenv("JWT_AUDIENCE", "gateway")
    response: Final = client.get("/v1/agents/identity/providers")
    assert response.status_code == 200
    assert response.json() == ["https://issuer.example"]
    forbidden: Final = _make_app_with_role(LitellmUserRoles.INTERNAL_USER).get("/v1/agents/identity/providers")
    assert forbidden.status_code == 403


def test_identity_evidence_is_persisted_and_never_taken_from_runtime_metadata(monkeypatch: pytest.MonkeyPatch) -> None:
    from litellm.proxy import proxy_server
    from litellm.types.proxy.agent_identity import AgentIdentityBinding

    binding: Final = AgentIdentityBinding(
        agent_id="bound",
        provider="microsoft_entra",
        tenant_id="11111111-1111-4111-8111-111111111111",
        client_id="22222222-2222-4222-8222-222222222222",
        issuer="https://issuer.example",
        revision="revision-one",
    )
    bound: Final = AgentResponse(
        agent_id="bound",
        agent_name="Readable name",
        agent_card_params={},
        identity=binding,
        identity_managed=True,
        litellm_params={"last_authenticated_at": "forged-proof"},
    )
    database: Final = MagicMock()
    database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=bound)
    monkeypatch.setattr(proxy_server, "prisma_client", database)
    pending: Final = client.get("/v1/agents/bound/identity")
    assert pending.status_code == 200
    assert pending.json()["last_authenticated_at"] is None
    verified_binding: Final = binding.model_copy(
        update={"last_authenticated_at": datetime(2026, 1, 1, tzinfo=timezone.utc)}
    )
    database.writer_db.litellm_agentstable.find_unique.return_value = bound.model_copy(update={"identity": verified_binding})
    verified: Final = client.get("/v1/agents/bound/identity")
    assert verified.json()["last_authenticated_at"] == "2026-01-01T00:00:00Z"
    assert verified.json()["identity"]["client_id"] == binding.client_id
    database.writer_db.litellm_agentstable.find_unique.return_value = None
    assert client.get("/v1/agents/missing/identity").status_code == 404
    database.writer_db.litellm_agentstable.find_unique.side_effect = RuntimeError("unavailable")
    assert client.get("/v1/agents/bound/identity").status_code == 503


@pytest.mark.parametrize("enabled", [True, False])
def test_identity_providers_honor_issuer_specific_audiences_and_global_fallback(
    monkeypatch: pytest.MonkeyPatch, enabled: bool
) -> None:
    from litellm.caching.dual_cache import DualCache
    from litellm.proxy import proxy_server
    from litellm.proxy._types import JWTIssuerConfig, LiteLLM_JWTAuth
    from litellm.proxy.auth.handle_jwt import JWTHandler

    handler: Final = JWTHandler()
    handler.update_environment(
        None,
        DualCache(),
        LiteLLM_JWTAuth(
            issuers=[
                JWTIssuerConfig(issuer="https://scoped.example", audience="gateway"),
                JWTIssuerConfig(issuer="https://unscoped.example", disable_audience_validation=True),
            ]
        ),
    )
    monkeypatch.setattr(proxy_server, "jwt_handler", handler)
    monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": enabled})
    monkeypatch.setenv("JWT_ISSUER", "https://global.example")
    monkeypatch.setenv("JWT_AUDIENCE", "gateway")
    assert client.get("/v1/agents/identity/providers").json() == (
        ["https://scoped.example", "https://global.example"] if enabled else []
    )
    monkeypatch.setenv("JWT_ISSUER", "https://unscoped.example")
    assert client.get("/v1/agents/identity/providers").json() == (["https://scoped.example"] if enabled else [])


@pytest.mark.parametrize("change", ({"execution_mode": "delegated"}, {"execution_mode": "both"}))
def test_mode_only_edit_requires_the_existing_identity_sso_tenant(
    monkeypatch: pytest.MonkeyPatch, change: PatchAgentRequest
) -> None:
    from tests.unit.proxy.agent_endpoints.test_managed_identity import BINDING, TENANT, managed_agent

    monkeypatch.setattr(agent_endpoints, "_trusted_agent_issuers", lambda: (BINDING.issuer,))
    monkeypatch.delenv("MICROSOFT_TENANT", raising=False)
    monkeypatch.setenv("MICROSOFT_CLIENT_ID", "gateway-client")
    with pytest.raises(HTTPException, match="Delegated agents require Microsoft SSO"):
        agent_endpoints._validate_managed_identity_request(change, managed_agent())
    monkeypatch.setenv("MICROSOFT_TENANT", TENANT)
    agent_endpoints._validate_managed_identity_request(change, managed_agent())


def test_identity_only_edit_preserves_delegated_mode_validation(monkeypatch: pytest.MonkeyPatch) -> None:
    from tests.unit.proxy.agent_endpoints.test_managed_identity import BINDING, managed_agent

    monkeypatch.setattr(agent_endpoints, "_trusted_agent_issuers", lambda: (BINDING.issuer,))
    monkeypatch.delenv("MICROSOFT_TENANT", raising=False)
    configuration: Final = BINDING.model_dump(
        exclude={"agent_id", "issuer", "revision", "last_authenticated_at", "active"}
    )
    delegated: Final = managed_agent().model_copy(update={"execution_mode": "delegated"})
    with pytest.raises(HTTPException, match="Delegated agents require Microsoft SSO"):
        agent_endpoints._validate_managed_identity_request({"identity": configuration}, delegated)

_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.writer_db = mock_prisma.db
            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


@pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY])
@pytest.mark.parametrize("path", ["/v1/agents", "/v1/agents/agent-123"])
def test_agent_identity_configuration_is_only_returned_to_admins(role, path, monkeypatch):
    from litellm.proxy.agent_endpoints import agent_registry

    binding = AgentIdentityBinding(
        agent_id="agent-123", provider="microsoft_entra", tenant_id="tenant", client_id="client",
        issuer="https://login.microsoftonline.com/tenant/v2.0", revision="revision",
    )
    agent = _sample_agent_response().model_copy(update={"identity": binding})
    registry = MagicMock()
    registry.get_agent_by_id.return_value = agent
    registry.get_agent_list.return_value = [agent]
    registry.ids_for_agent.return_value = frozenset({agent.agent_id})
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry)
    monkeypatch.setattr(agent_registry, "global_agent_registry", registry)
    monkeypatch.setattr(
        "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.resolve_agent_access",
        AsyncMock(return_value=RestrictedAgentAccess(frozenset({agent.agent_id}))),
    )
    with patch("litellm.proxy.proxy_server.prisma_client") as prisma:
        prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
        prisma.db.litellm_agentstable.find_many = AsyncMock(return_value=[])
        prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
        prisma.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
        response = _make_app_with_role(role).get(path, headers={"Authorization": "Bearer k"})
    assert response.status_code == 200
    payload = response.json()[0] if path == "/v1/agents" else response.json()
    assert payload["identity"] == (binding.model_dump(mode="json") if role == LitellmUserRoles.PROXY_ADMIN else None)
    assert agent.identity == binding


@pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.INTERNAL_USER])
def test_agent_detail_cache_miss_preserves_admin_identity_visibility(role, monkeypatch):
    binding = AgentIdentityBinding(
        agent_id="agent-123", provider="microsoft_entra", tenant_id="tenant", client_id="client",
        issuer="https://login.microsoftonline.com/tenant/v2.0", revision="revision",
    )
    agent = _sample_agent_response()
    registry = MagicMock()
    registry.get_agent_by_id.return_value = None
    registry.ids_for_agent.return_value = frozenset({agent.agent_id})
    monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", registry)
    monkeypatch.setattr(
        "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed",
        AsyncMock(return_value=True),
    )

    async def load_row(*, where, include):
        assert where == {"agent_id": agent.agent_id}
        return agent.model_copy(update={"identity": binding if include.get("identity") else None})

    with patch("litellm.proxy.proxy_server.prisma_client") as prisma:
        prisma.db.litellm_agentstable.find_unique = AsyncMock(side_effect=load_row)
        prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
        response = _make_app_with_role(role).get("/v1/agents/agent-123")
    assert response.status_code == 200
    assert response.json()["identity"] == (binding.model_dump(mode="json") if role == LitellmUserRoles.PROXY_ADMIN else None)


@pytest.mark.parametrize("trusted", [False, True])
def test_invalid_identity_and_untrusted_tenant_cannot_be_registered(
    monkeypatch: pytest.MonkeyPatch, trusted: bool
) -> None:
    from tests.unit.proxy.agent_endpoints.test_managed_identity import BINDING

    configuration: Final = BINDING.model_dump(
        exclude={"agent_id", "issuer", "revision", "last_authenticated_at", "active"}
    )
    monkeypatch.setattr(agent_endpoints, "_trusted_agent_issuers", lambda: (BINDING.issuer,) if trusted else ())
    request: Final = {"identity": {**configuration, "client_id": "invalid"} if trusted else configuration}
    message: Final = "Invalid Entra identity configuration" if trusted else "Configure trusted JWT issuer"
    with pytest.raises(HTTPException, match=message) as failure:
        agent_endpoints._validate_managed_identity_request(request)
    assert failure.value.status_code == 400