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

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

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

from litellm.constants import REDACTED_BY_LITELM_STRING
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.agent_endpoints import endpoints as agent_endpoints
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
    RestrictedAgentAccess,
    UnrestrictedAgentAccess,
)
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):
        """An unrestricted viewer 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 = UnrestrictedAgentAccess()
        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
        return ProxyConfig()._update_config_fields(
            current_config=config,
            param_name="litellm_settings",
            db_param_value=db_param_value,
        )

    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"]