mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
feat(agents): add identity storage and validation contracts (#43720)
* feat(agents): identity storage and contracts * fix(agents): cache positive identity lookups with fresh policy checks * test(agents): include identity attribution in spend fixture --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
parent
13d004fc5a
commit
6684256136
18 changed files with 1831 additions and 4 deletions
|
|
@ -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 $$;
|
||||
|
|
@ -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?
|
||||
|
|
|
|||
|
|
@ -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": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
||||
|
||||
|
|
|
|||
17
litellm/proxy/agent_endpoints/identity.py
Normal file
17
litellm/proxy/agent_endpoints/identity.py
Normal file
|
|
@ -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)
|
||||
250
litellm/proxy/agent_endpoints/identity_store.py
Normal file
250
litellm/proxy/agent_endpoints/identity_store.py
Normal file
|
|
@ -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
|
||||
220
litellm/proxy/agent_endpoints/managed_identity.py
Normal file
220
litellm/proxy/agent_endpoints/managed_identity.py
Normal file
|
|
@ -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")
|
||||
|
|
@ -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?
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
96
litellm/types/proxy/agent_identity.py
Normal file
96
litellm/types/proxy/agent_identity.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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?
|
||||
|
|
|
|||
27
tests/test_litellm/proxy/agent_endpoints/test_identity.py
Normal file
27
tests/test_litellm/proxy/agent_endpoints/test_identity.py
Normal file
|
|
@ -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
|
||||
450
tests/test_litellm/proxy/agent_endpoints/test_identity_store.py
Normal file
450
tests/test_litellm/proxy/agent_endpoints/test_identity_store.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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)
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
61
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
61
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue