diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260921190000_agent_identity/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260921190000_agent_identity/migration.sql new file mode 100644 index 00000000000..06cf03b26b5 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260921190000_agent_identity/migration.sql @@ -0,0 +1,97 @@ +-- AlterTable +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "enabled" BOOLEAN NOT NULL DEFAULT true, +ADD COLUMN IF NOT EXISTS "execution_mode" TEXT NOT NULL DEFAULT 'autonomous', +ADD COLUMN IF NOT EXISTS "identity_managed" BOOLEAN NOT NULL DEFAULT false; + +-- AlterTable +ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN IF NOT EXISTS "billing_agent_id" TEXT; + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_AgentIdentity" ( + "agent_id" TEXT NOT NULL, + "active" BOOLEAN NOT NULL DEFAULT true, + "provider" TEXT NOT NULL, + "issuer" TEXT NOT NULL, + "tenant_id" TEXT NOT NULL, + "client_id" TEXT NOT NULL, + "service_principal_id" TEXT, + "required_roles" TEXT[] DEFAULT ARRAY[]::TEXT[], + "required_scopes" TEXT[] DEFAULT ARRAY['user_impersonation']::TEXT[], + "revision" TEXT NOT NULL, + "last_authenticated_at" TIMESTAMP(3), + + CONSTRAINT "LiteLLM_AgentIdentity_pkey" PRIMARY KEY ("agent_id") +); + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_RetiredAgentIdentity" ( + "binding_id" TEXT NOT NULL, + "agent_id" TEXT, + "provider" TEXT NOT NULL, + "issuer" TEXT NOT NULL, + "tenant_id" TEXT NOT NULL, + "client_id" TEXT NOT NULL, + + CONSTRAINT "LiteLLM_RetiredAgentIdentity_pkey" PRIMARY KEY ("binding_id") +); + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_RetiredAgent" ( + "original_agent_id" TEXT NOT NULL, + "retired_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_RetiredAgent_pkey" PRIMARY KEY ("original_agent_id") +); + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_VerifiedSubject" ( + "subject_id" TEXT NOT NULL, + "issuer" TEXT NOT NULL, + "tenant_id" TEXT NOT NULL, + "oid" TEXT NOT NULL, + "kind" TEXT NOT NULL DEFAULT 'human', + "user_id" TEXT, + "verified_via" TEXT NOT NULL DEFAULT 'sso_interactive', + "verified_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_VerifiedSubject_pkey" PRIMARY KEY ("subject_id") +); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_AgentIdentity_provider_tenant_id_client_id_key" ON "LiteLLM_AgentIdentity"("provider", "tenant_id", "client_id"); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_AgentIdentity_issuer_service_principal_id_key" ON "LiteLLM_AgentIdentity"("issuer", "service_principal_id"); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_RetiredAgentIdentity_provider_tenant_id_client_id_key" ON "LiteLLM_RetiredAgentIdentity"("provider", "tenant_id", "client_id"); + +-- CreateIndex +CREATE INDEX IF NOT EXISTS "LiteLLM_VerifiedSubject_user_id_idx" ON "LiteLLM_VerifiedSubject"("user_id"); + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_VerifiedSubject_issuer_tenant_id_oid_key" ON "LiteLLM_VerifiedSubject"("issuer", "tenant_id", "oid"); + +-- AddForeignKey +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_AgentIdentity_agent_id_fkey') THEN + ALTER TABLE "LiteLLM_AgentIdentity" ADD CONSTRAINT "LiteLLM_AgentIdentity_agent_id_fkey" FOREIGN KEY ("agent_id") REFERENCES "LiteLLM_AgentsTable"("agent_id") ON DELETE CASCADE ON UPDATE CASCADE; + END IF; +END $$; + +-- AddForeignKey +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_RetiredAgentIdentity_agent_id_fkey') THEN + ALTER TABLE "LiteLLM_RetiredAgentIdentity" ADD CONSTRAINT "LiteLLM_RetiredAgentIdentity_agent_id_fkey" FOREIGN KEY ("agent_id") REFERENCES "LiteLLM_AgentsTable"("agent_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; + +-- AddForeignKey +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_VerifiedSubject_user_id_fkey') THEN + ALTER TABLE "LiteLLM_VerifiedSubject" ADD CONSTRAINT "LiteLLM_VerifiedSubject_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "LiteLLM_UserTable"("user_id") ON DELETE CASCADE ON UPDATE CASCADE; + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 03e59257f76..f29caa9ceb7 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -78,6 +78,11 @@ model LiteLLM_AgentsTable { object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) + identity_managed Boolean @default(false) + enabled Boolean @default(true) + execution_mode String @default("autonomous") + identity LiteLLM_AgentIdentity? + retired_identities LiteLLM_RetiredAgentIdentity[] tpm_limit Int? rpm_limit Int? session_tpm_limit Int? @@ -88,6 +93,56 @@ model LiteLLM_AgentsTable { updated_by String } +model LiteLLM_AgentIdentity { + agent_id String @id + active Boolean @default(true) + agent LiteLLM_AgentsTable @relation(fields: [agent_id], references: [agent_id], onDelete: Cascade) + provider String + issuer String + tenant_id String + client_id String + service_principal_id String? + required_roles String[] @default([]) + required_scopes String[] @default(["user_impersonation"]) + revision String @default(uuid()) + last_authenticated_at DateTime? + @@unique([provider, tenant_id, client_id]) + @@unique([issuer, service_principal_id]) +} + +model LiteLLM_RetiredAgentIdentity { + binding_id String @id @default(uuid()) + agent_id String? + agent LiteLLM_AgentsTable? @relation(fields: [agent_id], references: [agent_id], onDelete: SetNull) + provider String + issuer String + tenant_id String + client_id String + @@unique([provider, tenant_id, client_id]) +} + +model LiteLLM_RetiredAgent { + original_agent_id String @id + retired_at DateTime @default(now()) +} + +model LiteLLM_VerifiedSubject { + subject_id String @id @default(uuid()) + issuer String + tenant_id String + oid String + kind String @default("human") + user_id String? + user LiteLLM_UserTable? @relation(fields: [user_id], references: [user_id], onDelete: Cascade) + verified_via String @default("sso_interactive") + verified_at DateTime @default(now()) + @@unique([issuer, tenant_id, oid]) + @@index([user_id]) +} + + + + model LiteLLM_OrganizationTable { organization_id String @id @default(uuid()) organization_alias String @@ -241,6 +296,7 @@ model LiteLLM_DeletedTeamTable { // Track spend, rate limit, budget Users model LiteLLM_UserTable { + verified_subjects LiteLLM_VerifiedSubject[] user_id String @id user_alias String? team_id String? @@ -675,6 +731,7 @@ model LiteLLM_SpendLogs { session_id String? status String? mcp_namespaced_tool_name String? + billing_agent_id String? agent_id String? proxy_server_request Json? @default("{}") litellm_call_id String? diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index bb063f8f77a..78b7b729375 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -2521,6 +2521,91 @@ "title": "AgentExtension", "type": "object" }, + "AgentIdentityBinding": { + "properties": { + "active": { + "default": true, + "title": "Active", + "type": "boolean" + }, + "agent_id": { + "title": "Agent Id", + "type": "string" + }, + "client_id": { + "title": "Client Id", + "type": "string" + }, + "issuer": { + "title": "Issuer", + "type": "string" + }, + "last_authenticated_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Last Authenticated At" + }, + "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" + }, + "revision": { + "title": "Revision", + "type": "string" + }, + "service_principal_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Service Principal Id" + }, + "tenant_id": { + "title": "Tenant Id", + "type": "string" + } + }, + "required": [ + "agent_id", + "provider", + "tenant_id", + "client_id", + "issuer", + "revision" + ], + "title": "AgentIdentityBinding", + "type": "object" + }, "AgentInterface": { "description": "Declares a combination of a target URL and a transport protocol.", "properties": { @@ -2972,6 +3057,21 @@ ], "title": "Created By" }, + "enabled": { + "default": true, + "title": "Enabled", + "type": "boolean" + }, + "execution_mode": { + "default": "autonomous", + "enum": [ + "autonomous", + "delegated", + "both" + ], + "title": "Execution Mode", + "type": "string" + }, "extra_headers": { "anyOf": [ { @@ -2986,6 +3086,26 @@ ], "title": "Extra Headers" }, + "identity": { + "anyOf": [ + { + "$ref": "#/components/schemas/AgentIdentityBinding" + }, + { + "type": "null" + } + ] + }, + "identity_managed": { + "default": false, + "title": "Identity Managed", + "type": "boolean" + }, + "jwt_auth_configured": { + "default": false, + "title": "Jwt Auth Configured", + "type": "boolean" + }, "keys": { "anyOf": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d9fb053035b..d421363ee92 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -27,7 +27,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( validate_langfuse_span_scope_value, validate_no_callback_env_reference, ) -from litellm.types.agents import AgentCaller +from litellm.types.agents import AgentCaller, AgentResponse from litellm.types.integrations.compression_interception import ( CompressionSavingsMetadata, ) @@ -46,6 +46,7 @@ from litellm.types.mcp import ( MCPTransportType, ) from litellm.types.mcp_server.mcp_server_manager import MCPInfo +from litellm.types.proxy.agent_identity import ManagedAgentContext from litellm.types.proxy.carried_budget_state import ( OrgBudgetSnapshot, TeamBudgetSnapshot, @@ -567,6 +568,7 @@ class LiteLLMRoutes(enum.Enum): "/agents", "/a2a/{agent_id}", "/a2a/{agent_id}/message/send", + "/v1/a2a/{agent_id}/message/send", "/a2a/{agent_id}/message/stream", "/a2a/{agent_id}/.well-known/agent-card.json", ) @@ -3302,6 +3304,8 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob # metadata or JWT claims, so it cannot be forged to gain the team-inherited MCP grant union # or to escape the caller-Authorization egress scrub. exclude=True keeps it out of serialization. mcp_admitted_user_subject: bool = Field(default=False, exclude=True) + requires_fresh_policy: bool = Field(default=False, exclude=True) + mcp_explicit_grants_only: bool = Field(default=False, exclude=True) # team_id -> that team's mcp_rpm_limit map, for a keyless admitted subject that reaches MCP # servers through several teams at once and therefore has no single team_id for the limiter to # key off. Server-only and stripped from validated input for the same reason as the marker @@ -3326,6 +3330,12 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob "user id." ), ) + invoked_agent_id: str | None = Field(default=None, exclude=True) + invoked_agent_policy: AgentResponse | None = Field(default=None, exclude=True) + agent_invocation_cost: float | None = Field(default=None, exclude=True) + billing_agent_policy: AgentResponse | None = Field(default=None, exclude=True) + managed_agent_policy: AgentResponse | None = Field(default=None, exclude=True) + managed_agent_context: ManagedAgentContext | None = Field(default=None, exclude=True) agent_caller: AgentCaller | None = Field( default=None, exclude=True, @@ -3363,11 +3373,19 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob # path via post-construction assignment. Strip it from any validated input (constructor # kwargs, model_validate, a JWT/key claim splat) so it can never be forged from caller data. values.pop("mcp_admitted_user_subject", None) + values.pop("requires_fresh_policy", None) + values.pop("mcp_explicit_grants_only", None) values.pop("mcp_source_team_rpm_limits", None) values.pop("mcp_session_resource_server_id", None) values.pop("mcp_toolset_id", None) values.pop("via_virtual_key", None) values.pop("agent_caller", None) + values.pop("managed_agent_context", None) + values.pop("managed_agent_policy", None) + values.pop("invoked_agent_id", None) + values.pop("invoked_agent_policy", None) + values.pop("agent_invocation_cost", None) + values.pop("billing_agent_policy", None) if values.get("api_key") is not None: values.update({"token": cls._safe_hash_litellm_api_key(values.get("api_key"))}) if isinstance(values.get("api_key"), str): @@ -4063,6 +4081,11 @@ class SpendLogsRouterMetadata(TypedDict): class SpendLogsMetadata(TypedDict): + actor_agent_id: ReadOnly[NotRequired[str | None]] + target_agent_id: ReadOnly[NotRequired[str | None]] + billing_agent_id: ReadOnly[NotRequired[str | None]] + agent_execution_mode: ReadOnly[NotRequired[str | None]] + verified_human_user_id: ReadOnly[NotRequired[str | None]] autorouter_baseline_observation: ReadOnly[str | None] """ Specific metadata k,v pairs logged to spendlogs for easier cost tracking @@ -4126,6 +4149,7 @@ class SpendLogsPayload(TypedDict): model_id: str | None model_group: str | None mcp_namespaced_tool_name: str | None + billing_agent_id: ReadOnly[NotRequired[str | None]] agent_id: str | None api_base: str user: str @@ -5048,6 +5072,7 @@ class JWTAuthBuilderResult(TypedDict): org_id: str | None team_membership: LiteLLM_TeamMembership | None jwt_claims: dict # Decoded JWT token claims (avoids re-decoding) + managed_agent_context: ReadOnly[NotRequired[ManagedAgentContext | None]] agent_id: ReadOnly[str | None] diff --git a/litellm/proxy/agent_endpoints/identity.py b/litellm/proxy/agent_endpoints/identity.py new file mode 100644 index 00000000000..c0e5a748144 --- /dev/null +++ b/litellm/proxy/agent_endpoints/identity.py @@ -0,0 +1,17 @@ +from collections.abc import Mapping +from typing import Final + +from fastapi import HTTPException + +LEGACY_IDENTITY_MESSAGE: Final = ( + "litellm_params.identity is not supported: bind an Entra application through the top-level identity field" +) + + +def has_legacy_identity(params: Mapping[str, object] | None) -> bool: + return params is not None and "identity" in params + + +def reject_legacy_identity(params: Mapping[str, object] | None) -> None: + if has_legacy_identity(params): + raise HTTPException(400, LEGACY_IDENTITY_MESSAGE) diff --git a/litellm/proxy/agent_endpoints/identity_store.py b/litellm/proxy/agent_endpoints/identity_store.py new file mode 100644 index 00000000000..0d9d21108e5 --- /dev/null +++ b/litellm/proxy/agent_endpoints/identity_store.py @@ -0,0 +1,250 @@ +import json +from collections.abc import Mapping +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Final + +from litellm.proxy.agent_endpoints.managed_identity import classify_agent_subject +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, get_management_object_ttl +from litellm.repositories.table_repositories import ( + AgentIdentityRepository, + AgentsRepository, + RetiredAgentIdentityRepository, + RetiredAgentRepository, + VerifiedSubjectRepository, +) +from litellm.types.agents import AgentResponse +from litellm.types.proxy.agent_identity import ( + AgentIdentityFailure, + ManagedAgentContext, + MicrosoftInteractiveSubject, + VerifiedHumanSubject, +) + +if TYPE_CHECKING: + from prisma.models import LiteLLM_VerifiedSubject + from prisma.types import ( + LiteLLM_AgentIdentityUpdateManyMutationInput, + LiteLLM_AgentIdentityWhereInput, + LiteLLM_AgentIdentityWhereUniqueInput, + LiteLLM_AgentsTableInclude, + LiteLLM_AgentsTableWhereUniqueInput, + LiteLLM_VerifiedSubjectCreateInput, + LiteLLM_VerifiedSubjectUpsertInput, + LiteLLM_VerifiedSubjectWhereUniqueInput, + ) + + +class AgentIdentityStore: + @classmethod + def from_client(cls, client: object, *, cache: UserApiKeyCache | None = None) -> "AgentIdentityStore": + return cls( + AgentsRepository(client, use_writer=True), + AgentIdentityRepository(client, use_writer=True), + VerifiedSubjectRepository(client, use_writer=True), + RetiredAgentIdentityRepository(client, use_writer=True), + RetiredAgentRepository(client, use_writer=True), + cache=cache, + ) + + def __init__( + self, + agents: AgentsRepository, + identities: AgentIdentityRepository, + humans: VerifiedSubjectRepository, + retired: RetiredAgentIdentityRepository | None = None, + retired_agents: RetiredAgentRepository | None = None, + *, + cache: UserApiKeyCache | None = None, + ) -> None: + self.agents = agents + self.identities = identities + self.humans = humans + self.retired = retired + self.retired_agents = retired_agents + self.cache = cache + + async def agent(self, agent_id: str) -> AgentResponse | AgentIdentityFailure | None: + try: + where: Final[LiteLLM_AgentsTableWhereUniqueInput] = {"agent_id": agent_id} + include: Final[LiteLLM_AgentsTableInclude] = { + "identity": True, + "object_permission": True, + } + row: Final = await self.agents.table.find_unique(where=where, include=include) + if row is None: + return None + return AgentResponse.model_validate(row.model_dump()) + except Exception: + return AgentIdentityFailure(code="policy_unavailable", message="Agent policy could not be loaded") + + async def unbound_client(self, where: "LiteLLM_AgentIdentityWhereUniqueInput") -> AgentIdentityFailure | None: + if self.retired is not None: + try: + retired: Final = await self.retired.table.find_unique(where=where) + except Exception: + return AgentIdentityFailure( + code="policy_unavailable", message="Retired agent identity could not be checked" + ) + if retired is not None: + return AgentIdentityFailure(message="This agent identity binding has been retired") + return None + + async def _bound_agent_id(self, tenant_id: str, client_id: str) -> str | AgentIdentityFailure | None: + cache_key: Final = f"agent_identity:{json.dumps((tenant_id, client_id))}" + cached: Final[object] = await self.cache.async_get_cache(key=cache_key) if self.cache is not None else None + if isinstance(cached, str): + return cached + where: Final[LiteLLM_AgentIdentityWhereUniqueInput] = { + "provider_tenant_id_client_id": { + "provider": "microsoft_entra", + "tenant_id": tenant_id, + "client_id": client_id, + } + } + try: + row: Final = await self.identities.table.find_unique(where=where) + except Exception: + return AgentIdentityFailure(code="policy_unavailable", message="Agent identity could not be loaded") + if row is None: + return await self.unbound_client(where) + if self.cache is not None: + await self.cache.async_set_cache( + key=cache_key, value=row.agent_id, ttl=get_management_object_ttl(self.cache) + ) + return row.agent_id + + async def resolve_verified_claims( + self, claims: Mapping[str, object] + ) -> ManagedAgentContext | AgentIdentityFailure | None: + issuer: Final = claims.get("iss") + tenant: Final = claims.get("tid") + client: Final = claims.get("azp") + if not isinstance(issuer, str) or not isinstance(tenant, str) or not isinstance(client, str): + return None + agent_id: Final = await self._bound_agent_id(tenant, client) + if agent_id is None or isinstance(agent_id, AgentIdentityFailure): + return agent_id + agent: Final = await self.agent(agent_id) + if isinstance(agent, AgentIdentityFailure): + return agent + if ( + agent is None + or not agent.identity_managed + or not agent.enabled + or agent.identity is None + or not agent.identity.active + ): + return AgentIdentityFailure(message="Agent is disabled or no longer bound to an identity") + subject: Final = classify_agent_subject(agent.identity, claims, agent.execution_mode) + if isinstance(subject, AgentIdentityFailure): + return subject + if subject.kind == "application": + return ManagedAgentContext( + agent_id=agent.agent_id, + binding_revision=agent.identity.revision, + mode=subject.mode, + subject_oid=subject.oid, + ) + proven: Final = await self.subject(issuer, tenant, claims.get("oid")) + if isinstance(proven, AgentIdentityFailure): + return proven + human: Final = ( + VerifiedHumanSubject.model_validate(proven.model_dump()) + if proven is not None + and proven.kind == "human" + and proven.verified_via == "sso_interactive" + and proven.user_id is not None + else None + ) + if human is None: + return AgentIdentityFailure(message="The delegated user must first sign in through trusted Microsoft SSO") + return ManagedAgentContext( + agent_id=agent.agent_id, + binding_revision=agent.identity.revision, + mode=subject.mode, + user_id=human.user_id, + subject_oid=subject.oid, + ) + + async def subject( + self, issuer: str, tenant_id: str, oid: object + ) -> "LiteLLM_VerifiedSubject | AgentIdentityFailure | None": + if not isinstance(oid, str): + return None + try: + where: Final[LiteLLM_VerifiedSubjectWhereUniqueInput] = { + "issuer_tenant_id_oid": {"issuer": issuer, "tenant_id": tenant_id, "oid": oid} + } + return await self.humans.table.find_unique(where=where) + except Exception: + return AgentIdentityFailure(code="policy_unavailable", message="Subject classification is unavailable") + + async def retired_agent(self, agent_id: str) -> bool | AgentIdentityFailure: + if self.retired_agents is None: + return AgentIdentityFailure(code="policy_unavailable", message="Agent history is unavailable") + try: + return await self.retired_agents.table.find_unique(where={"original_agent_id": agent_id}) is not None + except Exception: + return AgentIdentityFailure(code="policy_unavailable", message="Agent history is unavailable") + + async def record_authentication(self, context: ManagedAgentContext) -> AgentIdentityFailure | None: + try: + if context.binding_revision is None: + return AgentIdentityFailure(message="Agent authentication requires a binding revision") + where: Final[LiteLLM_AgentIdentityWhereInput] = { + "agent_id": context.agent_id, + "revision": context.binding_revision, + "active": True, + "agent": {"is": {"enabled": True, "identity_managed": True}}, + } + data: Final[LiteLLM_AgentIdentityUpdateManyMutationInput] = { + "last_authenticated_at": datetime.now(timezone.utc) + } + count: Final = await self.identities.table.update_many(where=where, data=data) + if count != 1: + return AgentIdentityFailure(message="Agent identity changed during authentication; retry") + return None + except Exception: + return AgentIdentityFailure(code="policy_unavailable", message="Agent authentication could not be recorded") + + async def enroll_interactive_human( + self, + subject: MicrosoftInteractiveSubject, + user_id: str, + ) -> AgentIdentityFailure | None: + try: + where: Final[LiteLLM_VerifiedSubjectWhereUniqueInput] = { + "issuer_tenant_id_oid": {"issuer": subject.issuer, "tenant_id": subject.tenant_id, "oid": subject.oid} + } + create_data: Final[LiteLLM_VerifiedSubjectCreateInput] = { + "issuer": subject.issuer, + "tenant_id": subject.tenant_id, + "oid": subject.oid, + "user_id": user_id, + "verified_via": "sso_interactive", + } + data: Final[LiteLLM_VerifiedSubjectUpsertInput] = {"create": create_data, "update": {}} + row: Final = await self.humans.table.upsert(where=where, data=data) + if row.kind != "human" or row.user_id != user_id or row.verified_via != "sso_interactive": + return AgentIdentityFailure(message="Microsoft subject is already bound to another local identity") + return None + except Exception: + return AgentIdentityFailure( + code="policy_unavailable", message="Microsoft subject enrollment is unavailable" + ) + + +async def resolve_managed_agent( + claims: Mapping[str, object], + client: object, + *, + cache: UserApiKeyCache | None = None, +) -> ManagedAgentContext | None: + from litellm.proxy.agent_endpoints.managed_identity import raise_identity_failure + + if client is None: + return None + result: Final = await AgentIdentityStore.from_client(client, cache=cache).resolve_verified_claims(claims) + if isinstance(result, AgentIdentityFailure): + raise_identity_failure(result) + return result diff --git a/litellm/proxy/agent_endpoints/managed_identity.py b/litellm/proxy/agent_endpoints/managed_identity.py new file mode 100644 index 00000000000..260b74fcbd1 --- /dev/null +++ b/litellm/proxy/agent_endpoints/managed_identity.py @@ -0,0 +1,220 @@ +from collections.abc import Mapping +from datetime import datetime +from typing import Final, NoReturn, TypedDict +from uuid import uuid4 + +from fastapi import HTTPException +from pydantic import TypeAdapter, ValidationError +from typing_extensions import ReadOnly + +from litellm.types.agents import AgentResponse +from litellm.types.proxy.agent_identity import ( + AgentExecutionMode, + AgentIdentityBinding, + AgentIdentityFailure, + AgentSubject, + EntraIdentityConfig, +) + +_MODE: Final = TypeAdapter(AgentExecutionMode) + + +class IdentityFields(TypedDict, total=False): + provider: ReadOnly[str] + tenant_id: ReadOnly[str] + client_id: ReadOnly[str] + issuer: ReadOnly[str] + service_principal_id: ReadOnly[str | None] + required_roles: ReadOnly[tuple[str, ...]] + required_scopes: ReadOnly[tuple[str, ...]] + active: ReadOnly[bool] + revision: ReadOnly[str] + last_authenticated_at: ReadOnly[datetime | None] + + +class IdentityUpsert(TypedDict): + create: ReadOnly[IdentityFields] + update: ReadOnly[IdentityFields] + + +class IdentityRelationWrite(TypedDict, total=False): + create: ReadOnly[IdentityFields] + update: ReadOnly[IdentityFields] + upsert: ReadOnly[IdentityUpsert] + + +class IdentityHistoryKey(TypedDict): + provider: ReadOnly[str] + tenant_id: ReadOnly[str] + client_id: ReadOnly[str] + + +class IdentityHistoryWhere(TypedDict): + provider_tenant_id_client_id: ReadOnly[IdentityHistoryKey] + + +class IdentityHistoryEntry(IdentityHistoryKey): + issuer: ReadOnly[str] + + +class IdentityHistoryConnect(TypedDict): + where: ReadOnly[IdentityHistoryWhere] + create: ReadOnly[IdentityHistoryEntry] + + +class IdentityHistoryWrite(TypedDict): + connectOrCreate: ReadOnly[IdentityHistoryConnect] + + +class ManagedWriteFields(TypedDict, total=False): + enabled: ReadOnly[bool] + execution_mode: ReadOnly[AgentExecutionMode] + identity_managed: ReadOnly[bool] + identity: ReadOnly[IdentityRelationWrite] + retired_identities: ReadOnly[IdentityHistoryWrite] + + +def raise_identity_failure(failure: AgentIdentityFailure, status_code: int = 403) -> NoReturn: + raise HTTPException(503 if failure.code == "policy_unavailable" else status_code, failure.message) + + +def _configuration_failure( + identity: EntraIdentityConfig | AgentIdentityBinding | None, + mode: AgentExecutionMode, + enabling_without_binding: bool, +) -> AgentIdentityFailure | None: + if identity is not None and mode != "delegated" and not identity.service_principal_id: + return AgentIdentityFailure( + message="Autonomous mode requires the Enterprise application service-principal object ID" + ) + if enabling_without_binding and ( + identity is None or isinstance(identity, AgentIdentityBinding) and not identity.active + ): + return AgentIdentityFailure(message="Bind an identity before enabling this managed agent") + return None + + +def managed_write_fields( + incoming: Mapping[str, object], + existing: AgentResponse | None, + updated_by: str, +) -> ManagedWriteFields | AgentIdentityFailure: + try: + identity: Final = ( + EntraIdentityConfig.model_validate(incoming["identity"]) if incoming.get("identity") is not None else None + ) + mode: Final = _MODE.validate_python( + incoming.get("execution_mode", existing.execution_mode if existing else "autonomous") + ) + current_identity: Final = identity if "identity" in incoming else existing.identity if existing else None + failure: Final = _configuration_failure( + current_identity, + mode, + incoming.get("enabled") is True + and "identity" not in incoming + and bool(existing and existing.identity_managed), + ) + if failure is not None: + return failure + empty: Final[ManagedWriteFields] = {} + identity_fields: Final = _identity_write(identity, existing) if "identity" in incoming else empty + result: Final[ManagedWriteFields] = { + **({"enabled": incoming["enabled"] is True} if "enabled" in incoming else {}), + **({"execution_mode": mode} if "execution_mode" in incoming else {}), + **identity_fields, + } + return result + except (ValidationError, ValueError) as exc: + return AgentIdentityFailure(message=f"Invalid agent identity configuration: {exc}") + + +def _identity_write(identity: EntraIdentityConfig | None, existing: AgentResponse | None) -> ManagedWriteFields: + if identity is None: + unbind: Final[ManagedWriteFields] = { + **( + {"identity": {"update": {"active": False, "revision": str(uuid4()), "last_authenticated_at": None}}} + if existing and existing.identity + else {} + ), + **({"identity_managed": True, "enabled": False} if existing and existing.identity_managed else {}), + } + return unbind + if ( + existing + and existing.identity + and existing.identity.active + and all(getattr(existing.identity, name) == value for name, value in identity.model_dump().items()) + ): + unchanged: Final[ManagedWriteFields] = {} + return unchanged + binding: Final[IdentityFields] = { + "provider": identity.provider, + "tenant_id": identity.tenant_id, + "client_id": identity.client_id, + "service_principal_id": identity.service_principal_id, + "required_roles": identity.required_roles, + "required_scopes": identity.required_scopes, + "issuer": identity.issuer, + "active": True, + "revision": str(uuid4()), + "last_authenticated_at": None, + } + result: Final[ManagedWriteFields] = { + "retired_identities": { + "connectOrCreate": { + "where": { + "provider_tenant_id_client_id": { + "provider": identity.provider, + "tenant_id": identity.tenant_id, + "client_id": identity.client_id, + } + }, + "create": { + "provider": identity.provider, + "issuer": identity.issuer, + "tenant_id": identity.tenant_id, + "client_id": identity.client_id, + }, + } + }, + "identity_managed": True, + "identity": {"upsert": {"create": binding, "update": binding}} if existing else {"create": binding}, + } + return result + + +def classify_agent_subject( + binding: AgentIdentityBinding, + claims: Mapping[str, object], + allowed_mode: AgentExecutionMode, +) -> AgentSubject | AgentIdentityFailure: + if (claims.get("iss"), claims.get("tid"), claims.get("azp")) != ( + binding.issuer, + binding.tenant_id, + binding.client_id, + ): + return AgentIdentityFailure(message="Token does not match the registered Entra application") + oid: Final = claims.get("oid") + if not isinstance(oid, str) or not oid: + return AgentIdentityFailure(message="Entra token must identify its object subject") + scope: Final = claims.get("scp") + facets: Final = claims.get("xms_sub_fct") + if facets is not None and (not isinstance(facets, str) or "13" in facets.split()): + return AgentIdentityFailure(message="Native agent-user authentication is not supported by this binding") + if scope is not None and not isinstance(scope, str): + return AgentIdentityFailure(message="Invalid delegated scope claim") + if isinstance(scope, str) and scope: + if allowed_mode == "autonomous" or oid == binding.service_principal_id or claims.get("idtyp") == "app": + return AgentIdentityFailure(message="Delegated token contradicts the configured agent identity or mode") + granted_scopes: Final = frozenset(scope.split()) + if not granted_scopes or not frozenset(binding.required_scopes).issubset(granted_scopes): + return AgentIdentityFailure(message="Token lacks the required delegated scopes") + return AgentSubject(kind="delegated_subject", oid=oid, mode="delegated") + if allowed_mode == "delegated" or oid != binding.service_principal_id or claims.get("idtyp") == "user": + return AgentIdentityFailure(message="Application token contradicts the configured agent identity or mode") + roles: Final = claims.get("roles", ()) + if not isinstance(roles, (list, tuple)) or any(not isinstance(role, str) for role in roles): + return AgentIdentityFailure(message="Invalid application roles claim") + if not frozenset(binding.required_roles).issubset(roles): + return AgentIdentityFailure(message="Token lacks the required application roles") + return AgentSubject(kind="application", oid=oid, mode="autonomous") diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 03e59257f76..f29caa9ceb7 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -78,6 +78,11 @@ model LiteLLM_AgentsTable { object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) + identity_managed Boolean @default(false) + enabled Boolean @default(true) + execution_mode String @default("autonomous") + identity LiteLLM_AgentIdentity? + retired_identities LiteLLM_RetiredAgentIdentity[] tpm_limit Int? rpm_limit Int? session_tpm_limit Int? @@ -88,6 +93,56 @@ model LiteLLM_AgentsTable { updated_by String } +model LiteLLM_AgentIdentity { + agent_id String @id + active Boolean @default(true) + agent LiteLLM_AgentsTable @relation(fields: [agent_id], references: [agent_id], onDelete: Cascade) + provider String + issuer String + tenant_id String + client_id String + service_principal_id String? + required_roles String[] @default([]) + required_scopes String[] @default(["user_impersonation"]) + revision String @default(uuid()) + last_authenticated_at DateTime? + @@unique([provider, tenant_id, client_id]) + @@unique([issuer, service_principal_id]) +} + +model LiteLLM_RetiredAgentIdentity { + binding_id String @id @default(uuid()) + agent_id String? + agent LiteLLM_AgentsTable? @relation(fields: [agent_id], references: [agent_id], onDelete: SetNull) + provider String + issuer String + tenant_id String + client_id String + @@unique([provider, tenant_id, client_id]) +} + +model LiteLLM_RetiredAgent { + original_agent_id String @id + retired_at DateTime @default(now()) +} + +model LiteLLM_VerifiedSubject { + subject_id String @id @default(uuid()) + issuer String + tenant_id String + oid String + kind String @default("human") + user_id String? + user LiteLLM_UserTable? @relation(fields: [user_id], references: [user_id], onDelete: Cascade) + verified_via String @default("sso_interactive") + verified_at DateTime @default(now()) + @@unique([issuer, tenant_id, oid]) + @@index([user_id]) +} + + + + model LiteLLM_OrganizationTable { organization_id String @id @default(uuid()) organization_alias String @@ -241,6 +296,7 @@ model LiteLLM_DeletedTeamTable { // Track spend, rate limit, budget Users model LiteLLM_UserTable { + verified_subjects LiteLLM_VerifiedSubject[] user_id String @id user_alias String? team_id String? @@ -675,6 +731,7 @@ model LiteLLM_SpendLogs { session_id String? status String? mcp_namespaced_tool_name String? + billing_agent_id String? agent_id String? proxy_server_request Json? @default("{}") litellm_call_id String? diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index ab68f1a2bc7..4e511a2ec93 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -21,8 +21,9 @@ class PrismaTableRepository(Generic[RowT_co]): table_name: str - def __init__(self, prisma_client: object): + def __init__(self, prisma_client: object, *, use_writer: bool = False) -> None: self._prisma_client = prisma_client + self._use_writer = use_writer @property def prisma_client(self) -> Any: @@ -32,7 +33,9 @@ class PrismaTableRepository(Generic[RowT_co]): @property def table(self) -> TableActions[RowT_co]: - actions: Final[TableActions[RowT_co]] = getattr(self.prisma_client.db, self.table_name) + actions: Final[TableActions[RowT_co]] = getattr( + self.prisma_client.writer_db if self._use_writer else self.prisma_client.db, self.table_name + ) return wrap_table_actions_for_config_sync(actions=actions, table_name=self.table_name) @@ -44,6 +47,18 @@ class AgentsRepository(PrismaTableRepository["prisma_models.LiteLLM_AgentsTable" table_name = "litellm_agentstable" +class AgentIdentityRepository(PrismaTableRepository["prisma_models.LiteLLM_AgentIdentity"]): + table_name = "litellm_agentidentity" + + +class RetiredAgentIdentityRepository(PrismaTableRepository["prisma_models.LiteLLM_RetiredAgentIdentity"]): + table_name = "litellm_retiredagentidentity" + + +class VerifiedSubjectRepository(PrismaTableRepository["prisma_models.LiteLLM_VerifiedSubject"]): + table_name = "litellm_verifiedsubject" + + class ObjectPermissionRepository(PrismaTableRepository["prisma_models.LiteLLM_ObjectPermissionTable"]): table_name = "litellm_objectpermissiontable" @@ -250,3 +265,7 @@ class AuditLogRepository(PrismaTableRepository["prisma_models.LiteLLM_AuditLog"] class AdaptiveRouterSessionRepository(PrismaTableRepository["prisma_models.LiteLLM_AdaptiveRouterSession"]): table_name = "litellm_adaptiveroutersession" + + +class RetiredAgentRepository(PrismaTableRepository["prisma_models.LiteLLM_RetiredAgent"]): + table_name = "litellm_retiredagent" diff --git a/litellm/types/agents.py b/litellm/types/agents.py index f7aef09fa29..3b460bd66c6 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -7,6 +7,10 @@ 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, +) if TYPE_CHECKING: from a2a.types import SendMessageResponse @@ -301,6 +305,11 @@ class AgentKeySummary(BaseModel): class AgentResponse(BaseModel): + identity: AgentIdentityBinding | None = None + identity_managed: bool = False + enabled: bool = True + execution_mode: AgentExecutionMode = "autonomous" + jwt_auth_configured: bool = False agent_id: str agent_name: str litellm_params: dict[str, object] | None = None diff --git a/litellm/types/proxy/agent_identity.py b/litellm/types/proxy/agent_identity.py new file mode 100644 index 00000000000..a7fe0be37e1 --- /dev/null +++ b/litellm/types/proxy/agent_identity.py @@ -0,0 +1,96 @@ +from datetime import datetime +from typing import Literal, TypeAlias +from uuid import UUID + +from pydantic import BaseModel, ConfigDict, Field, field_validator + +AgentExecutionMode: TypeAlias = Literal["autonomous", "delegated", "both"] + + +class EntraIdentityConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + provider: Literal["microsoft_entra"] + tenant_id: str + client_id: str + service_principal_id: str | None = None + required_roles: tuple[str, ...] = () + required_scopes: tuple[str, ...] = Field( + default=("user_impersonation",), + description="Required delegated scopes. An empty list accepts any nonempty scope granted for this gateway.", + ) + + @field_validator("tenant_id", "client_id", "service_principal_id") + @classmethod + def normalize_identifier(cls, value: str | None) -> str | None: + return str(UUID(value)) if value is not None else None + + @property + def issuer(self) -> str: + return f"https://login.microsoftonline.com/{self.tenant_id}/v2.0" + + +class AgentIdentityBinding(BaseModel): + model_config = ConfigDict(frozen=True) + + agent_id: str + active: bool = True + provider: Literal["microsoft_entra"] + tenant_id: str + client_id: str + service_principal_id: str | None = None + issuer: str + required_roles: tuple[str, ...] = () + required_scopes: tuple[str, ...] = ("user_impersonation",) + revision: str + last_authenticated_at: datetime | None = None + + +class AgentSubject(BaseModel): + model_config = ConfigDict(frozen=True) + + kind: Literal["application", "delegated_subject"] + oid: str + mode: Literal["autonomous", "delegated"] + + +class AgentIdentityFailure(BaseModel): + model_config = ConfigDict(frozen=True) + + code: Literal["identity_denied", "policy_unavailable"] = "identity_denied" + message: str + + +class ManagedAgentContext(BaseModel): + model_config = ConfigDict(frozen=True) + + agent_id: str + binding_revision: str | None = None + mode: Literal["autonomous", "delegated"] + user_id: str | None = None + subject_oid: str | None = None + + +class VerifiedHumanSubject(BaseModel): + model_config = ConfigDict(frozen=True) + + issuer: str + tenant_id: str + oid: str + user_id: str + + +class MicrosoftInteractiveSubject(BaseModel): + model_config = ConfigDict(frozen=True) + + issuer: str + tenant_id: str + oid: str + + +class ManagedAgentIdentityStatus(BaseModel): + identity: AgentIdentityBinding | None = None + identity_managed: bool = False + enabled: bool = True + execution_mode: AgentExecutionMode = "autonomous" + last_authenticated_at: datetime | None = None diff --git a/schema.prisma b/schema.prisma index 03e59257f76..f29caa9ceb7 100644 --- a/schema.prisma +++ b/schema.prisma @@ -78,6 +78,11 @@ model LiteLLM_AgentsTable { object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) + identity_managed Boolean @default(false) + enabled Boolean @default(true) + execution_mode String @default("autonomous") + identity LiteLLM_AgentIdentity? + retired_identities LiteLLM_RetiredAgentIdentity[] tpm_limit Int? rpm_limit Int? session_tpm_limit Int? @@ -88,6 +93,56 @@ model LiteLLM_AgentsTable { updated_by String } +model LiteLLM_AgentIdentity { + agent_id String @id + active Boolean @default(true) + agent LiteLLM_AgentsTable @relation(fields: [agent_id], references: [agent_id], onDelete: Cascade) + provider String + issuer String + tenant_id String + client_id String + service_principal_id String? + required_roles String[] @default([]) + required_scopes String[] @default(["user_impersonation"]) + revision String @default(uuid()) + last_authenticated_at DateTime? + @@unique([provider, tenant_id, client_id]) + @@unique([issuer, service_principal_id]) +} + +model LiteLLM_RetiredAgentIdentity { + binding_id String @id @default(uuid()) + agent_id String? + agent LiteLLM_AgentsTable? @relation(fields: [agent_id], references: [agent_id], onDelete: SetNull) + provider String + issuer String + tenant_id String + client_id String + @@unique([provider, tenant_id, client_id]) +} + +model LiteLLM_RetiredAgent { + original_agent_id String @id + retired_at DateTime @default(now()) +} + +model LiteLLM_VerifiedSubject { + subject_id String @id @default(uuid()) + issuer String + tenant_id String + oid String + kind String @default("human") + user_id String? + user LiteLLM_UserTable? @relation(fields: [user_id], references: [user_id], onDelete: Cascade) + verified_via String @default("sso_interactive") + verified_at DateTime @default(now()) + @@unique([issuer, tenant_id, oid]) + @@index([user_id]) +} + + + + model LiteLLM_OrganizationTable { organization_id String @id @default(uuid()) organization_alias String @@ -241,6 +296,7 @@ model LiteLLM_DeletedTeamTable { // Track spend, rate limit, budget Users model LiteLLM_UserTable { + verified_subjects LiteLLM_VerifiedSubject[] user_id String @id user_alias String? team_id String? @@ -675,6 +731,7 @@ model LiteLLM_SpendLogs { session_id String? status String? mcp_namespaced_tool_name String? + billing_agent_id String? agent_id String? proxy_server_request Json? @default("{}") litellm_call_id String? diff --git a/tests/test_litellm/proxy/agent_endpoints/test_identity.py b/tests/test_litellm/proxy/agent_endpoints/test_identity.py new file mode 100644 index 00000000000..c9d803fdae7 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_identity.py @@ -0,0 +1,27 @@ +from collections.abc import Mapping + +import pytest +from fastapi import HTTPException + +from litellm.proxy.agent_endpoints.identity import has_legacy_identity, reject_legacy_identity + +TENANT = "11111111-1111-4111-8111-111111111111" +CLIENT = "22222222-2222-4222-8222-222222222222" + + +@pytest.mark.parametrize("params", [None, {}, {"model": "gpt-4o", "api_key": "sk-test"}]) +def test_runtime_params_without_identity_are_accepted(params: Mapping[str, object] | None) -> None: + assert has_legacy_identity(params) is False + reject_legacy_identity(params) + + +@pytest.mark.parametrize( + "identity", [None, {}, {"provider": "microsoft_entra", "tenant_id": TENANT, "client_id": CLIENT}] +) +def test_legacy_litellm_params_identity_is_rejected(identity: object) -> None: + params: Mapping[str, object] = {"model": "gpt-4o", "identity": identity} + assert has_legacy_identity(params) is True + with pytest.raises(HTTPException) as failure: + reject_legacy_identity(params) + assert failure.value.status_code == 400 + assert "top-level identity field" in failure.value.detail diff --git a/tests/test_litellm/proxy/agent_endpoints/test_identity_store.py b/tests/test_litellm/proxy/agent_endpoints/test_identity_store.py new file mode 100644 index 00000000000..005f0b4c074 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_identity_store.py @@ -0,0 +1,450 @@ +from datetime import datetime, timezone +from types import SimpleNamespace +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi import HTTPException +from prisma.models import LiteLLM_VerifiedSubject + +from litellm.proxy.agent_endpoints.identity_store import AgentIdentityStore, resolve_managed_agent +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.repositories.table_repositories import ( + AgentIdentityRepository, + AgentsRepository, + VerifiedSubjectRepository, +) +from litellm.types.agents import AgentResponse +from litellm.types.proxy.agent_identity import ( + AgentIdentityBinding, + AgentIdentityFailure, + ManagedAgentContext, + MicrosoftInteractiveSubject, +) + +TENANT: Final = "11111111-1111-4111-8111-111111111111" +CLIENT: Final = "22222222-2222-4222-8222-222222222222" +PRINCIPAL: Final = "33333333-3333-4333-8333-333333333333" +HUMAN: Final = "44444444-4444-4444-8444-444444444444" +ISSUER: Final = f"https://login.microsoftonline.com/{TENANT}/v2.0" +BINDING: Final = AgentIdentityBinding( + agent_id="agent-one", + provider="microsoft_entra", + tenant_id=TENANT, + client_id=CLIENT, + service_principal_id=PRINCIPAL, + issuer=ISSUER, + required_roles=("Agent.Invoke",), + revision="revision-one", +) +CLAIMS: Final = {"iss": ISSUER, "tid": TENANT, "azp": CLIENT, "oid": PRINCIPAL, "roles": ["Agent.Invoke"]} + + +def stored_agent(**overrides: object) -> AgentResponse: + return AgentResponse.model_validate( + { + "agent_id": "agent-one", + "agent_name": "Research", + "agent_card_params": {}, + "identity": BINDING, + "identity_managed": True, + "execution_mode": "both", + **overrides, + } + ) + + +def setup_store( + agent: AgentResponse | None = stored_agent(), + human: LiteLLM_VerifiedSubject | None = None, + cache: UserApiKeyCache | None = None, +) -> tuple[AgentIdentityStore, AsyncMock, AsyncMock, AsyncMock]: + agents: Final = AsyncMock() + identities: Final = AsyncMock() + humans: Final = AsyncMock() + agents.find_unique.return_value = agent + identities.find_unique.return_value = BINDING + identities.update_many.return_value = 1 + humans.find_unique.return_value = human + db: Final = SimpleNamespace( + db=SimpleNamespace( + litellm_agentstable=agents, + litellm_agentidentity=identities, + litellm_verifiedsubject=humans, + ) + ) + return ( + AgentIdentityStore(AgentsRepository(db), AgentIdentityRepository(db), VerifiedSubjectRepository(db), cache=cache), + agents, + identities, + humans, + ) + + +@pytest.mark.asyncio +async def test_application_authentication_has_no_fabricated_human() -> None: + store, _, _, humans = setup_store() + result: Final = await store.resolve_verified_claims(CLAIMS) + assert isinstance(result, ManagedAgentContext) + assert result.agent_id == "agent-one" + assert result.mode == "autonomous" + assert result.user_id is None + humans.upsert.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_shared_binding_lookup_cache_keeps_policy_reads_authoritative() -> None: + cache: Final = UserApiKeyCache() + store, agents, identities, _ = setup_store(cache=cache) + other: Final = AgentIdentityStore(store.agents, store.identities, store.humans, cache=cache) + assert isinstance(await store.resolve_verified_claims(CLAIMS), ManagedAgentContext) + assert isinstance(await other.resolve_verified_claims(CLAIMS), ManagedAgentContext) + identities.find_unique.assert_awaited_once() + assert agents.find_unique.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "changed", + [ + None, + stored_agent(enabled=False), + stored_agent(identity=None), + stored_agent(identity_managed=False), + stored_agent(execution_mode="delegated"), + stored_agent(identity=BINDING.model_copy(update={"active": False})), + stored_agent(identity=BINDING.model_copy(update={"client_id": HUMAN, "revision": "new-binding"})), + stored_agent(identity=BINDING.model_copy(update={"required_roles": ("New.Role",), "revision": "new-policy"})), + ], +) +async def test_lifecycle_is_read_on_every_request_without_cached_allow(changed: AgentResponse | None) -> None: + store, agents, identities, _ = setup_store(cache=UserApiKeyCache()) + agents.find_unique.side_effect = [stored_agent(), changed] + assert isinstance(await store.resolve_verified_claims(CLAIMS), ManagedAgentContext) + denial: Final = await store.resolve_verified_claims(CLAIMS) + assert isinstance(denial, AgentIdentityFailure) + assert denial.code == "identity_denied" + identities.find_unique.assert_awaited_once() + assert agents.find_unique.await_count == 2 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unavailable_table", ["agents", "identities", "humans"]) +async def test_identity_store_failure_never_becomes_a_legacy_allow(unavailable_table: str) -> None: + store, agents, identities, humans = setup_store() + table: Final = {"agents": agents, "identities": identities, "humans": humans}[unavailable_table] + table.find_unique.side_effect = RuntimeError("database unavailable") + result: Final = await store.resolve_verified_claims({**CLAIMS, "oid": HUMAN, "scp": "user_impersonation"}) + assert isinstance(result, AgentIdentityFailure) + assert result.code == "policy_unavailable" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unavailable_table", ["agents", "humans"]) +async def test_cached_binding_cannot_hide_authoritative_storage_failure(unavailable_table: str) -> None: + store, agents, identities, humans = setup_store(cache=UserApiKeyCache()) + assert isinstance(await store.resolve_verified_claims(CLAIMS), ManagedAgentContext) + table: Final = {"agents": agents, "humans": humans}[unavailable_table] + table.find_unique.side_effect = ConnectionError("writer unavailable") + result: Final = await store.resolve_verified_claims({**CLAIMS, "oid": HUMAN, "scp": "user_impersonation"}) + assert isinstance(result, AgentIdentityFailure) + assert result.code == "policy_unavailable" + identities.find_unique.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_unclassified_delegated_subject_cannot_authenticate_as_a_user() -> None: + store, _, _, _ = setup_store() + result: Final = await store.resolve_verified_claims( + {**CLAIMS, "oid": HUMAN, "scp": "user_impersonation", "idtyp": "user"} + ) + assert isinstance(result, AgentIdentityFailure) + assert "first sign in" in result.message + + +@pytest.mark.asyncio +async def test_delegated_subject_uses_canonical_sso_user_not_email_claim() -> None: + human: Final = LiteLLM_VerifiedSubject( + kind="human", + subject_id="subject-one", + issuer=ISSUER, + tenant_id=TENANT, + oid=HUMAN, + user_id="canonical-user", + verified_via="sso_interactive", + verified_at=datetime.now(timezone.utc), + ) + store, _, identities, humans = setup_store(human=human, cache=UserApiKeyCache()) + result: Final = await store.resolve_verified_claims( + { + **CLAIMS, + "oid": HUMAN, + "scp": "user_impersonation", + "email": "untrusted-alias@example.com", + } + ) + assert isinstance(result, ManagedAgentContext) + assert result.mode == "delegated" + assert result.user_id == "canonical-user" + humans.find_unique.assert_awaited_once_with( + where={"issuer_tenant_id_oid": {"issuer": ISSUER, "tenant_id": TENANT, "oid": HUMAN}} + ) + humans.find_unique.return_value = None + denied: Final = await store.resolve_verified_claims({**CLAIMS, "oid": HUMAN, "scp": "user_impersonation"}) + assert isinstance(denied, AgentIdentityFailure) + assert denied.code == "identity_denied" + identities.find_unique.assert_awaited_once() + assert humans.find_unique.await_count == 2 + + +@pytest.mark.asyncio +async def test_rebinding_during_authentication_does_not_mark_new_identity_verified() -> None: + store, _, identities, _ = setup_store() + identities.update_many.return_value = 0 + context: Final = ManagedAgentContext(agent_id="agent-one", binding_revision="old-revision", mode="autonomous") + result: Final = await store.record_authentication(context) + assert isinstance(result, AgentIdentityFailure) + assert "changed" in result.message + assert identities.update_many.call_args.kwargs["where"] == { + "agent_id": "agent-one", + "revision": "old-revision", + "active": True, + "agent": {"is": {"enabled": True, "identity_managed": True}}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("agent", [None, stored_agent(identity=None), stored_agent(identity_managed=False)]) +async def test_stale_binding_cannot_bypass_lifecycle(agent: AgentResponse | None) -> None: + store, _, _, _ = setup_store(agent=agent) + assert isinstance(await store.resolve_verified_claims(CLAIMS), AgentIdentityFailure) + + +@pytest.mark.asyncio +async def test_unrelated_non_entra_claims_do_not_query_identity_store() -> None: + store, agents, identities, _ = setup_store() + assert await store.resolve_verified_claims({"sub": "ordinary-user"}) is None + identities.find_unique.assert_not_awaited() + agents.find_unique.assert_not_awaited() + + +HUMAN_CLAIMS: Final = {"iss": ISSUER, "tid": TENANT, "azp": CLIENT, "oid": HUMAN, "scp": "user_impersonation"} + + +@pytest.mark.asyncio +async def test_bound_agents_and_policy_failures_are_never_served_from_the_miss_cache() -> None: + store, _, identities, _ = setup_store() + assert isinstance(await store.resolve_verified_claims(CLAIMS), ManagedAgentContext) + assert isinstance(await store.resolve_verified_claims(CLAIMS), ManagedAgentContext) + assert identities.find_unique.await_count == 2 + identities.find_unique.side_effect = ConnectionError("database down") + assert isinstance(await store.resolve_verified_claims(CLAIMS), AgentIdentityFailure) + assert isinstance(await store.resolve_verified_claims(CLAIMS), AgentIdentityFailure) + assert identities.find_unique.await_count == 4 + + +@pytest.mark.asyncio +async def test_retired_client_cannot_fall_back_to_ordinary_user_authentication() -> None: + from prisma.models import LiteLLM_RetiredAgentIdentity + + from litellm.repositories.table_repositories import RetiredAgentIdentityRepository + + identities: Final = AsyncMock() + identities.find_unique.return_value = None + retired: Final = AsyncMock() + retired.find_unique.return_value = LiteLLM_RetiredAgentIdentity( + binding_id="retired", + agent_id="agent-one", + provider="microsoft_entra", + issuer=ISSUER, + tenant_id=TENANT, + client_id=CLIENT, + ) + db: Final = SimpleNamespace( + db=SimpleNamespace( + litellm_agentidentity=identities, + litellm_retiredagentidentity=retired, + litellm_agentstable=AsyncMock(), + litellm_verifiedsubject=AsyncMock(), + ) + ) + store: Final = AgentIdentityStore( + AgentsRepository(db), + AgentIdentityRepository(db), + VerifiedSubjectRepository(db), + RetiredAgentIdentityRepository(db), + ) + result: Final = await store.resolve_verified_claims({**CLAIMS, "oid": HUMAN, "scp": "user_impersonation"}) + assert isinstance(result, AgentIdentityFailure) + assert result.code == "identity_denied" + assert "retired" in result.message + + +@pytest.mark.asyncio +async def test_missing_revision_cannot_create_entra_authentication_evidence() -> None: + store, _, identities, _ = setup_store() + result: Final = await store.record_authentication(ManagedAgentContext(agent_id="agent-one", mode="autonomous")) + assert isinstance(result, AgentIdentityFailure) + assert result.code == "identity_denied" + identities.update_many.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_authentication_evidence_write_failure_is_not_success() -> None: + store, _, identities, _ = setup_store() + identities.update_many.side_effect = RuntimeError("writer unavailable") + result: Final = await store.record_authentication( + ManagedAgentContext(agent_id="agent-one", binding_revision="revision-one", mode="autonomous") + ) + assert isinstance(result, AgentIdentityFailure) + assert result.code == "policy_unavailable" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("unavailable", [True, False]) +async def test_retired_binding_denies_and_history_outage_cannot_become_legacy_fallback(unavailable: bool) -> None: + + database: Final = MagicMock() + database.writer_db.litellm_agentidentity.find_unique = AsyncMock(return_value=None) + database.writer_db.litellm_verifiedsubject.find_unique = AsyncMock(return_value=None) + database.writer_db.litellm_retiredagentidentity.find_unique = AsyncMock( + return_value={"client_id": CLIENT}, side_effect=RuntimeError("unavailable") if unavailable else None + ) + result: Final = await AgentIdentityStore.from_client(database).resolve_verified_claims(CLAIMS) + assert isinstance(result, AgentIdentityFailure) + assert result.code == ("policy_unavailable" if unavailable else "identity_denied") + assert result.message == ( + "Retired agent identity could not be checked" if unavailable else "This agent identity binding has been retired" + ) + + +@pytest.mark.asyncio +async def test_new_binding_is_enforced_after_another_worker_commits_it() -> None: + _, agents, identities, humans = setup_store() + identities.find_unique.return_value = None + retired: Final = AsyncMock() + retired.find_unique.return_value = None + db: Final = SimpleNamespace( + writer_db=SimpleNamespace( + litellm_agentstable=agents, + litellm_agentidentity=identities, + litellm_verifiedsubject=humans, + litellm_retiredagentidentity=retired, + litellm_retiredagent=retired, + ) + ) + worker: Final = AgentIdentityStore.from_client(db, cache=UserApiKeyCache()) + claims: Final = {**CLAIMS, "oid": "55555555-5555-4555-8555-555555555555"} + assert await worker.resolve_verified_claims(claims) is None + identities.find_unique.return_value = BINDING + denied: Final = await worker.resolve_verified_claims(claims) + assert isinstance(denied, AgentIdentityFailure) + assert denied.code == "identity_denied" + assert "Application token contradicts" in denied.message + assert identities.find_unique.await_count == 2 + + +@pytest.mark.asyncio +async def test_non_string_subject_does_not_query_directory_ownership() -> None: + store, _, _, humans = setup_store() + assert await store.subject(ISSUER, TENANT, None) is None + humans.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("configured", [False, True]) +async def test_missing_or_unavailable_retirement_history_fails_closed(configured: bool) -> None: + database: Final = MagicMock() + database.writer_db.litellm_retiredagent.find_unique = AsyncMock(side_effect=RuntimeError("history unavailable")) + store: Final = AgentIdentityStore.from_client(database) if configured else setup_store()[0] + result: Final = await store.retired_agent("deleted-agent") + assert isinstance(result, AgentIdentityFailure) + assert result.code == "policy_unavailable" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("owner", ["canonical-user", "another-user"]) +async def test_interactive_enrollment_preserves_existing_subject_ownership(owner: str) -> None: + store, _, _, humans = setup_store() + humans.upsert.return_value = LiteLLM_VerifiedSubject( + subject_id="subject-one", + issuer=ISSUER, + tenant_id=TENANT, + oid=HUMAN, + user_id=owner, + kind="human", + verified_via="sso_interactive", + verified_at=datetime.now(timezone.utc), + ) + result: Final = await store.enroll_interactive_human( + MicrosoftInteractiveSubject(issuer=ISSUER, tenant_id=TENANT, oid=HUMAN), "canonical-user" + ) + if owner == "canonical-user": + assert result is None + else: + assert isinstance(result, AgentIdentityFailure) + assert result.code == "identity_denied" + assert humans.upsert.call_args.kwargs["data"]["update"] == {} + assert humans.upsert.call_args.kwargs["data"]["create"]["user_id"] == "canonical-user" + + +@pytest.mark.asyncio +async def test_interactive_enrollment_outage_fails_closed() -> None: + store, _, _, humans = setup_store() + humans.upsert.side_effect = ConnectionError("writer unavailable") + result: Final = await store.enroll_interactive_human( + MicrosoftInteractiveSubject(issuer=ISSUER, tenant_id=TENANT, oid=HUMAN), "canonical-user" + ) + assert isinstance(result, AgentIdentityFailure) + assert result.code == "policy_unavailable" + + +@pytest.mark.asyncio +async def test_matching_revision_records_successful_authentication() -> None: + store, _, identities, _ = setup_store() + assert ( + await store.record_authentication( + ManagedAgentContext(agent_id="agent-one", binding_revision="revision-one", mode="autonomous") + ) + is None + ) + identities.update_many.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outage", [False, True]) +async def test_resolver_maps_denials_and_outages_to_public_errors(outage: bool) -> None: + database: Final = MagicMock() + database.writer_db.litellm_agentidentity.find_unique = AsyncMock( + return_value=BINDING, side_effect=ConnectionError("unavailable") if outage else None + ) + database.writer_db.litellm_verifiedsubject.find_unique = AsyncMock(return_value=None) + database.writer_db.litellm_agentstable.find_unique = AsyncMock(return_value=stored_agent(enabled=False)) + with pytest.raises(HTTPException) as exc: + await resolve_managed_agent(CLAIMS, database) + assert exc.value.status_code == (503 if outage else 403) + + +@pytest.mark.asyncio +async def test_resolver_preserves_unconfigured_and_unrelated_authentication() -> None: + assert await resolve_managed_agent(CLAIMS, None) is None + assert await resolve_managed_agent({"sub": "ordinary-user"}, MagicMock()) is None + store, _, identities, _ = setup_store() + identities.find_unique.return_value = None + assert await store.resolve_verified_claims(CLAIMS) is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("registered", [True, False]) +async def test_application_and_unregistered_clients_do_not_depend_on_human_subject_storage(registered: bool) -> None: + store, _, identities, humans = setup_store() + identities.find_unique.return_value = BINDING if registered else None + humans.find_unique.side_effect = RuntimeError("subject database unavailable") + result: Final = await store.resolve_verified_claims(CLAIMS) + if registered: + assert isinstance(result, ManagedAgentContext) + assert result.mode == "autonomous" + assert result.user_id is None + else: + assert result is None + humans.find_unique.assert_not_awaited() diff --git a/tests/test_litellm/proxy/agent_endpoints/test_managed_identity.py b/tests/test_litellm/proxy/agent_endpoints/test_managed_identity.py new file mode 100644 index 00000000000..45fe4b0655f --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_managed_identity.py @@ -0,0 +1,257 @@ +from typing import Final + +import pytest + +from litellm.proxy.agent_endpoints.managed_identity import classify_agent_subject, managed_write_fields +from litellm.types.agents import AgentResponse +from litellm.types.proxy.agent_identity import ( + AgentExecutionMode, + AgentIdentityBinding, + AgentIdentityFailure, + AgentSubject, +) + +TENANT: Final = "11111111-1111-4111-8111-111111111111" +CLIENT: Final = "22222222-2222-4222-8222-222222222222" +PRINCIPAL: Final = "33333333-3333-4333-8333-333333333333" +HUMAN: Final = "44444444-4444-4444-8444-444444444444" +ISSUER: Final = f"https://login.microsoftonline.com/{TENANT}/v2.0" +BINDING: Final = AgentIdentityBinding( + agent_id="agent-one", + provider="microsoft_entra", + tenant_id=TENANT, + client_id=CLIENT, + service_principal_id=PRINCIPAL, + issuer=ISSUER, + required_roles=("Agent.Invoke",), + required_scopes=("user_impersonation",), + revision="binding-one", +) + + +def claims(**overrides: object) -> dict[str, object]: + return {"iss": ISSUER, "tid": TENANT, "azp": CLIENT, "oid": PRINCIPAL, "roles": ["Agent.Invoke"], **overrides} + + +def test_autonomous_identity_needs_no_human_and_checks_the_pinned_principal() -> None: + result: Final = classify_agent_subject(BINDING, claims(), "autonomous") + assert result == AgentSubject(kind="application", oid=PRINCIPAL, mode="autonomous") + assert isinstance(classify_agent_subject(BINDING, claims(oid=HUMAN), "autonomous"), AgentIdentityFailure) + + +@pytest.mark.parametrize( + "overrides", + [ + {"iss": "https://untrusted.example"}, + {"tid": CLIENT}, + {"azp": TENANT}, + {"roles": []}, + {"idtyp": "user"}, + {"scp": "user_impersonation"}, + {"scp": 1}, + {"oid": None}, + ], +) +def test_application_rejects_mismatched_or_contradictory_verified_claims(overrides: dict[str, object]) -> None: + assert isinstance(classify_agent_subject(BINDING, claims(**overrides), "both"), AgentIdentityFailure) + + +def test_delegated_profile_identifies_a_subject_without_asserting_that_it_is_human() -> None: + result: Final = classify_agent_subject(BINDING, claims(oid=HUMAN, scp="user_impersonation"), "delegated") + assert result == AgentSubject(kind="delegated_subject", oid=HUMAN, mode="delegated") + + +@pytest.mark.parametrize( + "overrides", + [ + {"scp": "unrelated"}, + {"scp": ""}, + {"idtyp": "app"}, + {"xms_sub_fct": "2 13 15"}, + {"xms_sub_fct": [13]}, + ], +) +def test_delegated_profile_rejects_unknown_scope_and_known_nonhuman_subjects(overrides: dict[str, object]) -> None: + assert isinstance( + classify_agent_subject(BINDING, claims(**{"oid": HUMAN, "scp": "user_impersonation", **overrides}), "both"), + AgentIdentityFailure, + ) + + +def test_allowed_mode_cannot_be_selected_by_the_caller() -> None: + assert isinstance(classify_agent_subject(BINDING, claims(), "delegated"), AgentIdentityFailure) + assert isinstance( + classify_agent_subject(BINDING, claims(oid=HUMAN, scp="user_impersonation"), "autonomous"), + AgentIdentityFailure, + ) + + +def test_native_facet_absence_does_not_establish_human_identity() -> None: + result: Final = classify_agent_subject( + BINDING, claims(oid=HUMAN, scp="user_impersonation", xms_sub_fct="113"), "both" + ) + assert isinstance(result, AgentSubject) + assert result.kind == "delegated_subject" + + +def managed_agent() -> AgentResponse: + return AgentResponse( + agent_id="agent-one", agent_name="Research", agent_card_params={}, identity=BINDING, identity_managed=True + ) + + +def test_unbinding_keeps_managed_state_and_disables_agent() -> None: + result: Final = managed_write_fields({"identity": None, "enabled": True}, managed_agent(), "admin") + assert not isinstance(result, AgentIdentityFailure) + assert result["identity_managed"] is True + assert result["enabled"] is False + assert result["identity"]["update"]["active"] is False + assert result["identity"]["update"]["last_authenticated_at"] is None + assert result["identity"]["update"]["revision"] != BINDING.revision + + +def test_rename_does_not_rewrite_binding_or_evidence() -> None: + assert managed_write_fields({"agent_name": "Renamed"}, managed_agent(), "admin") == {} + + +def test_autonomous_binding_requires_enterprise_application_object_id() -> None: + result: Final = managed_write_fields( + {"identity": {"provider": "microsoft_entra", "tenant_id": TENANT, "client_id": CLIENT}}, None, "admin" + ) + assert isinstance(result, AgentIdentityFailure) + assert "service-principal" in result.message + + +def test_rebinding_clears_evidence_and_uses_atomic_nested_write() -> None: + result: Final = managed_write_fields( + { + "identity": { + "provider": "microsoft_entra", + "tenant_id": TENANT, + "client_id": CLIENT, + "service_principal_id": PRINCIPAL, + } + }, + managed_agent(), + "admin", + ) + assert not isinstance(result, AgentIdentityFailure) + assert result["identity_managed"] is True + assert "upsert" in result["identity"] + assert result["identity"]["upsert"]["update"]["revision"] != BINDING.revision + assert result["identity"]["upsert"]["update"]["last_authenticated_at"] is None + + +def test_unbound_identity_can_be_reactivated_with_the_same_application() -> None: + disabled: Final = managed_agent().model_copy( + update={"identity": BINDING.model_copy(update={"active": False}), "enabled": False} + ) + configuration: Final = BINDING.model_dump( + exclude={"agent_id", "issuer", "revision", "last_authenticated_at", "active"} + ) + result: Final = managed_write_fields({"identity": configuration, "enabled": True}, disabled, "admin") + assert not isinstance(result, AgentIdentityFailure) + assert result["enabled"] is True + assert result["identity"]["upsert"]["update"]["active"] is True + assert result["identity"]["upsert"]["update"]["revision"] != BINDING.revision + + +def test_each_application_binding_records_its_history_atomically() -> None: + configuration: Final = BINDING.model_dump( + exclude={"agent_id", "issuer", "revision", "last_authenticated_at", "active"} + ) + created: Final = managed_write_fields({"identity": configuration}, None, "admin") + assert not isinstance(created, AgentIdentityFailure) + assert created["retired_identities"]["connectOrCreate"]["create"]["client_id"] == CLIENT + replacement: Final = managed_write_fields( + {"identity": {**configuration, "client_id": HUMAN}}, managed_agent(), "admin" + ) + assert not isinstance(replacement, AgentIdentityFailure) + assert replacement["retired_identities"]["connectOrCreate"]["create"]["client_id"] == HUMAN + + +def test_unchanged_binding_preserves_revision_and_authentication_evidence() -> None: + configuration: Final = BINDING.model_dump( + exclude={"agent_id", "issuer", "revision", "last_authenticated_at", "active"} + ) + assert managed_write_fields({"identity": configuration}, managed_agent(), "admin") == {} + + +@pytest.mark.parametrize("identity", [None, BINDING.model_copy(update={"active": False})]) +def test_enabling_unbound_or_inactive_identity_requires_rebinding(identity: AgentIdentityBinding | None) -> None: + agent: Final = managed_agent().model_copy(update={"identity": identity, "enabled": False}) + result: Final = managed_write_fields({"enabled": True}, agent, "admin") + assert isinstance(result, AgentIdentityFailure) + assert "Bind an identity" in result.message + + +@pytest.mark.parametrize("mode", ["delegated", "both"]) +def test_explicit_empty_scope_requirements_can_be_registered_and_preserved(mode: str) -> None: + from litellm.types.proxy.agent_identity import EntraIdentityConfig + + configuration: Final = EntraIdentityConfig( + provider="microsoft_entra", + tenant_id=TENANT, + client_id=CLIENT, + service_principal_id=PRINCIPAL, + required_scopes=(), + ) + created: Final = managed_write_fields( + {"identity": configuration.model_dump(), "execution_mode": mode}, None, "admin" + ) + assert not isinstance(created, AgentIdentityFailure) + assert created["identity"]["create"]["required_scopes"] == () + agent: Final = managed_agent().model_copy(update={"identity": BINDING.model_copy(update={"required_scopes": ()})}) + updated: Final = managed_write_fields({"execution_mode": mode}, agent, "admin") + assert not isinstance(updated, AgentIdentityFailure) + assert updated["execution_mode"] == mode + + +@pytest.mark.parametrize( + "incoming", + [ + {"identity": {"provider": "microsoft_entra", "tenant_id": "invalid", "client_id": CLIENT}}, + {"execution_mode": "unknown"}, + ], +) +def test_invalid_identity_configuration_returns_a_public_validation_failure(incoming: dict[str, object]) -> None: + result: Final = managed_write_fields(incoming, None, "admin") + assert isinstance(result, AgentIdentityFailure) + assert result.code == "identity_denied" + assert result.message.startswith("Invalid agent identity configuration:") + + +@pytest.mark.parametrize("roles", ["Agent.Invoke", [42], None]) +def test_malformed_application_roles_are_rejected(roles: object) -> None: + result: Final = classify_agent_subject(BINDING, claims(roles=roles), "autonomous") + assert isinstance(result, AgentIdentityFailure) + assert "Invalid application roles" in result.message + + +def test_entra_binding_normalizes_identifiers_and_rejects_invalid_configuration() -> None: + from pydantic import ValidationError + + from litellm.types.proxy.agent_identity import EntraIdentityConfig + + identifier = "ABCDEF00-1234-4234-9234-123456789ABC" + config = EntraIdentityConfig(provider="microsoft_entra", tenant_id=identifier, client_id=identifier) + assert config.tenant_id == identifier.lower() + assert config.client_id == identifier.lower() + assert config.service_principal_id is None + assert config.issuer == f"https://login.microsoftonline.com/{config.tenant_id}/v2.0" + with pytest.raises(ValidationError): + EntraIdentityConfig(provider="microsoft_entra", tenant_id="invalid", client_id=identifier) + + +@pytest.mark.parametrize("mode", ["delegated", "both"]) +def test_empty_required_scopes_allow_valid_delegated_scope(mode: AgentExecutionMode) -> None: + binding: Final = BINDING.model_copy(update={"required_scopes": ()}) + result: Final = classify_agent_subject(binding, claims(oid=HUMAN, scp="custom_scope"), mode) + assert result == AgentSubject(kind="delegated_subject", oid=HUMAN, mode="delegated") + + +@pytest.mark.parametrize("scope", [None, "", " \t ", 42]) +def test_empty_requirements_do_not_make_a_scope_less_human_token_valid(scope: object) -> None: + binding: Final = BINDING.model_copy(update={"required_scopes": ()}) + result: Final = classify_agent_subject(binding, claims(oid=HUMAN, scp=scope), "both") + assert isinstance(result, AgentIdentityFailure) diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 347adc421a2..7db5d7ca9a3 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -3762,7 +3762,7 @@ class TestSpendLogsPayload: "model": "gpt-4o", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', + "metadata": '{"actor_agent_id": null, "target_agent_id": null, "billing_agent_id": null, "agent_execution_mode": null, "verified_human_user_id": null, "applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, diff --git a/tests/test_litellm/proxy/test__types.py b/tests/test_litellm/proxy/test__types.py index b43a75d3323..adc3bc04bdf 100644 --- a/tests/test_litellm/proxy/test__types.py +++ b/tests/test_litellm/proxy/test__types.py @@ -20,6 +20,14 @@ from litellm.proxy._types import ( ) SERVER_ONLY_MARKERS = ( + "requires_fresh_policy", + "mcp_explicit_grants_only", + "managed_agent_context", + "managed_agent_policy", + "invoked_agent_id", + "invoked_agent_policy", + "agent_invocation_cost", + "billing_agent_policy", "mcp_admitted_user_subject", "mcp_source_team_rpm_limits", "mcp_session_resource_server_id", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index ecfabf33e87..8e11a9234e1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -24285,6 +24285,45 @@ export interface components { /** Uri */ uri?: string; }; + /** AgentIdentityBinding */ + AgentIdentityBinding: { + /** + * Active + * @default true + */ + active: boolean; + /** Agent Id */ + agent_id: string; + /** Client Id */ + client_id: string; + /** Issuer */ + issuer: string; + /** Last Authenticated At */ + last_authenticated_at?: string | null; + /** + * Provider + * @constant + */ + provider: "microsoft_entra"; + /** + * Required Roles + * @default [] + */ + required_roles: string[]; + /** + * Required Scopes + * @default [ + * "user_impersonation" + * ] + */ + required_scopes: string[]; + /** Revision */ + revision: string; + /** Service Principal Id */ + service_principal_id?: string | null; + /** Tenant Id */ + tenant_id: string; + }; /** * AgentInterface * @description Declares a combination of a target URL and a transport protocol. @@ -24440,8 +24479,30 @@ export interface components { created_at?: string | null; /** Created By */ created_by?: string | null; + /** + * Enabled + * @default true + */ + enabled: boolean; + /** + * Execution Mode + * @default autonomous + * @enum {string} + */ + execution_mode: "autonomous" | "delegated" | "both"; /** Extra Headers */ extra_headers?: string[] | null; + identity?: components["schemas"]["AgentIdentityBinding"] | null; + /** + * Identity Managed + * @default false + */ + identity_managed: boolean; + /** + * Jwt Auth Configured + * @default false + */ + jwt_auth_configured: boolean; /** Keys */ keys?: components["schemas"]["AgentKeySummary"][] | null; kill_switch?: components["schemas"]["AgentKillSwitchConfig"] | null;