mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-01 02:02:20 +00:00
feat(agents): register Entra identities through the agent API
This commit is contained in:
Joshua Valluru
parent
0a03f3d0c8
commit
afa0d37119
7 changed files with 1070 additions and 173 deletions
|
|
@ -2378,6 +2378,19 @@
|
|||
"title": "Agent Name",
|
||||
"type": "string"
|
||||
},
|
||||
"enabled": {
|
||||
"title": "Enabled",
|
||||
"type": "boolean"
|
||||
},
|
||||
"execution_mode": {
|
||||
"enum": [
|
||||
"autonomous",
|
||||
"delegated",
|
||||
"both"
|
||||
],
|
||||
"title": "Execution Mode",
|
||||
"type": "string"
|
||||
},
|
||||
"extra_headers": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
@ -2392,6 +2405,16 @@
|
|||
],
|
||||
"title": "Extra Headers"
|
||||
},
|
||||
"identity": {
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/components/schemas/EntraIdentityConfig"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
},
|
||||
"kill_switch": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
@ -2470,8 +2493,7 @@
|
|||
}
|
||||
},
|
||||
"required": [
|
||||
"agent_name",
|
||||
"agent_card_params"
|
||||
"agent_name"
|
||||
],
|
||||
"title": "AgentConfig",
|
||||
"type": "object"
|
||||
|
|
@ -3561,6 +3583,60 @@
|
|||
"title": "DailySpendMetadata",
|
||||
"type": "object"
|
||||
},
|
||||
"EntraIdentityConfig": {
|
||||
"additionalProperties": false,
|
||||
"properties": {
|
||||
"client_id": {
|
||||
"title": "Client Id",
|
||||
"type": "string"
|
||||
},
|
||||
"provider": {
|
||||
"const": "microsoft_entra",
|
||||
"title": "Provider",
|
||||
"type": "string"
|
||||
},
|
||||
"required_roles": {
|
||||
"default": [],
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"title": "Required Roles",
|
||||
"type": "array"
|
||||
},
|
||||
"required_scopes": {
|
||||
"default": [
|
||||
"user_impersonation"
|
||||
],
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"title": "Required Scopes",
|
||||
"type": "array"
|
||||
},
|
||||
"service_principal_id": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Service Principal Id"
|
||||
},
|
||||
"tenant_id": {
|
||||
"title": "Tenant Id",
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"provider",
|
||||
"tenant_id",
|
||||
"client_id"
|
||||
],
|
||||
"title": "EntraIdentityConfig",
|
||||
"type": "object"
|
||||
},
|
||||
"HTTPAuthSecurityScheme": {
|
||||
"description": "Defines a security scheme using HTTP authentication.",
|
||||
"properties": {
|
||||
|
|
@ -3710,6 +3786,54 @@
|
|||
"title": "MakeAgentsPublicRequest",
|
||||
"type": "object"
|
||||
},
|
||||
"ManagedAgentIdentityStatus": {
|
||||
"properties": {
|
||||
"enabled": {
|
||||
"default": true,
|
||||
"title": "Enabled",
|
||||
"type": "boolean"
|
||||
},
|
||||
"execution_mode": {
|
||||
"default": "autonomous",
|
||||
"enum": [
|
||||
"autonomous",
|
||||
"delegated",
|
||||
"both"
|
||||
],
|
||||
"title": "Execution Mode",
|
||||
"type": "string"
|
||||
},
|
||||
"identity": {
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/components/schemas/AgentIdentityBinding"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
},
|
||||
"identity_managed": {
|
||||
"default": false,
|
||||
"title": "Identity Managed",
|
||||
"type": "boolean"
|
||||
},
|
||||
"last_authenticated_at": {
|
||||
"anyOf": [
|
||||
{
|
||||
"format": "date-time",
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "Last Authenticated At"
|
||||
}
|
||||
},
|
||||
"title": "ManagedAgentIdentityStatus",
|
||||
"type": "object"
|
||||
},
|
||||
"MetricWithMetadata": {
|
||||
"properties": {
|
||||
"api_key_breakdown": {
|
||||
|
|
@ -3910,6 +4034,19 @@
|
|||
"title": "Agent Name",
|
||||
"type": "string"
|
||||
},
|
||||
"enabled": {
|
||||
"title": "Enabled",
|
||||
"type": "boolean"
|
||||
},
|
||||
"execution_mode": {
|
||||
"enum": [
|
||||
"autonomous",
|
||||
"delegated",
|
||||
"both"
|
||||
],
|
||||
"title": "Execution Mode",
|
||||
"type": "string"
|
||||
},
|
||||
"extra_headers": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
@ -3924,6 +4061,16 @@
|
|||
],
|
||||
"title": "Extra Headers"
|
||||
},
|
||||
"identity": {
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/components/schemas/EntraIdentityConfig"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
]
|
||||
},
|
||||
"kill_switch": {
|
||||
"anyOf": [
|
||||
{
|
||||
|
|
@ -4445,6 +4592,36 @@
|
|||
]
|
||||
}
|
||||
},
|
||||
"/v1/agents/identity/providers": {
|
||||
"get": {
|
||||
"operationId": "get_agent_identity_providers_v1_agents_identity_providers_get",
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"items": {
|
||||
"type": "string"
|
||||
},
|
||||
"title": "Response Get Agent Identity Providers V1 Agents Identity Providers Get",
|
||||
"type": "array"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Get Agent Identity Providers",
|
||||
"tags": [
|
||||
"agents"
|
||||
]
|
||||
}
|
||||
},
|
||||
"/v1/agents/make_public": {
|
||||
"post": {
|
||||
"description": "Make multiple agents publicly discoverable\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/v1/agents/make_public\" \\\n -H \"Authorization: Bearer <your_api_key>\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"agent_ids\": [\"123e4567-e89b-12d3-a456-426614174000\", \"123e4567-e89b-12d3-a456-426614174001\"]\n }'\n```\n\nExample Response:\n```json\n{\n \"agent_id\": \"123e4567-e89b-12d3-a456-426614174000\",\n \"agent_name\": \"my-custom-agent\",\n \"litellm_params\": {\n \"make_public\": true\n },\n \"agent_card_params\": {...},\n \"created_at\": \"2025-11-15T10:30:00Z\",\n \"updated_at\": \"2025-11-15T10:35:00Z\",\n \"created_by\": \"user123\",\n \"updated_by\": \"user123\"\n}\n```",
|
||||
|
|
@ -4696,6 +4873,53 @@
|
|||
]
|
||||
}
|
||||
},
|
||||
"/v1/agents/{agent_id}/identity": {
|
||||
"get": {
|
||||
"operationId": "get_agent_identity_status_v1_agents__agent_id__identity_get",
|
||||
"parameters": [
|
||||
{
|
||||
"in": "path",
|
||||
"name": "agent_id",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"title": "Agent Id",
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ManagedAgentIdentityStatus"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Successful Response"
|
||||
},
|
||||
"422": {
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/HTTPValidationError"
|
||||
}
|
||||
}
|
||||
},
|
||||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Get Agent Identity Status",
|
||||
"tags": [
|
||||
"agents"
|
||||
]
|
||||
}
|
||||
},
|
||||
"/v1/agents/{agent_id}/kill_switch": {
|
||||
"post": {
|
||||
"description": "Fire the agent's configured kill switch webhook. Proxy admin only.\n\nLiteLLM only makes the configured HTTP call and reports what came back; it\ndoes not change the agent's state in LiteLLM. Returns 200 when the webhook\nanswered 2xx, 502 with the same result body otherwise. Every attempt is\nwritten to the audit log as a `kill_switch_fired` row against the agent.\n\nExample Request:\n```bash\ncurl -X POST \"http://localhost:4000/v1/agents/123e4567-e89b-12d3-a456-426614174000/kill_switch\" \\\n -H \"Authorization: Bearer <your_api_key>\"\n```",
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from datetime import datetime, timezone
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, TypedDict
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
|
|
@ -14,13 +15,16 @@ from litellm.constants import REDACTED_BY_LITELM_STRING
|
|||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
|
||||
from litellm.proxy.agent_endpoints.kill_switch import restore_kill_switch
|
||||
from litellm.proxy.agent_endpoints.managed_identity import managed_write_fields, raise_identity_failure
|
||||
from litellm.proxy.management_helpers.object_permission_utils import (
|
||||
handle_update_object_permission_common,
|
||||
prepare_object_permission_upsert,
|
||||
)
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.base_repository import is_unique_violation
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import AgentsRepository, ObjectPermissionRepository
|
||||
from litellm.types.agents import AgentConfig, AgentKillSwitchConfig, AgentResponse, PatchAgentRequest
|
||||
from litellm.types.proxy.agent_identity import AgentIdentityFailure
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
|
@ -135,6 +139,39 @@ def object_permission_table(
|
|||
return table
|
||||
|
||||
|
||||
class AgentPermissionWrite(TypedDict, total=False):
|
||||
create: ReadOnly[Mapping[str, object]]
|
||||
update: ReadOnly[Mapping[str, object]]
|
||||
|
||||
|
||||
async def _permission_write(
|
||||
incoming: Mapping[str, object],
|
||||
existing_id: str | None,
|
||||
client: PrismaClient,
|
||||
) -> AgentPermissionWrite | None:
|
||||
raw: Final = incoming.get("object_permission")
|
||||
if raw is None:
|
||||
return None
|
||||
permission: Final = _AGENT_PARAMS_ADAPTER.validate_python(raw)
|
||||
prepared: Final = await prepare_object_permission_upsert(permission, existing_id, client)
|
||||
if existing_id is None:
|
||||
created: Final[AgentPermissionWrite] = {"create": prepared.record}
|
||||
return created
|
||||
updated: Final[AgentPermissionWrite] = {"update": prepared.record}
|
||||
return updated
|
||||
|
||||
|
||||
def _managed_fields(
|
||||
incoming: Mapping[str, object],
|
||||
existing: AgentResponse | None,
|
||||
updated_by: str,
|
||||
) -> Mapping[str, object]:
|
||||
result: Final = managed_write_fields(incoming, existing, updated_by)
|
||||
if isinstance(result, AgentIdentityFailure):
|
||||
raise_identity_failure(result, 400)
|
||||
return result
|
||||
|
||||
|
||||
def _dump_agent_params(raw: Mapping[str, object]) -> dict[str, object]:
|
||||
model_dump: Final[Callable[[], dict[str, object]] | None] = getattr(raw, "model_dump", None)
|
||||
if model_dump is not None:
|
||||
|
|
@ -552,11 +589,7 @@ class AgentRegistry:
|
|||
agent_card_params_dict: Final[dict[str, object]] = _dump_agent_params(agent_card_params_obj)
|
||||
agent_card_params: Final[str] = safe_dumps(agent_card_params_dict)
|
||||
|
||||
# Handle object_permission (MCP tool access for agent)
|
||||
object_permission_id: str | None = None
|
||||
if agent.get("object_permission") is not None:
|
||||
agent_copy: Final = dict(agent)
|
||||
object_permission_id = await handle_update_object_permission_common(agent_copy, None, prisma_client)
|
||||
permission_write: Final = await _permission_write(agent, None, prisma_client)
|
||||
|
||||
# Serialize static_headers
|
||||
static_headers_obj: Final = agent.get("static_headers")
|
||||
|
|
@ -583,8 +616,8 @@ class AgentRegistry:
|
|||
create_data["extra_headers"] = extra_headers_val
|
||||
if access_group_ids_val is not None:
|
||||
create_data["access_group_ids"] = tuple(dict.fromkeys(access_group_ids_val))
|
||||
if object_permission_id is not None:
|
||||
create_data["object_permission_id"] = object_permission_id
|
||||
if permission_write is not None:
|
||||
create_data["object_permission"] = permission_write
|
||||
|
||||
for rate_field in (
|
||||
"tpm_limit",
|
||||
|
|
@ -598,31 +631,46 @@ class AgentRegistry:
|
|||
|
||||
# Create agent in DB
|
||||
created_agent: Final = await agents_table(prisma_client).create(
|
||||
data=create_data,
|
||||
include={"object_permission": True},
|
||||
data={**create_data, **_managed_fields(agent, None, created_by)},
|
||||
include={"object_permission": True, "identity": True},
|
||||
)
|
||||
|
||||
created_agent_dict: Final = created_agent.model_dump()
|
||||
if created_agent.object_permission is not None:
|
||||
try:
|
||||
created_agent_dict["object_permission"] = created_agent.object_permission.model_dump()
|
||||
except Exception:
|
||||
created_agent_dict["object_permission"] = created_agent.object_permission.dict()
|
||||
return AgentResponse(**created_agent_dict)
|
||||
return AgentResponse.model_validate(created_agent.model_dump())
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise Exception(f"Error adding agent to DB: {e}")
|
||||
if is_unique_violation(e):
|
||||
raise HTTPException(409, "Agent name or Entra application is already registered") from e
|
||||
raise
|
||||
|
||||
async def delete_agent_from_db(self, agent_id: str, prisma_client: PrismaClient) -> Mapping[str, object]:
|
||||
"""
|
||||
Delete an agent from the database
|
||||
"""
|
||||
try:
|
||||
deleted_agent: Final = await agents_table(prisma_client).delete(where={"agent_id": agent_id})
|
||||
from prisma.types import (
|
||||
LiteLLM_AgentsTableWhereUniqueInput,
|
||||
LiteLLM_RetiredAgentCreateInput,
|
||||
LiteLLM_RetiredAgentUpsertInput,
|
||||
LiteLLM_RetiredAgentWhereUniqueInput,
|
||||
LiteLLM_VerificationTokenWhereInput,
|
||||
)
|
||||
|
||||
where: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_id": agent_id}
|
||||
async with prisma_client.tx() as tx:
|
||||
existing: Final = await tx.litellm_agentstable.find_unique(where=where)
|
||||
if existing is None:
|
||||
raise ValueError(f"Agent not found, passed agent_id={agent_id}")
|
||||
if existing.identity_managed:
|
||||
history_where: Final[LiteLLM_RetiredAgentWhereUniqueInput] = {"original_agent_id": agent_id}
|
||||
history_create: Final = LiteLLM_RetiredAgentCreateInput(original_agent_id=agent_id)
|
||||
history_data: Final[LiteLLM_RetiredAgentUpsertInput] = {"create": history_create, "update": {}}
|
||||
await tx.litellm_retiredagent.upsert(where=history_where, data=history_data)
|
||||
keys_where: Final[LiteLLM_VerificationTokenWhereInput] = {"agent_id": agent_id}
|
||||
await tx.litellm_verificationtoken.delete_many(where=keys_where)
|
||||
deleted_agent: Final = await tx.litellm_agentstable.delete(where=where)
|
||||
if deleted_agent is None:
|
||||
raise ValueError(f"Agent not found, passed agent_id={agent_id}")
|
||||
return dict(deleted_agent)
|
||||
except Exception as e:
|
||||
raise Exception(f"Error deleting agent from DB: {e}")
|
||||
return deleted_agent.model_dump()
|
||||
|
||||
async def patch_agent_in_db(
|
||||
self,
|
||||
|
|
@ -646,7 +694,9 @@ class AgentRegistry:
|
|||
The patched agent
|
||||
"""
|
||||
try:
|
||||
existing_record: Final = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
|
||||
existing_record: Final = await agents_table(prisma_client).find_unique(
|
||||
where={"agent_id": agent_id}, include={"identity": True}
|
||||
)
|
||||
if existing_record is None:
|
||||
raise Exception(f"Agent with ID {agent_id} not found")
|
||||
existing_agent: Final[Mapping[str, object]] = dict(existing_record)
|
||||
|
|
@ -683,37 +733,31 @@ class AgentRegistry:
|
|||
if "extra_headers" in agent:
|
||||
extra_headers_value: Final = agent.get("extra_headers")
|
||||
update_data["extra_headers"] = extra_headers_value if extra_headers_value is not None else []
|
||||
if agent.get("object_permission") is not None:
|
||||
agent_copy: Final = dict(augment_agent)
|
||||
existing_object_permission_id: Final = existing_record.object_permission_id
|
||||
object_permission_id: Final = await handle_update_object_permission_common(
|
||||
agent_copy,
|
||||
existing_object_permission_id,
|
||||
prisma_client,
|
||||
)
|
||||
if object_permission_id is not None:
|
||||
update_data["object_permission_id"] = object_permission_id
|
||||
permission_write: Final = await _permission_write(
|
||||
agent, existing_record.object_permission_id, prisma_client
|
||||
)
|
||||
if permission_write is not None:
|
||||
update_data["object_permission"] = permission_write
|
||||
# Patch agent in DB
|
||||
patched_agent: Final = await agents_table(prisma_client).update(
|
||||
where={"agent_id": agent_id},
|
||||
data={
|
||||
**update_data,
|
||||
**_managed_fields(agent, AgentResponse.model_validate(existing_record.model_dump()), updated_by),
|
||||
"updated_by": updated_by,
|
||||
"updated_at": datetime.now(timezone.utc),
|
||||
},
|
||||
include={"object_permission": True},
|
||||
include={"object_permission": True, "identity": True},
|
||||
)
|
||||
if patched_agent is None:
|
||||
raise ValueError(f"Agent not found, passed agent_id={agent_id}")
|
||||
patched_agent_dict: Final = patched_agent.model_dump()
|
||||
if patched_agent.object_permission is not None:
|
||||
try:
|
||||
patched_agent_dict["object_permission"] = patched_agent.object_permission.model_dump()
|
||||
except Exception:
|
||||
patched_agent_dict["object_permission"] = patched_agent.object_permission.dict()
|
||||
return AgentResponse(**patched_agent_dict)
|
||||
return AgentResponse.model_validate(patched_agent.model_dump())
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise Exception(f"Error patching agent in DB: {e}")
|
||||
if is_unique_violation(e):
|
||||
raise HTTPException(409, "Agent name or Entra application is already registered") from e
|
||||
raise
|
||||
|
||||
async def update_agent_in_db(
|
||||
self,
|
||||
|
|
@ -733,7 +777,7 @@ class AgentRegistry:
|
|||
# caller echoed back redacted (or omitted) rather than persisting
|
||||
# the marker -- or nothing -- over the real stored credential.
|
||||
existing_row: Final = await agents_table(prisma_client).find_unique(
|
||||
where={"agent_id": agent_id} # mutable-ok: prisma's query builder rejects a Mapping/MappingProxyType
|
||||
where={"agent_id": agent_id}, include={"identity": True}
|
||||
)
|
||||
existing_litellm_params: Final = parse_agent_litellm_params(
|
||||
existing_row.litellm_params if existing_row is not None else None
|
||||
|
|
@ -784,37 +828,35 @@ class AgentRegistry:
|
|||
if _val is not None:
|
||||
update_data[rate_field] = _val
|
||||
|
||||
if agent.get("object_permission") is not None:
|
||||
existing_object_permission_id: Final = (
|
||||
existing_row.object_permission_id if existing_row is not None else None
|
||||
)
|
||||
agent_copy: Final = dict(agent)
|
||||
object_permission_id: Final = await handle_update_object_permission_common(
|
||||
agent_copy,
|
||||
existing_object_permission_id,
|
||||
prisma_client,
|
||||
)
|
||||
if object_permission_id is not None:
|
||||
update_data["object_permission_id"] = object_permission_id
|
||||
permission_write: Final = await _permission_write(
|
||||
agent, existing_row.object_permission_id if existing_row is not None else None, prisma_client
|
||||
)
|
||||
if permission_write is not None:
|
||||
update_data["object_permission"] = permission_write
|
||||
|
||||
# Update agent in DB
|
||||
updated_agent: Final = await agents_table(prisma_client).update(
|
||||
where={"agent_id": agent_id},
|
||||
data=update_data,
|
||||
include={"object_permission": True},
|
||||
data={
|
||||
**update_data,
|
||||
**_managed_fields(
|
||||
agent,
|
||||
AgentResponse.model_validate(existing_row.model_dump()) if existing_row else None,
|
||||
updated_by,
|
||||
),
|
||||
},
|
||||
include={"object_permission": True, "identity": True},
|
||||
)
|
||||
|
||||
if updated_agent is None:
|
||||
raise ValueError(f"Agent not found, passed agent_id={agent_id}")
|
||||
updated_agent_dict: Final = updated_agent.model_dump()
|
||||
if updated_agent.object_permission is not None:
|
||||
try:
|
||||
updated_agent_dict["object_permission"] = updated_agent.object_permission.model_dump()
|
||||
except Exception:
|
||||
updated_agent_dict["object_permission"] = updated_agent.object_permission.dict()
|
||||
return AgentResponse(**updated_agent_dict)
|
||||
return AgentResponse.model_validate(updated_agent.model_dump())
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise Exception(f"Error updating agent in DB: {e}")
|
||||
if is_unique_violation(e):
|
||||
raise HTTPException(409, "Agent name or Entra application is already registered") from e
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
async def get_all_agents_from_db(
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from types import MappingProxyType
|
|||
from typing import Annotated, Final, TypedDict
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query, Request
|
||||
from pydantic import ValidationError
|
||||
from typing_extensions import ReadOnly, Required, assert_never
|
||||
|
||||
import litellm
|
||||
|
|
@ -47,6 +48,8 @@ from litellm.proxy.agent_endpoints.agent_search import (
|
|||
search_agents,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.auth.agent_permission_handler import accessible_agents
|
||||
from litellm.proxy.agent_endpoints.identity import reject_legacy_identity
|
||||
from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore
|
||||
from litellm.proxy.agent_endpoints.kill_switch import (
|
||||
KillSwitchAuditLogWriter,
|
||||
KillSwitchHttpClient,
|
||||
|
|
@ -56,6 +59,7 @@ from litellm.proxy.agent_endpoints.kill_switch import (
|
|||
fire_kill_switch,
|
||||
redact_kill_switch,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure
|
||||
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
|
||||
|
|
@ -72,6 +76,12 @@ from litellm.types.agents import (
|
|||
PatchAgentRequest,
|
||||
)
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
from litellm.types.proxy.agent_identity import (
|
||||
AgentIdentityBinding,
|
||||
AgentIdentityFailure,
|
||||
EntraIdentityConfig,
|
||||
ManagedAgentIdentityStatus,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import (
|
||||
DailySpendMetadata,
|
||||
SpendAnalyticsPaginatedResponse,
|
||||
|
|
@ -178,9 +188,15 @@ def _redact_sensitive_agent_fields(
|
|||
virtual-key, header and kill-switch fields stripped entirely. The original
|
||||
objects are not modified.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import general_settings, jwt_handler
|
||||
|
||||
redacted: Final[list[AgentResponse]] = []
|
||||
for agent in agents:
|
||||
copy = agent.model_copy(deep=True)
|
||||
copy.jwt_auth_configured = bool(
|
||||
general_settings.get("enable_jwt_auth")
|
||||
and (agent.identity is not None or jwt_handler.litellm_jwtauth.agent_id_jwt_field)
|
||||
)
|
||||
if not is_admin:
|
||||
copy.static_headers = None
|
||||
copy.extra_headers = None
|
||||
|
|
@ -430,6 +446,71 @@ from litellm.proxy.agent_endpoints.agent_registry import (
|
|||
)
|
||||
|
||||
|
||||
def _trusted_agent_issuers() -> tuple[str, ...]:
|
||||
from litellm.proxy.proxy_server import general_settings, jwt_handler
|
||||
|
||||
if not general_settings.get("enable_jwt_auth"):
|
||||
return ()
|
||||
configured: Final = jwt_handler.litellm_jwtauth.issuers or ()
|
||||
issuer: Final = os.getenv("JWT_ISSUER")
|
||||
global_issuers: Final = (
|
||||
(issuer,)
|
||||
if issuer and os.getenv("JWT_AUDIENCE") and not any(item.issuer == issuer for item in configured)
|
||||
else ()
|
||||
)
|
||||
return (
|
||||
tuple(item.issuer for item in configured if item.audience and not item.disable_audience_validation)
|
||||
+ global_issuers
|
||||
)
|
||||
|
||||
|
||||
def _validate_managed_identity_request(
|
||||
request: AgentConfig | PatchAgentRequest, existing: AgentResponse | None = None
|
||||
) -> None:
|
||||
raw: Final = request.get("identity") if "identity" in request else existing.identity if existing else None
|
||||
if raw is None:
|
||||
return
|
||||
try:
|
||||
identity: Final = raw if isinstance(raw, AgentIdentityBinding) else EntraIdentityConfig.model_validate(raw)
|
||||
except ValidationError as exc:
|
||||
raise HTTPException(400, "Invalid Entra identity configuration") from exc
|
||||
if identity.issuer not in _trusted_agent_issuers():
|
||||
raise HTTPException(400, "Configure trusted JWT issuer and audience validation for this Entra tenant first")
|
||||
if request.get("execution_mode", existing.execution_mode if existing else "autonomous") != "autonomous":
|
||||
if os.getenv("MICROSOFT_TENANT") != identity.tenant_id or not os.getenv("MICROSOFT_CLIENT_ID"):
|
||||
raise HTTPException(400, "Delegated agents require Microsoft SSO for the same trusted tenant")
|
||||
|
||||
|
||||
@router.get("/v1/agents/identity/providers", response_model=tuple[str, ...], tags=("[beta] A2A Agents",))
|
||||
async def get_agent_identity_providers(
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
) -> tuple[str, ...]:
|
||||
_check_agent_management_permission(user_api_key_dict)
|
||||
return _trusted_agent_issuers()
|
||||
|
||||
|
||||
@router.get("/v1/agents/{agent_id}/identity", response_model=ManagedAgentIdentityStatus, tags=("[beta] A2A Agents",))
|
||||
async def get_agent_identity_status(
|
||||
agent_id: str,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
|
||||
) -> ManagedAgentIdentityStatus:
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
_check_agent_management_permission(user_api_key_dict)
|
||||
agent: Final = await AgentIdentityStore.from_client(prisma_client).agent(agent_id)
|
||||
if isinstance(agent, AgentIdentityFailure):
|
||||
raise_identity_failure(agent)
|
||||
if agent is None:
|
||||
raise HTTPException(404, "Agent not found")
|
||||
return ManagedAgentIdentityStatus(
|
||||
identity=agent.identity,
|
||||
identity_managed=agent.identity_managed,
|
||||
enabled=agent.enabled,
|
||||
execution_mode=agent.execution_mode,
|
||||
last_authenticated_at=agent.identity.last_authenticated_at if agent.identity else None,
|
||||
)
|
||||
|
||||
|
||||
@router.post(
|
||||
"/v1/agents",
|
||||
tags=["[beta] A2A Agents"],
|
||||
|
|
@ -491,6 +572,9 @@ async def create_agent(
|
|||
# Get the user ID from the API key auth
|
||||
created_by: Final = user_api_key_dict.user_id or "unknown"
|
||||
|
||||
_validate_managed_identity_request(request)
|
||||
reject_legacy_identity(request.get("litellm_params"))
|
||||
|
||||
# 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:
|
||||
|
|
@ -681,13 +765,18 @@ async def update_agent(
|
|||
|
||||
try:
|
||||
# Check if agent exists
|
||||
existing_agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
|
||||
existing_agent = await agents_table(prisma_client).find_unique(
|
||||
where={"agent_id": agent_id}, include={"identity": True}
|
||||
)
|
||||
if existing_agent is not None:
|
||||
existing_agent = dict(existing_agent)
|
||||
existing_agent = existing_agent.model_dump()
|
||||
|
||||
if existing_agent is None:
|
||||
raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")
|
||||
|
||||
_validate_managed_identity_request(request, AgentResponse.model_validate(existing_agent))
|
||||
reject_legacy_identity(request.get("litellm_params"))
|
||||
|
||||
# Get the user ID from the API key auth
|
||||
updated_by: Final = user_api_key_dict.user_id or "unknown"
|
||||
|
||||
|
|
@ -783,13 +872,18 @@ async def patch_agent(
|
|||
|
||||
try:
|
||||
# Check if agent exists
|
||||
existing_agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
|
||||
existing_agent = await agents_table(prisma_client).find_unique(
|
||||
where={"agent_id": agent_id}, include={"identity": True}
|
||||
)
|
||||
if existing_agent is not None:
|
||||
existing_agent = dict(existing_agent)
|
||||
existing_agent = existing_agent.model_dump()
|
||||
|
||||
if existing_agent is None:
|
||||
raise HTTPException(status_code=404, detail=f"Agent with ID {agent_id} not found")
|
||||
|
||||
_validate_managed_identity_request(request, AgentResponse.model_validate(existing_agent))
|
||||
reject_legacy_identity(request.get("litellm_params"))
|
||||
|
||||
# Get the user ID from the API key auth
|
||||
updated_by: Final = user_api_key_dict.user_id or "unknown"
|
||||
|
||||
|
|
@ -870,7 +964,9 @@ async def delete_agent(
|
|||
|
||||
try:
|
||||
# Check if agent exists
|
||||
existing_agent = await agents_table(prisma_client).find_unique(where={"agent_id": agent_id})
|
||||
existing_agent = await agents_table(prisma_client).find_unique(
|
||||
where={"agent_id": agent_id}, include={"identity": True}
|
||||
)
|
||||
if existing_agent is not None:
|
||||
existing_agent = dict[str, object](existing_agent)
|
||||
|
||||
|
|
|
|||
|
|
@ -7,7 +7,11 @@ from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, StrictInt, field
|
|||
from typing_extensions import ReadOnly, Required, TypedDict
|
||||
|
||||
from litellm.types.llms.base import LiteLLMPydanticObjectBase
|
||||
from litellm.types.proxy.agent_identity import AgentExecutionMode, AgentIdentityBinding
|
||||
from litellm.types.proxy.agent_identity import (
|
||||
AgentExecutionMode,
|
||||
AgentIdentityBinding,
|
||||
EntraIdentityConfig,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from a2a.types import SendMessageResponse
|
||||
|
|
@ -249,8 +253,11 @@ class AgentKillSwitchResult(BaseModel):
|
|||
|
||||
|
||||
class AgentConfig(TypedDict, total=False):
|
||||
identity: ReadOnly[EntraIdentityConfig | None]
|
||||
enabled: ReadOnly[bool]
|
||||
execution_mode: ReadOnly[AgentExecutionMode]
|
||||
agent_name: Required[str]
|
||||
agent_card_params: Required[AgentCard]
|
||||
agent_card_params: ReadOnly[AgentCard]
|
||||
litellm_params: dict[str, object] # allow for any future litellm params
|
||||
object_permission: AgentObjectPermission
|
||||
tpm_limit: int | None
|
||||
|
|
@ -264,6 +271,9 @@ class AgentConfig(TypedDict, total=False):
|
|||
|
||||
|
||||
class PatchAgentRequest(TypedDict, total=False):
|
||||
identity: ReadOnly[EntraIdentityConfig | None]
|
||||
enabled: ReadOnly[bool]
|
||||
execution_mode: ReadOnly[AgentExecutionMode]
|
||||
agent_name: str
|
||||
agent_card_params: AgentCard
|
||||
litellm_params: dict[str, object]
|
||||
|
|
|
|||
|
|
@ -2,11 +2,14 @@
|
|||
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from prisma.models import LiteLLM_AgentsTable
|
||||
|
||||
from litellm.constants import REDACTED_BY_LITELM_STRING
|
||||
from litellm.proxy.agent_endpoints.agent_registry import (
|
||||
|
|
@ -451,11 +454,11 @@ async def test_update_agent_in_db_raises_when_row_deleted_mid_update():
|
|||
registry: Final = AgentRegistry()
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(litellm_params={}, object_permission_id=None, kill_switch=None)
|
||||
return_value=_stored_agent_row(SimpleNamespace(litellm_params={}, object_permission_id=None))
|
||||
)
|
||||
mock_prisma.db.litellm_agentstable.update = AsyncMock(return_value=None)
|
||||
|
||||
with pytest.raises(Exception, match="Error updating agent in DB") as exc_info:
|
||||
with pytest.raises(Exception, match="Agent not found") as exc_info:
|
||||
await registry.update_agent_in_db(
|
||||
agent_id="agent-123",
|
||||
agent={
|
||||
|
|
@ -467,7 +470,7 @@ async def test_update_agent_in_db_raises_when_row_deleted_mid_update():
|
|||
updated_by="test-user",
|
||||
)
|
||||
|
||||
assert str(exc_info.value) == "Error updating agent in DB: Agent not found, passed agent_id=agent-123"
|
||||
assert str(exc_info.value) == "Agent not found, passed agent_id=agent-123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -476,11 +479,13 @@ async def test_patch_agent_in_db_raises_when_row_deleted_mid_update():
|
|||
registry: Final = AgentRegistry()
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value={"agent_id": "agent-123", "agent_name": "Old Agent", "object_permission_id": None}
|
||||
return_value=_stored_agent_row(
|
||||
{"agent_id": "agent-123", "agent_name": "Old Agent", "object_permission_id": None}
|
||||
)
|
||||
)
|
||||
mock_prisma.db.litellm_agentstable.update = AsyncMock(return_value=None)
|
||||
|
||||
with pytest.raises(Exception, match="Error patching agent in DB") as exc_info:
|
||||
with pytest.raises(Exception, match="Agent not found") as exc_info:
|
||||
await registry.patch_agent_in_db(
|
||||
agent_id="agent-123",
|
||||
agent={"agent_name": "Patched Agent"},
|
||||
|
|
@ -488,20 +493,43 @@ async def test_patch_agent_in_db_raises_when_row_deleted_mid_update():
|
|||
updated_by="test-user",
|
||||
)
|
||||
|
||||
assert str(exc_info.value) == "Error patching agent in DB: Agent not found, passed agent_id=agent-123"
|
||||
assert str(exc_info.value) == "Agent not found, passed agent_id=agent-123"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_agent_from_db_raises_when_row_already_gone():
|
||||
"""Prisma's delete returns None for a missing row, which dict() cannot consume."""
|
||||
async def test_delete_agent_from_db_raises_when_row_already_gone() -> None:
|
||||
registry: Final = AgentRegistry()
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_prisma.db.litellm_agentstable.delete = AsyncMock(return_value=None)
|
||||
database: Final = MagicMock()
|
||||
tx: Final = database.tx.return_value.__aenter__.return_value
|
||||
tx.litellm_agentstable.find_unique = AsyncMock(return_value=None)
|
||||
with pytest.raises(ValueError, match="Agent not found, passed agent_id=agent-123"):
|
||||
await registry.delete_agent_from_db(agent_id="agent-123", prisma_client=database)
|
||||
tx.litellm_verificationtoken.delete_many.assert_not_called()
|
||||
|
||||
with pytest.raises(Exception, match="Error deleting agent from DB") as exc_info:
|
||||
await registry.delete_agent_from_db(agent_id="agent-123", prisma_client=mock_prisma)
|
||||
|
||||
assert str(exc_info.value) == "Error deleting agent from DB: Agent not found, passed agent_id=agent-123"
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("managed", [True, False])
|
||||
async def test_agent_deletion_revokes_managed_keys_and_keeps_identity_history(managed: bool) -> None:
|
||||
registry: Final = AgentRegistry()
|
||||
database: Final = MagicMock()
|
||||
tx: Final = database.tx.return_value.__aenter__.return_value
|
||||
row: Final = _stored_agent_row({"agent_id": "agent-123", "identity_managed": managed})
|
||||
tx.litellm_agentstable.find_unique = AsyncMock(return_value=row)
|
||||
tx.litellm_agentstable.delete = AsyncMock(return_value=row)
|
||||
tx.litellm_verificationtoken.delete_many = AsyncMock(return_value=2)
|
||||
tx.litellm_retiredagent.upsert = AsyncMock()
|
||||
result: Final = await registry.delete_agent_from_db("agent-123", database)
|
||||
assert result["agent_id"] == "agent-123"
|
||||
tx.litellm_agentstable.delete.assert_awaited_once_with(where={"agent_id": "agent-123"})
|
||||
if managed:
|
||||
tx.litellm_retiredagent.upsert.assert_awaited_once_with(
|
||||
where={"original_agent_id": "agent-123"},
|
||||
data={"create": {"original_agent_id": "agent-123"}, "update": {}},
|
||||
)
|
||||
tx.litellm_verificationtoken.delete_many.assert_awaited_once_with(where={"agent_id": "agent-123"})
|
||||
else:
|
||||
tx.litellm_retiredagent.upsert.assert_not_awaited()
|
||||
tx.litellm_verificationtoken.delete_many.assert_not_awaited()
|
||||
|
||||
|
||||
# ---------- LIT-6736: agent litellm_params secret redaction ----------
|
||||
|
|
@ -729,14 +757,15 @@ async def test_update_agent_in_db_preserves_secret_when_echoed_back_redacted():
|
|||
mock_prisma: Final = MagicMock()
|
||||
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
litellm_params={
|
||||
"aws_access_key_id": SENTINEL_AWS_ACCESS_KEY_ID,
|
||||
"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY,
|
||||
"model": "bedrock/agentcore/my-agent",
|
||||
},
|
||||
object_permission_id=None,
|
||||
kill_switch=None,
|
||||
return_value=_stored_agent_row(
|
||||
SimpleNamespace(
|
||||
litellm_params={
|
||||
"aws_access_key_id": SENTINEL_AWS_ACCESS_KEY_ID,
|
||||
"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY,
|
||||
"model": "bedrock/agentcore/my-agent",
|
||||
},
|
||||
object_permission_id=None,
|
||||
)
|
||||
)
|
||||
)
|
||||
updated_agent = MagicMock()
|
||||
|
|
@ -782,10 +811,11 @@ async def test_update_agent_in_db_preserves_secret_when_key_omitted_entirely():
|
|||
mock_prisma: Final = MagicMock()
|
||||
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
litellm_params={"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY},
|
||||
object_permission_id=None,
|
||||
kill_switch=None,
|
||||
return_value=_stored_agent_row(
|
||||
SimpleNamespace(
|
||||
litellm_params={"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY},
|
||||
object_permission_id=None,
|
||||
)
|
||||
)
|
||||
)
|
||||
updated_agent = MagicMock()
|
||||
|
|
@ -824,15 +854,16 @@ async def test_update_agent_in_db_preserves_secret_nested_under_a_non_sensitive_
|
|||
mock_prisma: Final = MagicMock()
|
||||
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
litellm_params={
|
||||
"provider_config": {
|
||||
"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY,
|
||||
"region": "us-east-1",
|
||||
}
|
||||
},
|
||||
object_permission_id=None,
|
||||
kill_switch=None,
|
||||
return_value=_stored_agent_row(
|
||||
SimpleNamespace(
|
||||
litellm_params={
|
||||
"provider_config": {
|
||||
"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY,
|
||||
"region": "us-east-1",
|
||||
}
|
||||
},
|
||||
object_permission_id=None,
|
||||
)
|
||||
)
|
||||
)
|
||||
updated_agent = MagicMock()
|
||||
|
|
@ -878,10 +909,11 @@ async def test_update_agent_in_db_clears_secret_on_explicit_empty_value():
|
|||
mock_prisma: Final = MagicMock()
|
||||
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
litellm_params={"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY},
|
||||
object_permission_id=None,
|
||||
kill_switch=None,
|
||||
return_value=_stored_agent_row(
|
||||
SimpleNamespace(
|
||||
litellm_params={"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY},
|
||||
object_permission_id=None,
|
||||
)
|
||||
)
|
||||
)
|
||||
updated_agent = MagicMock()
|
||||
|
|
@ -919,12 +951,14 @@ async def test_patch_agent_in_db_preserves_secret_when_litellm_params_omitted():
|
|||
mock_prisma: Final = MagicMock()
|
||||
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value={
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "Old Name",
|
||||
"litellm_params": {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY},
|
||||
"object_permission_id": None,
|
||||
}
|
||||
return_value=_stored_agent_row(
|
||||
{
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "Old Name",
|
||||
"litellm_params": {"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY},
|
||||
"object_permission_id": None,
|
||||
}
|
||||
)
|
||||
)
|
||||
patched_agent = MagicMock()
|
||||
patched_agent.model_dump.return_value = {
|
||||
|
|
@ -958,15 +992,17 @@ async def test_patch_agent_in_db_preserves_secret_when_echoed_back_redacted():
|
|||
mock_prisma: Final = MagicMock()
|
||||
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value={
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "Test Agent",
|
||||
"litellm_params": {
|
||||
"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY,
|
||||
"is_public": False,
|
||||
},
|
||||
"object_permission_id": None,
|
||||
}
|
||||
return_value=_stored_agent_row(
|
||||
{
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "Test Agent",
|
||||
"litellm_params": {
|
||||
"aws_secret_access_key": SENTINEL_AWS_SECRET_ACCESS_KEY,
|
||||
"is_public": False,
|
||||
},
|
||||
"object_permission_id": None,
|
||||
}
|
||||
)
|
||||
)
|
||||
patched_agent = MagicMock()
|
||||
patched_agent.model_dump.return_value = {
|
||||
|
|
@ -997,6 +1033,48 @@ async def test_patch_agent_in_db_preserves_secret_when_echoed_back_redacted():
|
|||
assert stored_params["is_public"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("operation", ["patch", "put"])
|
||||
async def test_runtime_update_drops_legacy_identity_and_keeps_agent_id(operation: str) -> None:
|
||||
registry: Final = AgentRegistry()
|
||||
prisma: Final = MagicMock()
|
||||
identity: Final = {
|
||||
"provider": "microsoft_entra",
|
||||
"tenant_id": "11111111-1111-4111-8111-111111111111",
|
||||
"client_id": "22222222-2222-4222-8222-222222222222",
|
||||
}
|
||||
existing_params: Final = {"identity": identity, "model": "old"}
|
||||
existing: Final = (
|
||||
SimpleNamespace(litellm_params=existing_params, object_permission_id=None)
|
||||
if operation == "put"
|
||||
else {"agent_name": "Readable agent", "litellm_params": existing_params}
|
||||
)
|
||||
prisma.db.litellm_agentstable.find_unique = AsyncMock(return_value=_stored_agent_row(existing))
|
||||
saved: Final = MagicMock()
|
||||
saved.object_permission = None
|
||||
saved.model_dump.return_value = {
|
||||
"agent_id": "unchanged-id",
|
||||
"agent_name": "Renamed agent",
|
||||
"agent_card_params": {},
|
||||
"litellm_params": {"model": "new"},
|
||||
}
|
||||
prisma.db.litellm_agentstable.update = AsyncMock(return_value=saved)
|
||||
update: Final = registry.patch_agent_in_db if operation == "patch" else registry.update_agent_in_db
|
||||
result: Final = await update(
|
||||
agent_id="unchanged-id",
|
||||
agent={"agent_name": "Renamed agent", "agent_card_params": {}, "litellm_params": {"model": "new"}},
|
||||
prisma_client=prisma,
|
||||
updated_by="admin",
|
||||
)
|
||||
stored: Final = prisma.db.litellm_agentstable.update.call_args.kwargs
|
||||
assert stored["where"] == {"agent_id": "unchanged-id"}
|
||||
assert json.loads(stored["data"]["litellm_params"]) == {"model": "new"}, (
|
||||
"a stored litellm_params.identity must not be resurrected once the JWT path no longer honours it"
|
||||
)
|
||||
assert result.agent_id == "unchanged-id"
|
||||
assert "object_permission_id" not in stored["data"]
|
||||
|
||||
|
||||
def _agent_row_mock(access_group_ids: list[str]) -> MagicMock:
|
||||
row: Final = MagicMock()
|
||||
row.model_dump.return_value = {
|
||||
|
|
@ -1063,13 +1141,15 @@ async def test_patch_agent_in_db_replaces_access_group_ids_when_provided(
|
|||
registry: Final = AgentRegistry()
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value={
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "Test Agent",
|
||||
"litellm_params": {},
|
||||
"object_permission_id": None,
|
||||
"access_group_ids": ["ag-1"],
|
||||
}
|
||||
return_value=_stored_agent_row(
|
||||
{
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "Test Agent",
|
||||
"litellm_params": {},
|
||||
"object_permission_id": None,
|
||||
"access_group_ids": ["ag-1"],
|
||||
}
|
||||
)
|
||||
)
|
||||
mock_update = AsyncMock(return_value=_agent_row_mock(expected))
|
||||
mock_prisma.db.litellm_agentstable.update = mock_update
|
||||
|
|
@ -1086,13 +1166,15 @@ async def test_patch_agent_in_db_keeps_access_group_ids_when_omitted():
|
|||
registry: Final = AgentRegistry()
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value={
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "Old Name",
|
||||
"litellm_params": {},
|
||||
"object_permission_id": None,
|
||||
"access_group_ids": ["ag-1"],
|
||||
}
|
||||
return_value=_stored_agent_row(
|
||||
{
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "Old Name",
|
||||
"litellm_params": {},
|
||||
"object_permission_id": None,
|
||||
"access_group_ids": ["ag-1"],
|
||||
}
|
||||
)
|
||||
)
|
||||
mock_update = AsyncMock(return_value=_agent_row_mock(["ag-1"]))
|
||||
mock_prisma.db.litellm_agentstable.update = mock_update
|
||||
|
|
@ -1114,8 +1196,8 @@ async def test_update_agent_in_db_always_writes_access_group_ids(body_access_gro
|
|||
registry: Final = AgentRegistry()
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
litellm_params={}, object_permission_id=None, kill_switch=None, access_group_ids=["ag-1"]
|
||||
return_value=_stored_agent_row(
|
||||
SimpleNamespace(litellm_params={}, object_permission_id=None, access_group_ids=["ag-1"])
|
||||
)
|
||||
)
|
||||
mock_update = AsyncMock(return_value=_agent_row_mock(expected))
|
||||
|
|
@ -1134,6 +1216,34 @@ async def test_update_agent_in_db_always_writes_access_group_ids(body_access_gro
|
|||
assert tuple(mock_update.call_args.kwargs["data"]["access_group_ids"]) == tuple(expected)
|
||||
|
||||
|
||||
def _stored_agent_row(values: Mapping[str, object] | SimpleNamespace) -> LiteLLM_AgentsTable:
|
||||
fields: Final = vars(values) if isinstance(values, SimpleNamespace) else values
|
||||
return LiteLLM_AgentsTable.model_validate(
|
||||
{
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "Test Agent",
|
||||
"agent_card_params": "{}",
|
||||
"extra_headers": [],
|
||||
"agent_access_groups": [],
|
||||
"access_group_ids": [],
|
||||
"created_at": datetime.now(timezone.utc),
|
||||
"updated_at": datetime.now(timezone.utc),
|
||||
"created_by": "admin",
|
||||
"updated_by": "admin",
|
||||
"spend": 0,
|
||||
"identity_managed": False,
|
||||
"enabled": True,
|
||||
"execution_mode": "autonomous",
|
||||
**{
|
||||
key: json.dumps(value)
|
||||
if key in ("litellm_params", "agent_card_params", "kill_switch") and not isinstance(value, str)
|
||||
else value
|
||||
for key, value in fields.items()
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
_KILL_SWITCH: Final = {
|
||||
"url": "https://ops.example.com/kill",
|
||||
"method": "POST",
|
||||
|
|
@ -1194,13 +1304,15 @@ async def test_patch_agent_in_db_keeps_kill_switch_when_omitted_and_clears_it_on
|
|||
registry: Final = AgentRegistry()
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value={
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "Old",
|
||||
"litellm_params": {},
|
||||
"object_permission_id": None,
|
||||
"kill_switch": _KILL_SWITCH,
|
||||
}
|
||||
return_value=_stored_agent_row(
|
||||
{
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "Old",
|
||||
"litellm_params": {},
|
||||
"object_permission_id": None,
|
||||
"kill_switch": _KILL_SWITCH,
|
||||
}
|
||||
)
|
||||
)
|
||||
mock_update = AsyncMock(return_value=_agent_row_mock([]))
|
||||
mock_prisma.db.litellm_agentstable.update = mock_update
|
||||
|
|
@ -1223,13 +1335,15 @@ async def test_patch_agent_in_db_restores_the_stored_kill_switch_secret_behind_t
|
|||
registry: Final = AgentRegistry()
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value={
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "A",
|
||||
"litellm_params": {},
|
||||
"object_permission_id": None,
|
||||
"kill_switch": _KILL_SWITCH,
|
||||
}
|
||||
return_value=_stored_agent_row(
|
||||
{
|
||||
"agent_id": "agent-123",
|
||||
"agent_name": "A",
|
||||
"litellm_params": {},
|
||||
"object_permission_id": None,
|
||||
"kill_switch": _KILL_SWITCH,
|
||||
}
|
||||
)
|
||||
)
|
||||
mock_update = AsyncMock(return_value=_agent_row_mock([]))
|
||||
mock_prisma.db.litellm_agentstable.update = mock_update
|
||||
|
|
@ -1258,7 +1372,9 @@ async def test_update_agent_in_db_clears_kill_switch_when_omitted_and_restores_s
|
|||
registry: Final = AgentRegistry()
|
||||
mock_prisma: Final = MagicMock()
|
||||
mock_prisma.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(litellm_params={}, object_permission_id=None, kill_switch=json.dumps(_KILL_SWITCH))
|
||||
return_value=_stored_agent_row(
|
||||
SimpleNamespace(litellm_params={}, object_permission_id=None, kill_switch=json.dumps(_KILL_SWITCH))
|
||||
)
|
||||
)
|
||||
mock_update = AsyncMock(return_value=_agent_row_mock([]))
|
||||
mock_prisma.db.litellm_agentstable.update = mock_update
|
||||
|
|
@ -1337,3 +1453,92 @@ async def test_agent_listing_preserves_stored_identity_bindings(bound: bool) ->
|
|||
order={"created_at": "desc"},
|
||||
include={"object_permission": True, "identity": True},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("operation", ["create", "patch", "put"])
|
||||
async def test_agent_permissions_are_written_atomically_with_the_registration(operation: str) -> None:
|
||||
from litellm.proxy._types import LiteLLM_ObjectPermissionTable
|
||||
|
||||
registry: Final = AgentRegistry()
|
||||
client: Final = MagicMock()
|
||||
existing: Final = _stored_agent_row({"agent_id": "agent-123", "object_permission_id": "permissions"})
|
||||
client.db.litellm_agentstable.find_unique = AsyncMock(return_value=existing)
|
||||
client.db.litellm_agentstable.create = AsyncMock(return_value=existing)
|
||||
client.db.litellm_agentstable.update = AsyncMock(return_value=existing)
|
||||
client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
|
||||
return_value=(
|
||||
LiteLLM_ObjectPermissionTable(object_permission_id="permissions", models=["prior"], mcp_servers=["slack"])
|
||||
if operation != "create"
|
||||
else None
|
||||
)
|
||||
)
|
||||
incoming: Final = {"agent_name": "Agent", "agent_card_params": {}, "object_permission": {"models": ["new"]}}
|
||||
if operation == "create":
|
||||
await registry.add_agent_to_db(incoming, client, created_by="admin")
|
||||
else:
|
||||
update: Final = registry.patch_agent_in_db if operation == "patch" else registry.update_agent_in_db
|
||||
await update("agent-123", incoming, client, updated_by="admin")
|
||||
write: Final = (
|
||||
client.db.litellm_agentstable.create if operation == "create" else client.db.litellm_agentstable.update
|
||||
)
|
||||
permission: Final = write.call_args.kwargs["data"]["object_permission"][
|
||||
"create" if operation == "create" else "update"
|
||||
]
|
||||
assert permission["models"] == ["new"]
|
||||
if operation != "create":
|
||||
assert permission["mcp_servers"] == ["slack"]
|
||||
assert permission["object_permission_id"] == "permissions"
|
||||
client.db.litellm_objectpermissiontable.update.assert_not_called()
|
||||
client.db.litellm_objectpermissiontable.create.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("operation", ["create", "patch", "put"])
|
||||
async def test_invalid_identity_fails_before_registration_is_written(operation: str) -> None:
|
||||
from fastapi import HTTPException
|
||||
|
||||
registry: Final = AgentRegistry()
|
||||
client: Final = MagicMock()
|
||||
client.db.litellm_agentstable.create = AsyncMock()
|
||||
client.db.litellm_agentstable.update = AsyncMock()
|
||||
client.db.litellm_agentstable.find_unique = AsyncMock(return_value=_stored_agent_row({"agent_id": "agent-123"}))
|
||||
incoming: Final = {"agent_name": "Agent", "agent_card_params": {}, "identity": {"provider": "unknown"}}
|
||||
write: Final = (
|
||||
registry.add_agent_to_db(incoming, client, created_by="admin")
|
||||
if operation == "create"
|
||||
else (registry.patch_agent_in_db if operation == "patch" else registry.update_agent_in_db)(
|
||||
"agent-123", incoming, client, updated_by="admin"
|
||||
)
|
||||
)
|
||||
with pytest.raises(HTTPException) as failure:
|
||||
await write
|
||||
assert failure.value.status_code == 400
|
||||
client.db.litellm_agentstable.create.assert_not_awaited()
|
||||
client.db.litellm_agentstable.update.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("operation", ["create", "patch", "put"])
|
||||
async def test_duplicate_agent_binding_returns_conflict_for_every_write(operation: str) -> None:
|
||||
from fastapi import HTTPException
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
registry: Final = AgentRegistry()
|
||||
client: Final = MagicMock()
|
||||
client.db.litellm_agentstable.find_unique = AsyncMock(return_value=_stored_agent_row({"agent_id": "agent-123"}))
|
||||
failure: Final = UniqueViolationError({"user_facing_error": {"message": "Unique constraint failed", "meta": {"target": ["client_id"]}, "error_code": "P2002"}})
|
||||
client.db.litellm_agentstable.create = AsyncMock(side_effect=failure)
|
||||
client.db.litellm_agentstable.update = AsyncMock(side_effect=failure)
|
||||
incoming: Final = {"agent_name": "Agent", "agent_card_params": {}}
|
||||
write: Final = (
|
||||
registry.add_agent_to_db(incoming, client, created_by="admin")
|
||||
if operation == "create"
|
||||
else (registry.patch_agent_in_db if operation == "patch" else registry.update_agent_in_db)(
|
||||
"agent-123", incoming, client, updated_by="admin"
|
||||
)
|
||||
)
|
||||
with pytest.raises(HTTPException) as denied:
|
||||
await write
|
||||
assert denied.value.status_code == 409
|
||||
assert denied.value.detail == "Agent name or Entra application is already registered"
|
||||
|
|
|
|||
|
|
@ -1,11 +1,13 @@
|
|||
import json
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from litellm.constants import REDACTED_BY_LITELM_STRING
|
||||
|
|
@ -21,7 +23,7 @@ from litellm.proxy.agent_endpoints.endpoints import (
|
|||
router,
|
||||
user_api_key_auth,
|
||||
)
|
||||
from litellm.types.agents import AgentResponse
|
||||
from litellm.types.agents import AgentResponse, PatchAgentRequest
|
||||
from litellm.types.proxy.agent_identity import AgentIdentityBinding
|
||||
|
||||
|
||||
|
|
@ -98,7 +100,7 @@ def test_update_agent_success(mock_prisma_client, mock_user_api_key_auth, monkey
|
|||
"agent_card_params": _sample_agent_card_params(),
|
||||
}
|
||||
mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
|
||||
return_value=existing_agent
|
||||
return_value=AgentResponse.model_validate(existing_agent)
|
||||
)
|
||||
|
||||
mock_registry = MagicMock()
|
||||
|
|
@ -595,6 +597,24 @@ class TestAgentRBACProxyAdmin:
|
|||
)
|
||||
assert resp.status_code == 200
|
||||
|
||||
def test_create_agent_rejects_legacy_litellm_params_identity(self):
|
||||
with patch("litellm.proxy.proxy_server.prisma_client"): # test-quality-ok: proxy_server module global is the endpoint's only injection point
|
||||
self.mock_registry.get_agent_by_name = MagicMock(return_value=None)
|
||||
self.mock_registry.add_agent_to_db = AsyncMock(return_value=_sample_agent_response())
|
||||
config = _sample_agent_config()
|
||||
config["litellm_params"] = {
|
||||
**config["litellm_params"],
|
||||
"identity": {
|
||||
"provider": "microsoft_entra",
|
||||
"tenant_id": "11111111-1111-4111-8111-111111111111",
|
||||
"client_id": "22222222-2222-4222-8222-222222222222",
|
||||
},
|
||||
}
|
||||
resp = self.admin_client.post("/v1/agents", json=config, headers={"Authorization": "Bearer k"})
|
||||
assert resp.status_code == 400, resp.text
|
||||
assert "top-level identity field" in resp.json()["detail"]
|
||||
self.mock_registry.add_agent_to_db.assert_not_awaited()
|
||||
|
||||
def test_create_agent_applies_litellm_merge_to_stored_card(self):
|
||||
"""The card stored in the DB must reflect the LiteLLM-fronting merge."""
|
||||
with patch("litellm.proxy.proxy_server.prisma_client"):
|
||||
|
|
@ -666,11 +686,9 @@ class TestAgentRBACProxyAdmin:
|
|||
"""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(),
|
||||
}
|
||||
return_value=AgentResponse(
|
||||
agent_id="agent-123", agent_name="Existing Agent", agent_card_params=_sample_agent_card_params()
|
||||
)
|
||||
)
|
||||
self.mock_registry.update_agent_in_db = AsyncMock(
|
||||
return_value=AgentResponse(
|
||||
|
|
@ -701,11 +719,9 @@ class TestAgentRBACProxyAdmin:
|
|||
"""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(),
|
||||
}
|
||||
return_value=AgentResponse(
|
||||
agent_id="agent-123", agent_name="Existing Agent", agent_card_params=_sample_agent_card_params()
|
||||
)
|
||||
)
|
||||
self.mock_registry.patch_agent_in_db = AsyncMock(
|
||||
return_value=AgentResponse(
|
||||
|
|
@ -1143,6 +1159,143 @@ def test_make_agent_public_rejects_an_agent_published_only_in_the_db(monkeypatch
|
|||
assert "already in public agent groups" in duplicate.json()["detail"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enabled, claim_field, expected", [(True, "azp", True), (False, "azp", False), (True, None, False)])
|
||||
def test_jwt_authentication_status_does_not_require_virtual_keys(
|
||||
monkeypatch: pytest.MonkeyPatch, enabled: bool, claim_field: str | None, expected: bool
|
||||
) -> None:
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import LiteLLM_JWTAuth
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
|
||||
handler: Final = JWTHandler()
|
||||
handler.update_environment(None, DualCache(), LiteLLM_JWTAuth(agent_id_jwt_field=claim_field))
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": enabled})
|
||||
monkeypatch.setattr(proxy_server, "jwt_handler", handler)
|
||||
agent: Final = _sample_agent_response()
|
||||
response: Final = agent_endpoints._redact_sensitive_agent_fields((agent,), is_admin=True)[0]
|
||||
assert response.jwt_auth_configured is expected
|
||||
assert agent.jwt_auth_configured is False
|
||||
|
||||
|
||||
def test_identity_providers_require_configured_issuer_and_audience(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import LiteLLM_JWTAuth
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
|
||||
handler: Final = JWTHandler()
|
||||
handler.update_environment(None, DualCache(), LiteLLM_JWTAuth())
|
||||
monkeypatch.setattr(proxy_server, "jwt_handler", handler)
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": True})
|
||||
monkeypatch.setenv("JWT_ISSUER", "https://issuer.example")
|
||||
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
|
||||
assert client.get("/v1/agents/identity/providers").json() == []
|
||||
monkeypatch.setenv("JWT_AUDIENCE", "gateway")
|
||||
response: Final = client.get("/v1/agents/identity/providers")
|
||||
assert response.status_code == 200
|
||||
assert response.json() == ["https://issuer.example"]
|
||||
forbidden: Final = _make_app_with_role(LitellmUserRoles.INTERNAL_USER).get("/v1/agents/identity/providers")
|
||||
assert forbidden.status_code == 403
|
||||
|
||||
|
||||
def test_identity_evidence_is_persisted_and_never_taken_from_runtime_metadata(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.types.proxy.agent_identity import AgentIdentityBinding
|
||||
|
||||
binding: Final = AgentIdentityBinding(
|
||||
agent_id="bound",
|
||||
provider="microsoft_entra",
|
||||
tenant_id="11111111-1111-4111-8111-111111111111",
|
||||
client_id="22222222-2222-4222-8222-222222222222",
|
||||
issuer="https://issuer.example",
|
||||
revision="revision-one",
|
||||
)
|
||||
bound: Final = AgentResponse(
|
||||
agent_id="bound",
|
||||
agent_name="Readable name",
|
||||
agent_card_params={},
|
||||
identity=binding,
|
||||
identity_managed=True,
|
||||
litellm_params={"last_authenticated_at": "forged-proof"},
|
||||
)
|
||||
database: Final = MagicMock()
|
||||
database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=bound)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", database)
|
||||
pending: Final = client.get("/v1/agents/bound/identity")
|
||||
assert pending.status_code == 200
|
||||
assert pending.json()["last_authenticated_at"] is None
|
||||
verified_binding: Final = binding.model_copy(
|
||||
update={"last_authenticated_at": datetime(2026, 1, 1, tzinfo=timezone.utc)}
|
||||
)
|
||||
database.writer_db.litellm_agentstable.find_unique.return_value = bound.model_copy(update={"identity": verified_binding})
|
||||
verified: Final = client.get("/v1/agents/bound/identity")
|
||||
assert verified.json()["last_authenticated_at"] == "2026-01-01T00:00:00Z"
|
||||
assert verified.json()["identity"]["client_id"] == binding.client_id
|
||||
database.writer_db.litellm_agentstable.find_unique.return_value = None
|
||||
assert client.get("/v1/agents/missing/identity").status_code == 404
|
||||
database.writer_db.litellm_agentstable.find_unique.side_effect = RuntimeError("unavailable")
|
||||
assert client.get("/v1/agents/bound/identity").status_code == 503
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enabled", [True, False])
|
||||
def test_identity_providers_honor_issuer_specific_audiences_and_global_fallback(
|
||||
monkeypatch: pytest.MonkeyPatch, enabled: bool
|
||||
) -> None:
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import JWTIssuerConfig, LiteLLM_JWTAuth
|
||||
from litellm.proxy.auth.handle_jwt import JWTHandler
|
||||
|
||||
handler: Final = JWTHandler()
|
||||
handler.update_environment(
|
||||
None,
|
||||
DualCache(),
|
||||
LiteLLM_JWTAuth(
|
||||
issuers=[
|
||||
JWTIssuerConfig(issuer="https://scoped.example", audience="gateway"),
|
||||
JWTIssuerConfig(issuer="https://unscoped.example", disable_audience_validation=True),
|
||||
]
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "jwt_handler", handler)
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": enabled})
|
||||
monkeypatch.setenv("JWT_ISSUER", "https://global.example")
|
||||
monkeypatch.setenv("JWT_AUDIENCE", "gateway")
|
||||
assert client.get("/v1/agents/identity/providers").json() == (
|
||||
["https://scoped.example", "https://global.example"] if enabled else []
|
||||
)
|
||||
monkeypatch.setenv("JWT_ISSUER", "https://unscoped.example")
|
||||
assert client.get("/v1/agents/identity/providers").json() == (["https://scoped.example"] if enabled else [])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("change", ({"execution_mode": "delegated"}, {"execution_mode": "both"}))
|
||||
def test_mode_only_edit_requires_the_existing_identity_sso_tenant(
|
||||
monkeypatch: pytest.MonkeyPatch, change: PatchAgentRequest
|
||||
) -> None:
|
||||
from tests.test_litellm.proxy.agent_endpoints.test_managed_identity import BINDING, TENANT, managed_agent
|
||||
|
||||
monkeypatch.setattr(agent_endpoints, "_trusted_agent_issuers", lambda: (BINDING.issuer,))
|
||||
monkeypatch.delenv("MICROSOFT_TENANT", raising=False)
|
||||
monkeypatch.setenv("MICROSOFT_CLIENT_ID", "gateway-client")
|
||||
with pytest.raises(HTTPException, match="Delegated agents require Microsoft SSO"):
|
||||
agent_endpoints._validate_managed_identity_request(change, managed_agent())
|
||||
monkeypatch.setenv("MICROSOFT_TENANT", TENANT)
|
||||
agent_endpoints._validate_managed_identity_request(change, managed_agent())
|
||||
|
||||
|
||||
def test_identity_only_edit_preserves_delegated_mode_validation(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from tests.test_litellm.proxy.agent_endpoints.test_managed_identity import BINDING, managed_agent
|
||||
|
||||
monkeypatch.setattr(agent_endpoints, "_trusted_agent_issuers", lambda: (BINDING.issuer,))
|
||||
monkeypatch.delenv("MICROSOFT_TENANT", raising=False)
|
||||
configuration: Final = BINDING.model_dump(
|
||||
exclude={"agent_id", "issuer", "revision", "last_authenticated_at", "active"}
|
||||
)
|
||||
delegated: Final = managed_agent().model_copy(update={"execution_mode": "delegated"})
|
||||
with pytest.raises(HTTPException, match="Delegated agents require Microsoft SSO"):
|
||||
agent_endpoints._validate_managed_identity_request({"identity": configuration}, delegated)
|
||||
|
||||
_KILL_SWITCH: Final = {
|
||||
"url": "https://ops.example.com/kill",
|
||||
"method": "POST",
|
||||
|
|
@ -1421,3 +1574,21 @@ def test_agent_detail_cache_miss_preserves_admin_identity_visibility(role, monke
|
|||
response = _make_app_with_role(role).get("/v1/agents/agent-123")
|
||||
assert response.status_code == 200
|
||||
assert response.json()["identity"] == (binding.model_dump(mode="json") if role == LitellmUserRoles.PROXY_ADMIN else None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("trusted", [False, True])
|
||||
def test_invalid_identity_and_untrusted_tenant_cannot_be_registered(
|
||||
monkeypatch: pytest.MonkeyPatch, trusted: bool
|
||||
) -> None:
|
||||
from tests.test_litellm.proxy.agent_endpoints.test_managed_identity import BINDING
|
||||
|
||||
configuration: Final = BINDING.model_dump(
|
||||
exclude={"agent_id", "issuer", "revision", "last_authenticated_at", "active"}
|
||||
)
|
||||
monkeypatch.setattr(agent_endpoints, "_trusted_agent_issuers", lambda: (BINDING.issuer,) if trusted else ())
|
||||
request: Final = {"identity": {**configuration, "client_id": "invalid"} if trusted else configuration}
|
||||
message: Final = "Invalid Entra identity configuration" if trusted else "Configure trusted JWT issuer"
|
||||
with pytest.raises(HTTPException, match=message) as failure:
|
||||
agent_endpoints._validate_managed_identity_request(request)
|
||||
assert failure.value.status_code == 400
|
||||
|
||||
|
|
|
|||
151
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
151
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -18071,6 +18071,23 @@ export interface paths {
|
|||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/v1/agents/identity/providers": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
/** Get Agent Identity Providers */
|
||||
get: operations["get_agent_identity_providers_v1_agents_identity_providers_get"];
|
||||
put?: never;
|
||||
post?: never;
|
||||
delete?: never;
|
||||
options?: never;
|
||||
head?: never;
|
||||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/v1/agents/make_public": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
@ -18220,6 +18237,23 @@ export interface paths {
|
|||
patch: operations["patch_agent_v1_agents__agent_id__patch"];
|
||||
trace?: never;
|
||||
};
|
||||
"/v1/agents/{agent_id}/identity": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
/** Get Agent Identity Status */
|
||||
get: operations["get_agent_identity_status_v1_agents__agent_id__identity_get"];
|
||||
put?: never;
|
||||
post?: never;
|
||||
delete?: never;
|
||||
options?: never;
|
||||
head?: never;
|
||||
patch?: never;
|
||||
trace?: never;
|
||||
};
|
||||
"/v1/agents/{agent_id}/kill_switch": {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
@ -24206,11 +24240,19 @@ export interface components {
|
|||
AgentConfig: {
|
||||
/** Access Group Ids */
|
||||
access_group_ids?: string[] | null;
|
||||
agent_card_params: components["schemas"]["AgentCard"];
|
||||
agent_card_params?: components["schemas"]["AgentCard"];
|
||||
/** Agent Name */
|
||||
agent_name: string;
|
||||
/** Enabled */
|
||||
enabled?: boolean;
|
||||
/**
|
||||
* Execution Mode
|
||||
* @enum {string}
|
||||
*/
|
||||
execution_mode?: "autonomous" | "delegated" | "both";
|
||||
/** Extra Headers */
|
||||
extra_headers?: string[] | null;
|
||||
identity?: components["schemas"]["EntraIdentityConfig"] | null;
|
||||
kill_switch?: components["schemas"]["AgentKillSwitchConfig"] | null;
|
||||
/** Litellm Params */
|
||||
litellm_params?: {
|
||||
|
|
@ -30193,6 +30235,32 @@ export interface components {
|
|||
/** Template Id */
|
||||
template_id: string;
|
||||
};
|
||||
/** EntraIdentityConfig */
|
||||
EntraIdentityConfig: {
|
||||
/** Client Id */
|
||||
client_id: string;
|
||||
/**
|
||||
* Provider
|
||||
* @constant
|
||||
*/
|
||||
provider: "microsoft_entra";
|
||||
/**
|
||||
* Required Roles
|
||||
* @default []
|
||||
*/
|
||||
required_roles: string[];
|
||||
/**
|
||||
* Required Scopes
|
||||
* @default [
|
||||
* "user_impersonation"
|
||||
* ]
|
||||
*/
|
||||
required_scopes: string[];
|
||||
/** Service Principal Id */
|
||||
service_principal_id?: string | null;
|
||||
/** Tenant Id */
|
||||
tenant_id: string;
|
||||
};
|
||||
/** EnvironmentReport */
|
||||
EnvironmentReport: {
|
||||
/** Config Lines */
|
||||
|
|
@ -35600,6 +35668,28 @@ export interface components {
|
|||
/** Mcp Server Ids */
|
||||
mcp_server_ids: string[];
|
||||
};
|
||||
/** ManagedAgentIdentityStatus */
|
||||
ManagedAgentIdentityStatus: {
|
||||
/**
|
||||
* Enabled
|
||||
* @default true
|
||||
*/
|
||||
enabled: boolean;
|
||||
/**
|
||||
* Execution Mode
|
||||
* @default autonomous
|
||||
* @enum {string}
|
||||
*/
|
||||
execution_mode: "autonomous" | "delegated" | "both";
|
||||
identity?: components["schemas"]["AgentIdentityBinding"] | null;
|
||||
/**
|
||||
* Identity Managed
|
||||
* @default false
|
||||
*/
|
||||
identity_managed: boolean;
|
||||
/** Last Authenticated At */
|
||||
last_authenticated_at?: string | null;
|
||||
};
|
||||
/**
|
||||
* Mcp
|
||||
* @description Give the model access to additional tools via remote Model Context Protocol
|
||||
|
|
@ -37674,8 +37764,16 @@ export interface components {
|
|||
agent_card_params?: components["schemas"]["AgentCard"];
|
||||
/** Agent Name */
|
||||
agent_name?: string;
|
||||
/** Enabled */
|
||||
enabled?: boolean;
|
||||
/**
|
||||
* Execution Mode
|
||||
* @enum {string}
|
||||
*/
|
||||
execution_mode?: "autonomous" | "delegated" | "both";
|
||||
/** Extra Headers */
|
||||
extra_headers?: string[] | null;
|
||||
identity?: components["schemas"]["EntraIdentityConfig"] | null;
|
||||
kill_switch?: components["schemas"]["AgentKillSwitchConfig"] | null;
|
||||
/** Litellm Params */
|
||||
litellm_params?: {
|
||||
|
|
@ -70133,6 +70231,26 @@ export interface operations {
|
|||
};
|
||||
};
|
||||
};
|
||||
get_agent_identity_providers_v1_agents_identity_providers_get: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path?: never;
|
||||
cookie?: never;
|
||||
};
|
||||
requestBody?: never;
|
||||
responses: {
|
||||
/** @description Successful Response */
|
||||
200: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": string[];
|
||||
};
|
||||
};
|
||||
};
|
||||
};
|
||||
make_agents_public_v1_agents_make_public_post: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
@ -70298,6 +70416,37 @@ export interface operations {
|
|||
};
|
||||
};
|
||||
};
|
||||
get_agent_identity_status_v1_agents__agent_id__identity_get: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
header?: never;
|
||||
path: {
|
||||
agent_id: string;
|
||||
};
|
||||
cookie?: never;
|
||||
};
|
||||
requestBody?: never;
|
||||
responses: {
|
||||
/** @description Successful Response */
|
||||
200: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["ManagedAgentIdentityStatus"];
|
||||
};
|
||||
};
|
||||
/** @description Validation Error */
|
||||
422: {
|
||||
headers: {
|
||||
[name: string]: unknown;
|
||||
};
|
||||
content: {
|
||||
"application/json": components["schemas"]["HTTPValidationError"];
|
||||
};
|
||||
};
|
||||
};
|
||||
};
|
||||
trigger_agent_kill_switch_v1_agents__agent_id__kill_switch_post: {
|
||||
parameters: {
|
||||
query?: never;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue