litellm/litellm/proxy/agent_endpoints/endpoints.py

"""
Agent endpoints for registering + discovering agents via LiteLLM.

Follows the A2A Spec.

1. Register an agent via POST `/v1/agents`
2. Discover agents via GET `/v1/agents`
3. Get specific agent via GET `/v1/agents/{agent_id}`
"""

import asyncio
import os
import uuid
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Annotated, Final, TypedDict

from fastapi import APIRouter, Depends, HTTPException, Query, Request
from typing_extensions import ReadOnly, Required, assert_never

import litellm
from litellm._logging import verbose_proxy_logger
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
from litellm.proxy._types import (
    CommonProxyErrors,
    LitellmUserRoles,
    UserAPIKeyAuth,
    user_api_key_has_admin_view,
)
from litellm.proxy.a2a.agent_card import (
    SUPPORTED_A2A_PROTOCOL_VERSIONS,
    merge_agent_card,
    normalize_protocol_version,
)
from litellm.proxy.agent_endpoints.agent_registry import (
    AgentIdWhere,
    parse_agent_kill_switch,
    parse_agent_litellm_params,
    redact_sensitive_agent_litellm_params,
)
from litellm.proxy.agent_endpoints.agent_search import (
    DEFAULT_AGENT_SEARCH_TOP_K,
    AgentSearchEmbeddingFailed,
    AgentSearchHits,
    AgentSearchNotConfigured,
    global_agent_search_index,
    search_agents,
)
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents
from litellm.proxy.agent_endpoints.kill_switch import (
    KillSwitchAuditLogWriter,
    KillSwitchHttpClient,
    build_kill_switch_audit_log,
    default_kill_switch_audit_log_writer,
    default_kill_switch_http_client,
    fire_kill_switch,
    redact_kill_switch,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user
from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
from litellm.proxy.utils import get_custom_url
from litellm.types.agents import (
    AgentCard,
    AgentConfig,
    AgentKeySummary,
    AgentKillSwitchConfig,
    AgentKillSwitchResult,
    AgentMakePublicResponse,
    AgentResponse,
    MakeAgentsPublicRequest,
    PatchAgentRequest,
)
from litellm.types.llms.custom_http import httpxSpecialProvider
from litellm.types.proxy.management_endpoints.common_daily_activity import (
    DailySpendMetadata,
    SpendAnalyticsPaginatedResponse,
)


def _proxy_base_url(http_request: Request) -> str:
    """Return the proxy's public base URL, preferring PROXY_BASE_URL when set."""
    return get_custom_url(str(http_request.base_url), route=None)


def _validate_protocol_version(upstream_card: AgentCard | None) -> None:
    """Reject an agent card pinning an unsupported A2A protocol version."""
    version: Final = upstream_card.get("protocolVersion") if upstream_card else None
    if version is not None and normalize_protocol_version(version) is None:
        raise HTTPException(
            status_code=400,
            detail=(
                f"Unsupported protocolVersion '{version}'. "
                f"Supported versions: {', '.join(SUPPORTED_A2A_PROTOCOL_VERSIONS)}."
            ),
        )


def _build_merged_agent_card(
    upstream_card: AgentCard | None,
    *,
    agent_id: str,
    http_request: Request,
    agent_name: str | None = None,
) -> dict[str, object]:
    """Apply the LiteLLM-fronting merge to ``upstream_card`` for ``agent_id``."""
    proxy_base: Final = _proxy_base_url(http_request)
    _validate_protocol_version(upstream_card)
    # Prefer a card-supplied ``name`` (the discovery UI exposes an editable
    # "Name (shown to API clients)" field that flows into
    # ``agent_card_params.name``) over the internal ``agent_name`` identifier.
    # Fall back to ``agent_name`` only when the card itself has no name.
    card_name: Final = upstream_card.get("name") if upstream_card else None
    return merge_agent_card(
        upstream_card,
        proxy_url=f"{proxy_base}/a2a/{agent_id}",
        proxy_base_url=proxy_base,
        name=card_name or agent_name,
    )


router: Final = APIRouter()


async def _attach_keys_to_agents(agents: Sequence[AgentResponse], prisma_client) -> None:
    """Attach each agent's virtual keys, derived from the key table's agent_id
    foreign key. Mirrors how spend is joined into the agent response so the UI
    never has to cross-reference a full key dump client-side. Only non-secret
    fields are exposed (alias, masked key_name, hashed token)."""
    from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry

    agent_ids: Final = tuple(
        alias_id for agent in agents for alias_id in global_agent_registry.ids_for_agent(agent.agent_id)
    )
    if not agent_ids:
        return
    key_rows: Final = await prisma_client.replica_db.litellm_verificationtoken.find_many(
        where={"agent_id": {"in": agent_ids}},
    )
    keys_by_agent: Final[dict[str, list[AgentKeySummary]]] = {}
    for row in key_rows:
        keys_by_agent.setdefault(row.agent_id, []).append(
            AgentKeySummary(
                token=row.token,
                key_alias=row.key_alias,
                key_name=row.key_name,
            )
        )
    for agent in agents:
        matched_keys = [
            key_summary
            for alias_id in global_agent_registry.ids_for_agent(agent.agent_id)
            for key_summary in keys_by_agent.get(alias_id) or ()
        ]
        agent.keys = matched_keys or None


def _redact_agent_litellm_params_dict(
    litellm_params: Mapping[str, object],
) -> dict[str, object]:  # mutable-ok: AgentResponse.litellm_params is declared as a plain dict, not Mapping
    """Type-narrowing wrapper: a dict in always yields a dict back from
    ``redact_sensitive_agent_litellm_params``, which the function's general
    (possible-JSON-string, possibly-None) signature can't express."""
    return dict(  # mutable-ok: AgentResponse.litellm_params is declared as a plain dict, not Mapping
        parse_agent_litellm_params(redact_sensitive_agent_litellm_params(litellm_params))
    )


def _redact_sensitive_agent_fields(
    agents: Sequence[AgentResponse],
    *,
    is_admin: bool,
) -> list[AgentResponse]:
    """
    Return copies of the given agents with credential-bearing litellm_params
    values and kill-switch auth secrets replaced by a fixed marker (never
    returned to ANY caller, admin included) and, for non-admin callers,
    virtual-key, header and kill-switch fields stripped entirely. The original
    objects are not modified.
    """
    redacted: Final[list[AgentResponse]] = []
    for agent in agents:
        copy = agent.model_copy(deep=True)
        if not is_admin:
            copy.static_headers = None
            copy.extra_headers = None
            copy.keys = None
            copy.kill_switch = None
        if copy.litellm_params:
            copy.litellm_params = _redact_agent_litellm_params_dict(copy.litellm_params)
        copy.kill_switch = redact_kill_switch(copy.kill_switch)
        redacted.append(copy)
    return redacted


def _check_agent_management_permission(user_api_key_dict: UserAPIKeyAuth) -> None:
    """
    Raises HTTP 403 if the caller does not have permission to create, update,
    or delete agents.  Only PROXY_ADMIN users are allowed to perform these
    write operations.
    """
    if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
        raise HTTPException(
            status_code=403,
            detail={
                "error": f"Only proxy admins can create, update, or delete agents. Your role={user_api_key_dict.user_role}"
            },
        )


AGENT_HEALTH_CHECK_TIMEOUT_SECONDS: Final = float(os.environ.get("LITELLM_AGENT_HEALTH_CHECK_TIMEOUT", "5.0"))
AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS = float(os.environ.get("LITELLM_AGENT_HEALTH_CHECK_GATHER_TIMEOUT", "30.0"))


class _AgentHealthResult(TypedDict, total=False):
    agent_id: Required[str]
    healthy: Required[bool]
    error: str


async def _check_agent_url_health(
    agent: AgentResponse,
) -> _AgentHealthResult:
    """
    Perform a GET request against the agent's URL and return the health result.

    Returns a dict with ``agent_id``, ``healthy`` (bool), and an optional
    ``error`` message.
    """
    url: Final = (agent.agent_card_params or {}).get("url")
    if not url:
        return {"agent_id": agent.agent_id, "healthy": True}

    try:
        client: Final = get_async_httpx_client(
            llm_provider=httpxSpecialProvider.AgentHealthCheck,
            params={"timeout": AGENT_HEALTH_CHECK_TIMEOUT_SECONDS},
        )
        response: Final = await client.get(url)
        if response.status_code >= 500:
            return {
                "agent_id": agent.agent_id,
                "healthy": False,
                "error": f"HTTP {response.status_code}",
            }
        return {"agent_id": agent.agent_id, "healthy": True}
    except Exception as exc:
        return {
            "agent_id": agent.agent_id,
            "healthy": False,
            "error": str(exc),
        }


class _AgentSearchErrorDetail(TypedDict):
    error: ReadOnly[str]
    message: ReadOnly[str]


def _agent_search_error(status_code: int, error: str, message: str) -> HTTPException:
    detail: Final[_AgentSearchErrorDetail] = {"error": error, "message": message}
    return HTTPException(status_code=status_code, detail=detail)


async def _rank_agents_by_query(
    query: str, agents: Sequence[AgentResponse], top_k: int, user_api_key_dict: UserAPIKeyAuth
) -> tuple[AgentResponse, ...]:
    from litellm.proxy.proxy_server import llm_router, proxy_logging_obj

    outcome: Final = await search_agents(
        query=query,
        agents=agents,
        top_k=top_k,
        router=llm_router,
        embedding_model=litellm.agent_search_embedding_model,
        index=global_agent_search_index,
        user_api_key_dict=user_api_key_dict,
        proxy_logging_obj=proxy_logging_obj,
    )
    match outcome:
        case AgentSearchHits(hits):
            return tuple(hit.agent.model_copy(update=MappingProxyType({"search_score": hit.score})) for hit in hits)
        case AgentSearchNotConfigured(reason):
            raise _agent_search_error(400, "agent_search_not_configured", reason)
        case AgentSearchEmbeddingFailed(reason):
            raise _agent_search_error(503, "agent_search_unavailable", reason)
        case _:
            assert_never(outcome)


@router.get(
    "/v1/agents",
    tags=["[beta] A2A Agents"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=list[AgentResponse],
)
async def get_agents(
    request: Request,
    health_check: bool = Query(
        False,
        description="When true, performs a GET request to each agent's URL. Agents with reachable URLs (HTTP status < 500) and agents without a URL are returned; unreachable agents are filtered out.",
    ),
    query: Annotated[
        str | None,
        Query(
            min_length=1,
            description="Describe the task in natural language to rank the agents you can reach by semantic similarity over their name, description, and skills. Each result carries a search_score. Requires litellm_settings.agent_search_embedding_model.",
        ),
    ] = None,
    top_k: Annotated[
        int,
        Query(ge=1, le=100, description="With query: the maximum number of ranked agents to return."),
    ] = DEFAULT_AGENT_SEARCH_TOP_K,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),  # Used for auth
):
    """
    Example usage:
    ```
    curl -X GET "http://localhost:4000/v1/agents" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer your-key" \
    ```

    Pass `?health_check=true` to filter out agents whose URL is unreachable:
    ```
    curl -X GET "http://localhost:4000/v1/agents?health_check=true" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer your-key" \
    ```

    Pass `?query=<task>` to get the best matching agents ranked by semantic similarity:
    ```
    curl -X GET "http://localhost:4000/v1/agents?query=translate+a+PDF+document&top_k=5" \
      -H "Content-Type: application/json" \
      -H "Authorization: Bearer your-key" \
    ```

    Returns: List[AgentResponse]

    """
    await check_feature_access_for_user(user_api_key_dict, "agents")

    from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry

    try:
        returned_agents: Sequence[AgentResponse] = await accessible_agents(user_api_key_dict)

        # Fetch current spend from DB for all returned agents
        from litellm.proxy.proxy_server import prisma_client

        if prisma_client is not None:
            agent_ids: Final = tuple(
                alias_id
                for agent in returned_agents
                for alias_id in global_agent_registry.ids_for_agent(agent.agent_id)
            )
            if agent_ids:
                db_agents: Final = await agents_table(prisma_client).find_many(
                    where={"agent_id": {"in": agent_ids}},
                )
                spend_map: Final = {a.agent_id: a.spend for a in db_agents}
                for agent in returned_agents:
                    matched_spends = tuple(
                        spend_map[alias_id]
                        for alias_id in global_agent_registry.ids_for_agent(agent.agent_id)
                        if alias_id in spend_map
                    )
                    if matched_spends:
                        agent.spend = sum(matched_spends)
                await _attach_keys_to_agents(returned_agents, prisma_client)

        # add is_public field to each agent - we do it this way, to allow setting config agents as public
        for agent in returned_agents:
            if agent.litellm_params is None:
                agent.litellm_params = {}
            agent.litellm_params["is_public"] = litellm.public_agent_groups is not None and not (
                global_agent_registry.ids_for_agent(agent.agent_id).isdisjoint(litellm.public_agent_groups)
            )

        # litellm_params secrets are always redacted; keys/headers stay
        # admin-only.
        is_admin: Final = (
            user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
            or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
        )
        returned_agents = _redact_sensitive_agent_fields(returned_agents, is_admin=is_admin)

        if health_check:
            agents_with_url: Final = [agent for agent in returned_agents if (agent.agent_card_params or {}).get("url")]
            agents_without_url = [agent for agent in returned_agents if not (agent.agent_card_params or {}).get("url")]
            try:
                health_results: Sequence[_AgentHealthResult] = await asyncio.wait_for(
                    asyncio.gather(*[_check_agent_url_health(agent) for agent in agents_with_url]),
                    timeout=AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS,
                )
            except asyncio.TimeoutError:
                verbose_proxy_logger.warning(
                    "Agent health check gather timed out after %s seconds",
                    AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS,
                )
                health_results = [
                    {
                        "agent_id": agent.agent_id,
                        "healthy": False,
                        "error": "Health check timed out",
                    }
                    for agent in agents_with_url
                ]
            healthy_ids: Final = {result["agent_id"] for result in health_results if result["healthy"]}
            returned_agents = [agent for agent in agents_with_url if agent.agent_id in healthy_ids] + agents_without_url

        if query is None:
            return returned_agents
        return await _rank_agents_by_query(query, returned_agents, top_k, user_api_key_dict)
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("litellm.proxy.agent_endpoints.get_agents(): Exception occurred - %s", e)
        raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"})


#### CRUD ENDPOINTS FOR AGENTS ####

from litellm.proxy.agent_endpoints.agent_registry import (
    agents_table,
)
from litellm.proxy.agent_endpoints.agent_registry import (
    global_agent_registry as AGENT_REGISTRY,
)


@router.post(
    "/v1/agents",
    tags=["[beta] A2A Agents"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=AgentResponse,
)
async def create_agent(
    request: AgentConfig,
    http_request: Request,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Create a new agent

    Example Request:
    ```bash
    curl -X POST "http://localhost:4000/v1/agents" \\
        -H "Authorization: Bearer <your_api_key>" \\
        -H "Content-Type: application/json" \\
        -d '{
            "agent_name": "my-custom-agent",
            "agent_card_params": {
                "protocolVersion": "1.0",
                "name": "Hello World Agent",
                "description": "Just a hello world agent",
                "url": "http://localhost:9999/",
                "version": "1.0.0",
                "defaultInputModes": ["text"],
                "defaultOutputModes": ["text"],
                "capabilities": {
                    "streaming": true
                },
                "skills": [
                    {
                        "id": "hello_world",
                        "name": "Returns hello world",
                        "description": "just returns hello world",
                        "tags": ["hello world"],
                        "examples": ["hi", "hello world"]
                    }
                ]
            },
            "litellm_params": {
                "make_public": true
            }
        }'
    ```
    """
    await check_feature_access_for_user(user_api_key_dict, "agents")

    from litellm.proxy.proxy_server import prisma_client

    _check_agent_management_permission(user_api_key_dict)

    if prisma_client is None:
        raise HTTPException(status_code=500, detail="Prisma client not initialized")

    try:
        # Get the user ID from the API key auth
        created_by: Final = user_api_key_dict.user_id or "unknown"

        # check for naming conflicts
        existing_agent: Final = AGENT_REGISTRY.get_agent_by_name(agent_name=request.get("agent_name"))
        if existing_agent is not None:
            raise HTTPException(
                status_code=400,
                detail=f"Agent with name {request.get('agent_name')} already exists",
            )

        # Apply the LiteLLM-fronting merge only when the admin actually
        # provided an agent card. Plain chat/LLM agents register without
        # ``agent_card_params``, and synthesising a default A2A card for them
        # would advertise capabilities (``supportedInterfaces``, security
        # schemes, default skills) the agent doesn't actually expose.
        upstream_card: Final = request.get("agent_card_params")
        agent_to_create: AgentConfig = request
        new_agent_id: str | None = None
        if upstream_card is not None:
            # Pre-generate the agent_id so the merged card can reference it
            # in ``supportedInterfaces`` before the DB row exists.
            new_agent_id = str(uuid.uuid4())
            merged_card: Final = _build_merged_agent_card(
                upstream_card,
                agent_id=new_agent_id,
                http_request=http_request,
                agent_name=request.get("agent_name"),
            )
            agent_to_create = {**request, "agent_card_params": merged_card}

        result: Final = await AGENT_REGISTRY.add_agent_to_db(
            agent=agent_to_create,
            prisma_client=prisma_client,
            created_by=created_by,
            agent_id=new_agent_id,
        )

        agent_name: Final = result.agent_name
        agent_id: Final = result.agent_id

        # Also register in memory
        try:
            AGENT_REGISTRY.register_agent(agent_config=result)
            verbose_proxy_logger.info("Successfully registered agent '%s' (ID: %s) in memory", agent_name, agent_id)
        except Exception as reg_error:
            verbose_proxy_logger.warning(
                "Failed to register agent '%s' (ID: %s) in memory: %s", agent_name, agent_id, reg_error
            )

        # The caller is a proxy admin (enforced above); litellm_params
        # secrets are still never echoed back in the response.
        return _redact_sensitive_agent_fields((result,), is_admin=True)[0]

    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("Error adding agent to db: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.get(
    "/v1/agents/{agent_id}",
    tags=["[beta] A2A Agents"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=AgentResponse,
)
async def get_agent_by_id(
    agent_id: str,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Get a specific agent by ID

    Example Request:
    ```bash
    curl -X GET "http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000" \\
        -H "Authorization: Bearer <your_api_key>"
    ```
    """
    await check_feature_access_for_user(user_api_key_dict, "agents")

    if not user_api_key_has_admin_view(user_api_key_dict):
        from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
            AgentRequestHandler,
        )

        is_allowed = await AgentRequestHandler.is_agent_allowed(agent_id=agent_id, user_api_key_auth=user_api_key_dict)
        if not is_allowed:
            raise HTTPException(
                status_code=403,
                detail=f"Agent '{agent_id}' is not allowed for your key/team. Contact proxy admin for access.",
            )

    from litellm.proxy.proxy_server import prisma_client

    if prisma_client is None:
        raise HTTPException(status_code=500, detail="Prisma client not initialized")

    try:
        agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id)
        if agent is None:
            agent_row: Final = await agents_table(prisma_client).find_unique(
                where={"agent_id": agent_id},
                include={"object_permission": True},
            )
            if agent_row is not None:
                agent_dict: Final = agent_row.model_dump()
                if agent_row.object_permission is not None:
                    try:
                        agent_dict["object_permission"] = agent_row.object_permission.model_dump()
                    except Exception:
                        agent_dict["object_permission"] = agent_row.object_permission.dict()
                agent = AgentResponse(**agent_dict)
        else:
            # Agent found in memory — refresh spend from DB
            db_row: Final = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
            if db_row is not None:
                agent.spend = db_row.spend

        if agent is None:
            raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")

        await _attach_keys_to_agents([agent], prisma_client)

        # litellm_params secrets are always redacted; keys/headers stay
        # admin-only.
        is_admin = (
            user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
            or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
        )
        agent = _redact_sensitive_agent_fields((agent,), is_admin=is_admin)[0]

        return agent
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("Error getting agent from db: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.put(
    "/v1/agents/{agent_id}",
    tags=["[beta] A2A Agents"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=AgentResponse,
)
async def update_agent(
    agent_id: str,
    request: AgentConfig,
    http_request: Request,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Update an existing agent

    Example Request:
    ```bash
    curl -X PUT "http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000" \\
        -H "Authorization: Bearer <your_api_key>" \\
        -H "Content-Type: application/json" \\
        -d '{
            "agent_name": "updated-agent",
            "agent_card_params": {
                "protocolVersion": "1.0",
                "name": "Updated Agent",
                "description": "Updated description",
                "url": "http://localhost:9999/",
                "version": "1.1.0",
                "defaultInputModes": ["text"],
                "defaultOutputModes": ["text"],
                "capabilities": {
                    "streaming": true
                },
                "skills": []
            },
            "litellm_params": {
                "make_public": false
            }
        }'
    ```
    """
    await check_feature_access_for_user(user_api_key_dict, "agents")

    from litellm.proxy.proxy_server import prisma_client

    _check_agent_management_permission(user_api_key_dict)

    if prisma_client is None:
        raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)

    try:
        # Check if agent exists
        existing_agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
        if existing_agent is not None:
            existing_agent = dict(existing_agent)

        if existing_agent is None:
            raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")

        # Get the user ID from the API key auth
        updated_by: Final = user_api_key_dict.user_id or "unknown"

        # Re-apply the LiteLLM-fronting merge — an update is a re-registration,
        # so any new upstream card the admin pasted must go through the same
        # transformation as initial create. Plain agents without an
        # ``agent_card_params`` skip the merge so we don't synthesise an A2A
        # card for them.
        upstream_card: Final = request.get("agent_card_params")
        agent_to_update: AgentConfig = request
        if upstream_card is not None:
            merged_card: Final = _build_merged_agent_card(
                upstream_card,
                agent_id=agent_id,
                http_request=http_request,
                agent_name=request.get("agent_name"),
            )
            agent_to_update = {**request, "agent_card_params": merged_card}

        result: Final = await AGENT_REGISTRY.update_agent_in_db(
            agent_id=agent_id,
            agent=agent_to_update,
            prisma_client=prisma_client,
            updated_by=updated_by,
        )

        # deregister in memory
        AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name"))
        # register in memory
        AGENT_REGISTRY.register_agent(agent_config=result)

        verbose_proxy_logger.info(
            "Successfully updated agent '%s' (ID: %s) in memory", existing_agent.get("agent_name"), agent_id
        )

        return _redact_sensitive_agent_fields((result,), is_admin=True)[0]
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("Error updating agent: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.patch(
    "/v1/agents/{agent_id}",
    tags=["[beta] A2A Agents"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=AgentResponse,
)
async def patch_agent(
    agent_id: str,
    request: PatchAgentRequest,
    http_request: Request,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Update an existing agent

    Example Request:
    ```bash
    curl -X PATCH "http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000" \\
        -H "Authorization: Bearer <your_api_key>" \\
        -H "Content-Type: application/json" \\
        -d '{
            "agent_name": "updated-agent",
            "agent_card_params": {
                "protocolVersion": "1.0",
                "name": "Updated Agent",
                "description": "Updated description",
                "url": "http://localhost:9999/",
                "version": "1.1.0",
                "defaultInputModes": ["text"],
                "defaultOutputModes": ["text"],
                "capabilities": {
                    "streaming": true
                },
                "skills": []
            },
            "litellm_params": {
                "make_public": false
            }
        }'
    ```
    """
    await check_feature_access_for_user(user_api_key_dict, "agents")

    from litellm.proxy.proxy_server import prisma_client

    _check_agent_management_permission(user_api_key_dict)

    if prisma_client is None:
        raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)

    try:
        # Check if agent exists
        existing_agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
        if existing_agent is not None:
            existing_agent = dict(existing_agent)

        if existing_agent is None:
            raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")

        # Get the user ID from the API key auth
        updated_by: Final = user_api_key_dict.user_id or "unknown"

        # Re-merge only when the patch actually touches agent_card_params; a
        # patch updating just litellm_params/rate limits (``agent_card_params``
        # omitted) shouldn't rewrite the stored card. An explicitly provided
        # ``agent_card_params`` — even an empty dict — still goes through the
        # merge so LiteLLM applies its security schemes and supported
        # interfaces instead of storing a bare card.
        patch_payload: PatchAgentRequest = request
        upstream_card: Final = request.get("agent_card_params")
        if upstream_card is not None:
            merged_card: Final = _build_merged_agent_card(
                upstream_card,
                agent_id=agent_id,
                http_request=http_request,
                agent_name=request.get("agent_name"),
            )
            patch_payload = {**request, "agent_card_params": merged_card}

        result: Final = await AGENT_REGISTRY.patch_agent_in_db(
            agent_id=agent_id,
            agent=patch_payload,
            prisma_client=prisma_client,
            updated_by=updated_by,
        )

        # deregister in memory
        AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name"))
        # register in memory
        AGENT_REGISTRY.register_agent(agent_config=result)

        verbose_proxy_logger.info(
            "Successfully updated agent '%s' (ID: %s) in memory", existing_agent.get("agent_name"), agent_id
        )

        return _redact_sensitive_agent_fields((result,), is_admin=True)[0]
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("Error updating agent: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.delete(
    "/v1/agents/{agent_id}",
    tags=["Agents"],
    dependencies=[Depends(user_api_key_auth)],
)
async def delete_agent(
    agent_id: str,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Delete an agent

    Example Request:
    ```bash
    curl -X DELETE "http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000" \\
        -H "Authorization: Bearer <your_api_key>"
    ```

    Example Response:
    ```json
    {
        "message": "Agent 123e4567-e89b-12d3-a456-426614174000 deleted successfully"
    }
    ```
    """
    await check_feature_access_for_user(user_api_key_dict, "agents")

    from litellm.proxy.proxy_server import prisma_client

    _check_agent_management_permission(user_api_key_dict)

    if prisma_client is None:
        raise HTTPException(status_code=500, detail="Prisma client not initialized")

    try:
        # Check if agent exists
        existing_agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
        if existing_agent is not None:
            existing_agent = dict[str, object](existing_agent)

        if existing_agent is None:
            raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found in DB.")

        await AGENT_REGISTRY.delete_agent_from_db(agent_id=agent_id, prisma_client=prisma_client)

        AGENT_REGISTRY.deregister_agent(agent_name=existing_agent.get("agent_name"))

        return {"message": f"Agent {agent_id} deleted successfully"}
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("Error deleting agent: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.post(
    "/v1/agents/{agent_id}/kill_switch",
    tags=["[beta] A2A Agents"],  # mutable-ok: fastapi types tags as list[str | Enum]
    dependencies=(Depends(user_api_key_auth),),
    response_model=AgentKillSwitchResult,
)
async def trigger_agent_kill_switch(
    agent_id: str,
    user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
    http_client: Annotated[KillSwitchHttpClient, Depends(default_kill_switch_http_client)],
    audit_log_writer: Annotated[KillSwitchAuditLogWriter, Depends(default_kill_switch_audit_log_writer)],
):
    """
    Fire the agent's configured kill switch webhook. Proxy admin only.

    LiteLLM only makes the configured HTTP call and reports what came back; it
    does not change the agent's state in LiteLLM. Returns 200 when the webhook
    answered 2xx, 502 with the same result body otherwise. Every attempt is
    written to the audit log as a `kill_switch_fired` row against the agent.

    Example Request:
    ```bash
    curl -X POST "http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000/kill_switch" \\
        -H "Authorization: Bearer <your_api_key>"
    ```
    """
    from litellm.proxy.proxy_server import litellm_proxy_admin_name

    await check_feature_access_for_user(user_api_key_dict, "agents")
    _check_agent_management_permission(user_api_key_dict)

    resolved: Final = await _resolve_agent_kill_switch(agent_id)
    if resolved is None:
        raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")
    resolved_agent_id, config = resolved
    if config is None:
        raise HTTPException(status_code=400, detail=f"Agent with ID {agent_id} has no kill_switch configured")

    result: Final = await fire_kill_switch(agent_id=resolved_agent_id, config=config, http_client=http_client)
    await audit_log_writer(
        build_kill_switch_audit_log(
            result=result,
            user_api_key_dict=user_api_key_dict,
            litellm_proxy_admin_name=litellm_proxy_admin_name,
        )
    )
    if not result.succeeded:
        raise HTTPException(status_code=502, detail=result.model_dump())
    return result


async def _resolve_agent_kill_switch(agent_id: str) -> tuple[str, AgentKillSwitchConfig | None] | None:
    """The DB row wins over this replica's in-memory registry so a trigger never fires a webhook another
    replica has since changed; config.yaml agents have no row and fall back to the registry."""
    from litellm.proxy.proxy_server import prisma_client

    if prisma_client is not None:
        where: Final[AgentIdWhere] = {"agent_id": agent_id}
        row: Final = await agents_table(prisma_client).find_unique(where=where)
        if row is not None:
            return row.agent_id, parse_agent_kill_switch(row.kill_switch)

    agent: Final = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id)
    if agent is None:
        return None
    return agent.agent_id, agent.kill_switch


@router.post(
    "/v1/agents/{agent_id}/make_public",
    tags=["[beta] A2A Agents"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=AgentMakePublicResponse,
)
async def make_agent_public(
    agent_id: str,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Make an agent publicly discoverable

    Example Request:
    ```bash
    curl -X POST "http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000/make_public" \\
        -H "Authorization: Bearer <your_api_key>" \\
        -H "Content-Type: application/json"
    ```

    Example Response:
    ```json
    {
        "agent_id": "123e4567-e89b-12d3-a456-426614174000",
        "agent_name": "my-custom-agent",
        "litellm_params": {
            "make_public": true
        },
        "agent_card_params": {...},
        "created_at": "2025-11-15T10:30:00Z",
        "updated_at": "2025-11-15T10:35:00Z",
        "created_by": "user123",
        "updated_by": "user123"
    }
    ```
    """
    from litellm.proxy.proxy_server import prisma_client

    if prisma_client is None:
        raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)

    try:
        # Update the public model groups
        import litellm
        from litellm.proxy.agent_endpoints.agent_registry import (
            global_agent_registry as AGENT_REGISTRY,
        )
        from litellm.proxy.proxy_server import proxy_config

        # Check if user has admin permissions
        if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
            raise HTTPException(
                status_code=403,
                detail={
                    "error": f"Only proxy admins can update public model groups. Your role={user_api_key_dict.user_role}"
                },
            )

        agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id)
        if agent is None:
            # check if agent exists in DB
            agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
            if agent is not None:
                agent = AgentResponse(**agent.model_dump())

            if agent is None:
                raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")

        config: Final = await proxy_config.get_config()

        current_public_agent_groups: Final = list(litellm.public_agent_groups or [])
        if not AGENT_REGISTRY.ids_for_agent(agent.agent_id).isdisjoint(current_public_agent_groups):
            raise HTTPException(
                status_code=400,
                detail=f"Agent with name {agent.agent_name} already in public agent groups",
            )
        updated_public_agent_groups: Final = [*current_public_agent_groups, agent.agent_id]

        if "litellm_settings" not in config or config["litellm_settings"] is None:
            config["litellm_settings"] = {}

        config["litellm_settings"]["public_agent_groups"] = updated_public_agent_groups

        await proxy_config.save_config(new_config=config)

        litellm.public_agent_groups = updated_public_agent_groups

        verbose_proxy_logger.debug(
            "Updated public agent groups to: %s by user: %s", updated_public_agent_groups, user_api_key_dict.user_id
        )

        return {
            "message": "Successfully updated public agent groups",
            "public_agent_groups": updated_public_agent_groups,
            "updated_by": user_api_key_dict.user_id,
        }
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("Error making agent public: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.post(
    "/v1/agents/make_public",
    tags=["[beta] A2A Agents"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=AgentMakePublicResponse,
)
async def make_agents_public(
    request: MakeAgentsPublicRequest,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Make multiple agents publicly discoverable

    Example Request:
    ```bash
    curl -X POST "http://localhost:4000/v1/agents/make_public" \\
        -H "Authorization: Bearer <your_api_key>" \\
        -H "Content-Type: application/json" \\
        -d '{
            "agent_ids": ["123e4567-e89b-12d3-a456-426614174000", "123e4567-e89b-12d3-a456-426614174001"]
        }'
    ```

    Example Response:
    ```json
    {
        "agent_id": "123e4567-e89b-12d3-a456-426614174000",
        "agent_name": "my-custom-agent",
        "litellm_params": {
            "make_public": true
        },
        "agent_card_params": {...},
        "created_at": "2025-11-15T10:30:00Z",
        "updated_at": "2025-11-15T10:35:00Z",
        "created_by": "user123",
        "updated_by": "user123"
    }
    ```
    """
    from litellm.proxy.proxy_server import prisma_client

    if prisma_client is None:
        raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value)

    try:
        # Update the public model groups
        import litellm
        from litellm.proxy.agent_endpoints.agent_registry import (
            global_agent_registry as AGENT_REGISTRY,
        )
        from litellm.proxy.proxy_server import proxy_config

        # Load existing config
        config: Final = await proxy_config.get_config()
        # Check if user has admin permissions
        if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
            raise HTTPException(
                status_code=403,
                detail={
                    "error": f"Only proxy admins can update public model groups. Your role={user_api_key_dict.user_role}"
                },
            )

        if litellm.public_agent_groups is None:
            litellm.public_agent_groups = []

        for agent_id in request.agent_ids:
            agent = AGENT_REGISTRY.get_agent_by_id(agent_id=agent_id)
            if agent is None:
                # check if agent exists in DB
                agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
                if agent is not None:
                    agent = AgentResponse(**agent.model_dump())

                if agent is None:
                    raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")

        litellm.public_agent_groups = request.agent_ids

        # Update config with new settings
        if "litellm_settings" not in config or config["litellm_settings"] is None:
            config["litellm_settings"] = {}

        config["litellm_settings"]["public_agent_groups"] = litellm.public_agent_groups

        # Save the updated config
        await proxy_config.save_config(new_config=config)

        verbose_proxy_logger.debug(
            "Updated public agent groups to: %s by user: %s", litellm.public_agent_groups, user_api_key_dict.user_id
        )

        return {
            "message": "Successfully updated public agent groups",
            "public_agent_groups": litellm.public_agent_groups,
            "updated_by": user_api_key_dict.user_id,
        }
    except HTTPException:
        raise
    except Exception as e:
        verbose_proxy_logger.exception("Error making agent public: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@router.get(
    "/agent/daily/activity",
    tags=["Agent Management"],
    dependencies=[Depends(user_api_key_auth)],
    response_model=SpendAnalyticsPaginatedResponse,
)
async def get_agent_daily_activity(
    agent_ids: str | None = None,
    start_date: str | None = None,
    end_date: str | None = None,
    model: str | None = None,
    api_key: str | None = None,
    page: int = 1,
    page_size: int = 10,
    exclude_agent_ids: str | None = None,
    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
    """
    Get daily activity for specific agents or all accessible agents.
    """
    await check_feature_access_for_user(user_api_key_dict, "agents")

    from litellm.proxy.proxy_server import prisma_client

    if prisma_client is None:
        raise HTTPException(
            status_code=500,
            detail={"error": CommonProxyErrors.db_not_connected_error.value},
        )

    agent_ids_list = agent_ids.split(",") if agent_ids else None
    exclude_agent_ids_list: list[str] | None = None
    if exclude_agent_ids:
        exclude_agent_ids_list = exclude_agent_ids.split(",") if exclude_agent_ids else None

    # Without scoping, an empty `agent_ids` query returned every agent's
    # spend/token rows on the proxy. Restrict non-admin callers to the
    # agents they're permitted to invoke (or that they created), and
    # intersect their explicit `agent_ids` filter with the same allowlist.
    from litellm.proxy.agent_endpoints.auth.agent_permission_handler import (
        AgentRequestHandler,
        RestrictedAgentAccess,
        UnrestrictedAgentAccess,
    )
    from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view

    where_condition: Final[dict[str, object]] = {}
    if not _user_has_admin_view(user_api_key_dict):
        permitted_agent_ids: list[str] = []
        # An unrestricted caller is not "see everything" for activity scoping. Fall
        # back to the agents the caller created so they cannot enumerate other
        # tenants' agents.
        # Guard against `user_id is None`: a literal None in Prisma
        # `where={"created_by": None}` resolves to ``created_by IS NULL``
        # and would expose every ownerless agent's rows.
        match await AgentRequestHandler.resolve_agent_access(user_api_key_auth=user_api_key_dict):
            case RestrictedAgentAccess(allowed_agent_ids):
                permitted_agent_ids = list(allowed_agent_ids)
            case UnrestrictedAgentAccess():
                if user_api_key_dict.user_id is not None:
                    owned_records: Final = await agents_table(prisma_client).find_many(
                        where={"created_by": user_api_key_dict.user_id}
                    )
                    permitted_agent_ids = [a.agent_id for a in owned_records]

        if agent_ids_list:
            permitted_agent_id_set: Final = set(permitted_agent_ids)
            agent_ids_list = [aid for aid in agent_ids_list if aid in permitted_agent_id_set]
        else:
            agent_ids_list = list(permitted_agent_ids)

        # No accessible agents → return an empty page without querying.
        if not agent_ids_list:
            return SpendAnalyticsPaginatedResponse(
                results=[],
                metadata=DailySpendMetadata(
                    total_spend=0.0,
                    total_prompt_tokens=0,
                    total_completion_tokens=0,
                    total_tokens=0,
                    total_api_requests=0,
                    total_successful_requests=0,
                    total_failed_requests=0,
                    total_cache_read_input_tokens=0,
                    total_cache_creation_input_tokens=0,
                    total_compression_saved_tokens=0,
                    page=page,
                    total_pages=0,
                    has_more=False,
                ),
            )

    if agent_ids_list:
        where_condition["agent_id"] = {"in": list(agent_ids_list)}

    agent_records: Final = await agents_table(prisma_client).find_many(where=where_condition)
    agent_metadata: Final[Mapping[str, dict[str, object]]] = {
        agent.agent_id: {"agent_name": agent.agent_name} for agent in agent_records
    }

    return await get_daily_activity(
        prisma_client=prisma_client,
        table_name="litellm_dailyagentspend",
        entity_id_field="agent_id",
        entity_id=agent_ids_list,
        entity_metadata_field=agent_metadata,
        exclude_entity_ids=exclude_agent_ids_list,
        start_date=start_date,
        end_date=end_date,
        model=model,
        api_key=api_key,
        page=page,
        page_size=page_size,
    )