From 6370104c53ef5ae02a5e499993461d2135dbd5ce Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 3 Oct 2026 17:28:18 -0700 Subject: [PATCH] feat(enterprise): bundle LiteAdmin Slack with native gateway login (#44444) * feat(enterprise): bundle LiteAdmin Slack worker with native gateway login * fix(enterprise): preserve gateway prefixes during native Slack linking * fix(enterprise): retain native Slack linking on the admin backend * fix(enterprise): reuse shared native Slack connection services --- .circleci/scripts/unit_selection.sh | 1 + Dockerfile | 11 + backend/routes/allowlist.py | 1 + docker-compose.liteadmin.yml | 44 +++ docker/prod_entrypoint.sh | 6 + .../proxy/enterprise_routes.py | 2 + .../litellm_enterprise/proxy/liteadmin.py | 283 ++++++++++++++++++ helm/litellm-helm/templates/deployment.yaml | 13 + helm/litellm-helm/templates/liteadmin.yaml | 112 +++++++ helm/litellm-helm/values.yaml | 14 + .../mcp_server/byok_oauth_endpoints.py | 36 +++ tests/unit/enterprise/proxy/test_liteadmin.py | 244 +++++++++++++++ .../mcp_server/test_byok_oauth_endpoints.py | 47 ++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 80 +++++ 14 files changed, 893 insertions(+), 1 deletion(-) create mode 100644 docker-compose.liteadmin.yml create mode 100644 enterprise/litellm_enterprise/proxy/liteadmin.py create mode 100644 helm/litellm-helm/templates/liteadmin.yaml create mode 100644 tests/unit/enterprise/proxy/test_liteadmin.py diff --git a/.circleci/scripts/unit_selection.sh b/.circleci/scripts/unit_selection.sh index 542984dd2e0..02df32d5eab 100755 --- a/.circleci/scripts/unit_selection.sh +++ b/.circleci/scripts/unit_selection.sh @@ -41,6 +41,7 @@ legacy_paths() { echo tests/unit/enterprise/proxy/hooks echo tests/unit/enterprise/proxy/management_endpoints echo tests/unit/enterprise/proxy/test_audit_logging_endpoints.py + echo tests/unit/enterprise/proxy/test_liteadmin.py echo tests/unit/enterprise/enterprise_callbacks/test_prometheus_logging_callbacks.py ;; enterprise-routing) echo tests/unit/google_genai diff --git a/Dockerfile b/Dockerfile index 35dbaa4d41b..be507f6efb4 100644 --- a/Dockerfile +++ b/Dockerfile @@ -114,6 +114,16 @@ RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/ RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \ sed -i 's/\r$//' docker/prod_entrypoint.sh && chmod +x docker/prod_entrypoint.sh +FROM $LITELLM_BUILD_IMAGE AS liteadmin-builder +COPY --from=uvbin /uv /usr/local/bin/uv +RUN apk add --no-cache python-3.13 +ADD --checksum=sha256:2f7ae5cdd9d91731c0990e74a58239dc3e3fd2bf28dab23b55eafcdc47aaf87e \ + https://github.com/BerriAI/litellm-admin-agent/archive/ef501e94bc9fbacb9233b922abf71427f030408c.tar.gz /tmp/liteadmin.tar.gz +RUN mkdir /tmp/liteadmin && tar xzf /tmp/liteadmin.tar.gz --strip-components=1 -C /tmp/liteadmin && \ + uv venv /opt/liteadmin --python python3.13 && \ + uv pip install --python /opt/liteadmin/bin/python --require-hashes -r /tmp/liteadmin/requirements.txt && \ + uv pip install --python /opt/liteadmin/bin/python --no-deps /tmp/liteadmin + # Runtime stage FROM $LITELLM_RUNTIME_IMAGE AS runtime ARG LITELLM_RELEASE_TAG="" @@ -143,6 +153,7 @@ ENV PATH="/app/.venv/bin:${PATH}" \ # ship (manifest-scanning tools attribute everything in it to this image). # entrypoint.sh invokes litellm/proxy/prisma_migration.py by source path. COPY --from=builder /app/.venv /app/.venv +COPY --from=liteadmin-builder /opt/liteadmin /opt/liteadmin COPY --from=builder /app/docker /app/docker COPY --from=builder /app/schema.prisma /app/schema.prisma COPY --from=builder /app/litellm/proxy/prisma_migration.py /app/litellm/proxy/prisma_migration.py diff --git a/backend/routes/allowlist.py b/backend/routes/allowlist.py index d7f3e615c67..2c651277dab 100644 --- a/backend/routes/allowlist.py +++ b/backend/routes/allowlist.py @@ -22,6 +22,7 @@ BACKEND_PATH_PREFIXES: tuple[str, ...] = ( "/customer/", "/end_user/", "/sso/", + "/liteadmin/slack/connect/", "/login", "/v2/login", "/v3/login", diff --git a/docker-compose.liteadmin.yml b/docker-compose.liteadmin.yml new file mode 100644 index 00000000000..a66846429de --- /dev/null +++ b/docker-compose.liteadmin.yml @@ -0,0 +1,44 @@ +services: + litellm: + image: ${LITELLM_IMAGE:?Set the native-enabled gateway image} + environment: + LITELLM_ADMIN_AGENT_URL: http://liteadmin:10000 + ADMIN_AGENT_SERVICE_TOKEN: ${ADMIN_AGENT_SERVICE_TOKEN:?Set a shared worker token} + PROXY_BASE_URL: ${LITELLM_PUBLIC_URL:?Set the existing HTTPS gateway URL} + + liteadmin: + image: ${LITELLM_IMAGE:?Set the same native-enabled image used by the gateway} + command: ["--admin-agent"] + restart: unless-stopped + init: true + read_only: true + cap_drop: [ALL] + security_opt: [no-new-privileges:true] + stop_grace_period: 75s + environment: + CONNECTION_AUTH_MODE: native + LITELLM_BASE_URL: ${LITELLM_PUBLIC_URL:?Set the existing HTTPS gateway URL} + LITELLM_MODEL: ${LITELLM_ADMIN_MODEL:?Set a gateway model with tool support} + SLACK_BOT_TOKEN: ${SLACK_BOT_TOKEN:?Install the Slack app} + SLACK_APP_TOKEN: ${SLACK_APP_TOKEN:?Enable Socket Mode} + SLACK_WORKSPACE_ID: ${SLACK_WORKSPACE_ID:?Set the Slack workspace ID} + ADMIN_AGENT_SERVICE_TOKEN: ${ADMIN_AGENT_SERVICE_TOKEN:?Set a shared worker token} + CREDENTIAL_ENCRYPTION_KEY: ${CREDENTIAL_ENCRYPTION_KEY:?Set a persistent Fernet key} + STATE_DB: /var/data/events.sqlite3 + ADMIN_READ_ONLY: ${ADMIN_READ_ONLY:-false} + OPENAI_AGENTS_DISABLE_TRACING: "1" + volumes: + - liteadmin_state:/var/data + tmpfs: + - /tmp:rw,noexec,nosuid,size=64m + healthcheck: + test: ["CMD", "/opt/liteadmin/bin/python", "-c", "import urllib.request; urllib.request.urlopen('http://127.0.0.1:10000/readyz', timeout=3)"] + interval: 30s + timeout: 5s + start_period: 30s + depends_on: + litellm: + condition: service_healthy + +volumes: + liteadmin_state: diff --git a/docker/prod_entrypoint.sh b/docker/prod_entrypoint.sh index 630eb6b065b..4386be65a32 100644 --- a/docker/prod_entrypoint.sh +++ b/docker/prod_entrypoint.sh @@ -1,5 +1,11 @@ #!/bin/sh +if [ "$1" = "--admin-agent" ]; then + shift + export CONNECTION_AUTH_MODE=native + exec /opt/liteadmin/bin/litellm-admin-agent --web "$@" +fi + case "$USE_DDTRACE" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED="False" diff --git a/enterprise/litellm_enterprise/proxy/enterprise_routes.py b/enterprise/litellm_enterprise/proxy/enterprise_routes.py index ec37c049809..a76b8d01f0e 100644 --- a/enterprise/litellm_enterprise/proxy/enterprise_routes.py +++ b/enterprise/litellm_enterprise/proxy/enterprise_routes.py @@ -6,6 +6,7 @@ from litellm_enterprise.enterprise_callbacks.send_emails.endpoints import ( from . import ui_crud_endpoints # side-effect: registers extra UI settings from .audit_logging_endpoints import router as audit_logging_router +from .liteadmin import router as liteadmin_router from .management_endpoints import management_endpoints_router from .utils import _should_block_robots @@ -14,6 +15,7 @@ __all__ = ["router", "ui_crud_endpoints"] router = APIRouter() router.include_router(email_events_router) router.include_router(audit_logging_router) +router.include_router(liteadmin_router) router.include_router(management_endpoints_router) diff --git a/enterprise/litellm_enterprise/proxy/liteadmin.py b/enterprise/litellm_enterprise/proxy/liteadmin.py new file mode 100644 index 00000000000..6a9110f1460 --- /dev/null +++ b/enterprise/litellm_enterprise/proxy/liteadmin.py @@ -0,0 +1,283 @@ +from __future__ import annotations + +import hashlib +import hmac +import html +import os +import re +import secrets +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Annotated, Final +from urllib.parse import urlencode, urlsplit + +import httpx +from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi.responses import HTMLResponse, RedirectResponse, Response +from pydantic import BaseModel, ConfigDict, Field, SecretStr, TypeAdapter, ValidationError + +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.proxy._experimental.mcp_server.oauth_utils import get_request_base_url +from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth +from litellm.types.proxy.auth.auth_checks import UserNotFoundError + +router: Final = APIRouter() +_PREFIX: Final = "/liteadmin/slack/connect/" +_COOKIE: Final = "__Host-litellm-slack-connect-" +_HEADERS: Final = { + "Cache-Control": "no-store", + "Referrer-Policy": "same-origin", + "X-Frame-Options": "DENY", + "X-Content-Type-Options": "nosniff", + "Content-Security-Policy": "default-src 'none'; style-src 'unsafe-inline'; form-action 'self'; frame-ancestors 'none'; base-uri 'none'", +} + + +class LinkDetails(BaseModel): + model_config = ConfigDict(frozen=True, strict=True, extra="forbid") + workspace_id: str = Field(min_length=1, max_length=64) + slack_user_id: str = Field(min_length=1, max_length=64) + email: str = Field(min_length=1, max_length=320) + + +class AdminSession(BaseModel): + model_config = ConfigDict(frozen=True) + user_id: str + credential: SecretStr + expires_at: float + + +@dataclass(frozen=True, slots=True) +class NativeAdminContext: + worker_url: str + service_token: SecretStr + client: httpx.AsyncClient + session_user: Callable[[Request], Awaitable[str | None]] + load_user: Callable[[str], Awaitable[LiteLLM_UserTable | None]] + mint_session: Callable[[LiteLLM_UserTable], AdminSession] + + async def worker_request(self, token: str, session: AdminSession | None = None) -> httpx.Response: + if re.fullmatch(r"[A-Za-z0-9_-]{43}", token) is None: + raise HTTPException(410, "Connection link expired. Send connect in Slack for a new link") + try: + response: Final = await self.client.request( + "GET" if session is None else "POST", + f"{self.worker_url}/internal/liteadmin/links/{token}", + headers={"X-LiteLLM-Admin-Agent-Token": self.service_token.get_secret_value()}, + json=None + if session is None + else { + "user_id": session.user_id, + "credential": session.credential.get_secret_value(), + "expires_at": session.expires_at, + }, + timeout=15, + follow_redirects=False, + ) + except httpx.HTTPError: + raise HTTPException(503, "LiteAdmin is temporarily unavailable") from None + if response.status_code == 410: + raise HTTPException(410, "Connection link expired. Send connect in Slack for a new link") + if response.status_code == 403: + raise HTTPException(403, "Connect your own active LiteLLM proxy-admin account with the same email as Slack") + if response.status_code != 200: + raise HTTPException(503, "LiteAdmin could not verify this connection") + return response + + async def details(self, token: str) -> LinkDetails: + response: Final = await self.worker_request(token) + try: + return LinkDetails.model_validate_json(response.content) + except ValidationError: + raise HTTPException(503, "LiteAdmin could not verify this connection") from None + + async def admin(self, user_id: str, details: LinkDetails) -> LiteLLM_UserTable: + user: Final = await self.load_user(user_id) + if ( + user is None + or user.user_role != LitellmUserRoles.PROXY_ADMIN.value + or not user.user_email + or user.user_email.strip().casefold() != details.email.strip().casefold() + ): + raise HTTPException(403, "Connect your own active LiteLLM proxy-admin account with the same email as Slack") + return user + + +def _page(title: str, body: str) -> HTMLResponse: + return HTMLResponse( + f'' + f'{html.escape(title)}' + "" + f"

{html.escape(title)}

{body}
", + headers=_HEADERS, + ) + + +def _cookie_name(token: str) -> str: + return _COOKIE + hashlib.sha256(token.encode()).hexdigest()[:16] + + +async def _session_user(request: Request) -> str | None: + from litellm.proxy._experimental.mcp_server.byok_oauth_endpoints import ( + get_authenticated_browser_user_id, + ) + + return await get_authenticated_browser_user_id(request) + + +async def _load_user(user_id: str) -> LiteLLM_UserTable | None: + from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if prisma_client is None: + raise HTTPException(503, "LiteAdmin requires a database") + try: + return await get_user_object( + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + check_db_only=True, + ) + except UserNotFoundError: + return None + except Exception: + raise HTTPException(503, "LiteAdmin could not verify your current permissions") from None + + +def mint_admin_session(user: LiteLLM_UserTable) -> AdminSession: + from litellm.proxy.auth.auth_checks import LITELLM_SESSION_TOKEN_PREFIX + from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_bearer_token + + expires: Final = datetime.now(timezone.utc) + timedelta(hours=24) + auth: Final = UserAPIKeyAuth( + token="liteadmin-" + secrets.token_urlsafe(24), + key_name="LiteAdmin Slack", + key_alias="LiteAdmin Slack", + user_id=user.user_id, + user_role=LitellmUserRoles.PROXY_ADMIN, + models=TypeAdapter(list[str]).validate_python(user.model_dump().get("models", [])), + expires=expires, + is_session_token=True, + ) + return AdminSession( + user_id=user.user_id, + credential=SecretStr( + encrypt_bearer_token(auth.model_dump_json(exclude_none=True), LITELLM_SESSION_TOKEN_PREFIX) + ), + expires_at=expires.timestamp(), + ) + + +def validate_native_configuration( + worker_url: str, service_token: str, enterprise: bool, database_available: bool +) -> None: + if not worker_url: + raise HTTPException(404, "LiteAdmin Slack is not enabled") + if not enterprise: + raise HTTPException(403, "LiteAdmin Slack requires LiteLLM Enterprise") + if not database_available: + raise HTTPException(503, "LiteAdmin requires a database") + try: + parsed: Final = urlsplit(worker_url) + port: Final = parsed.port + except ValueError: + raise HTTPException(503, "LiteAdmin worker configuration is invalid") from None + if ( + parsed.scheme not in {"http", "https"} + or not parsed.hostname + or port == 0 + or parsed.username + or parsed.password + or parsed.path + or parsed.query + or parsed.fragment + or len(service_token) < 32 + or any(character.isspace() for character in service_token) + ): + raise HTTPException(503, "LiteAdmin worker configuration is invalid") + + +async def native_admin_context() -> NativeAdminContext: + from litellm.proxy.proxy_server import premium_user, prisma_client + + worker_url: Final = os.getenv("LITELLM_ADMIN_AGENT_URL", "").rstrip("/") + service_token: Final = os.getenv("ADMIN_AGENT_SERVICE_TOKEN", "") + validate_native_configuration(worker_url, service_token, premium_user is True, prisma_client is not None) + client: Final = get_async_httpx_client( + llm_provider="liteadmin_native", params={"timeout": 15.0, "follow_redirects": False} + ).client + return NativeAdminContext( + worker_url, SecretStr(service_token), client, _session_user, _load_user, mint_admin_session + ) + + +@router.get(_PREFIX + "{token}", include_in_schema=False, response_class=HTMLResponse) +async def connect_page( + request: Request, + token: str, + context: Annotated[NativeAdminContext, Depends(native_admin_context)], +) -> Response: + details: Final = await context.details(token) + base_url: Final = get_request_base_url(request) + parsed_base: Final = urlsplit(base_url) + if parsed_base.scheme != "https": + raise HTTPException(400, "LiteAdmin account connections require HTTPS") + user_id: Final = await context.session_user(request) + if user_id is None: + return RedirectResponse( + base_url + "/sso/key/generate?" + urlencode({"return_to": parsed_base.path + _PREFIX + token}), + status_code=303, + headers=_HEADERS, + ) + await context.admin(user_id, details) + csrf: Final = secrets.token_urlsafe(32) + page: Final = _page( + "Connect LiteAdmin to Slack", + f"

Connect {html.escape(details.email)} to LiteAdmin in your Slack workspace?

" + "

Model requests and administrative actions will use your own LiteLLM account and current permissions

" + f'
' + '
' + "

This connection lasts 24 hours. Send disconnect in Slack to remove the saved session

", + ) + page.set_cookie(_cookie_name(token), csrf, max_age=600, secure=True, httponly=True, samesite="strict", path="/") + return page + + +@router.post(_PREFIX + "{token}", include_in_schema=False, response_class=HTMLResponse) +async def connect_account( + request: Request, + token: str, + context: Annotated[NativeAdminContext, Depends(native_admin_context)], +) -> Response: + base_url: Final = get_request_base_url(request) + parsed_base: Final = urlsplit(base_url) + origin: Final = f"{parsed_base.scheme}://{parsed_base.netloc}" + if parsed_base.scheme != "https" or request.headers.get("Origin") != origin: + raise HTTPException(403, "Reopen your private Slack connection link") + if request.headers.get("Content-Type", "").split(";", 1)[0] != "application/x-www-form-urlencoded": + raise HTTPException(400, "Expected a connection form") + form: Final = await request.form(max_fields=1, max_files=0, max_part_size=1024) + supplied: Final = form.get("csrf") + expected: Final = request.cookies.get(_cookie_name(token), "") + if ( + not isinstance(supplied, str) + or len(expected) != 43 + or len(supplied) != 43 + or not hmac.compare_digest(supplied.encode(), expected.encode()) + ): + raise HTTPException(403, "Reopen your private Slack connection link") + user_id: Final = await context.session_user(request) + if user_id is None: + raise HTTPException(401, "Your login expired. Reopen your private Slack connection link") + details: Final = await context.details(token) + user: Final = await context.admin(user_id, details) + await context.worker_request(token, context.mint_session(user)) + page: Final = _page( + "Account connected", "

Return to Slack and ask LiteAdmin to list your teams or check a budget

" + ) + page.delete_cookie(_cookie_name(token), path="/", secure=True, httponly=True, samesite="strict") + return page diff --git a/helm/litellm-helm/templates/deployment.yaml b/helm/litellm-helm/templates/deployment.yaml index cf7b3f8a38d..299d41e2019 100644 --- a/helm/litellm-helm/templates/deployment.yaml +++ b/helm/litellm-helm/templates/deployment.yaml @@ -57,6 +57,19 @@ spec: imagePullPolicy: {{ .Values.image.pullPolicy }} env: {{- include "litellm.proxyEnv" . | nindent 12 }} + {{- if .Values.liteadmin.enabled }} + - name: LITELLM_ADMIN_AGENT_URL + value: {{ printf "http://%s-liteadmin:10000" (include "litellm.fullname" . | trunc 53 | trimSuffix "-") | quote }} + - name: ADMIN_AGENT_SERVICE_TOKEN + valueFrom: + secretKeyRef: + name: {{ required "liteadmin.existingSecret is required" .Values.liteadmin.existingSecret }} + key: ADMIN_AGENT_SERVICE_TOKEN + {{- if not (hasKey (default dict .Values.envVars) "PROXY_BASE_URL") }} + - name: PROXY_BASE_URL + value: {{ required "liteadmin.gatewayUrl is required" .Values.liteadmin.gatewayUrl | quote }} + {{- end }} + {{- end }} {{- include "litellm.proxyMetricsEnv" . | nindent 12 }} {{- if .Values.collector.enabled }} {{- include "litellm.collectorEnv" . | nindent 12 }} diff --git a/helm/litellm-helm/templates/liteadmin.yaml b/helm/litellm-helm/templates/liteadmin.yaml new file mode 100644 index 00000000000..711edaf6913 --- /dev/null +++ b/helm/litellm-helm/templates/liteadmin.yaml @@ -0,0 +1,112 @@ +{{- if .Values.liteadmin.enabled }} +{{- $name := printf "%s-liteadmin" (include "litellm.fullname" . | trunc 53 | trimSuffix "-") }} +{{- $secret := required "liteadmin.existingSecret is required" .Values.liteadmin.existingSecret }} +apiVersion: apps/v1 +kind: Deployment +metadata: + name: {{ $name }} +spec: + replicas: 1 + strategy: + type: Recreate + selector: + matchLabels: + app.kubernetes.io/name: {{ $name }} + app.kubernetes.io/instance: {{ .Release.Name }} + template: + metadata: + labels: + app.kubernetes.io/name: {{ $name }} + app.kubernetes.io/instance: {{ .Release.Name }} + spec: + automountServiceAccountToken: false + terminationGracePeriodSeconds: 75 + {{- with .Values.imagePullSecrets }} + imagePullSecrets: + {{- toYaml . | nindent 8 }} + {{- end }} + securityContext: + runAsUser: 10001 + runAsGroup: 10001 + fsGroup: 10001 + runAsNonRoot: true + containers: + - name: liteadmin + image: "{{ .Values.image.repository }}:{{ .Values.image.tag | default .Chart.AppVersion }}" + imagePullPolicy: {{ .Values.image.pullPolicy }} + args: ["--admin-agent"] + securityContext: + allowPrivilegeEscalation: false + readOnlyRootFilesystem: true + capabilities: + drop: [ALL] + envFrom: + - secretRef: + name: {{ $secret }} + env: + - name: CONNECTION_AUTH_MODE + value: native + - name: LITELLM_BASE_URL + value: {{ required "liteadmin.gatewayUrl is required" .Values.liteadmin.gatewayUrl | quote }} + - name: LITELLM_MODEL + value: {{ required "liteadmin.model is required" .Values.liteadmin.model | quote }} + - name: STATE_DB + value: /var/data/events.sqlite3 + - name: ADMIN_READ_ONLY + value: {{ .Values.liteadmin.readOnly | quote }} + - name: OPENAI_AGENTS_DISABLE_TRACING + value: "1" + ports: + - name: health + containerPort: 10000 + readinessProbe: + httpGet: + path: /readyz + port: health + periodSeconds: 15 + livenessProbe: + httpGet: + path: /healthz + port: health + periodSeconds: 30 + resources: + {{- toYaml .Values.liteadmin.resources | nindent 12 }} + volumeMounts: + - name: state + mountPath: /var/data + - name: tmp + mountPath: /tmp + volumes: + - name: state + persistentVolumeClaim: + claimName: {{ $name }} + - name: tmp + emptyDir: + sizeLimit: 64Mi +--- +apiVersion: v1 +kind: Service +metadata: + name: {{ $name }} +spec: + type: ClusterIP + selector: + app.kubernetes.io/name: {{ $name }} + app.kubernetes.io/instance: {{ .Release.Name }} + ports: + - port: 10000 + targetPort: health +--- +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: {{ $name }} +spec: + accessModes: [ReadWriteOnce] + {{- with .Values.liteadmin.storageClassName }} + storageClassName: {{ . | quote }} + {{- end }} + resources: + requests: + storage: {{ .Values.liteadmin.storageSize }} +{{- end }} diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index 03d2a66a2b5..83dbb3c5aa0 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -3,6 +3,20 @@ # Declare variables to be passed into your templates. replicaCount: 1 +liteadmin: + enabled: false + existingSecret: "" + gatewayUrl: "" + model: "" + readOnly: false + storageSize: 1Gi + storageClassName: "" + resources: + requests: + cpu: 100m + memory: 256Mi + limits: + memory: 1Gi # numWorkers: 2 image: diff --git a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py index b51da626a60..5db9a92e51d 100644 --- a/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/byok_oauth_endpoints.py @@ -171,6 +171,42 @@ async def _session_key_is_live(session_key: str | None) -> bool: return True +async def get_authenticated_browser_user_id(request: Request) -> str | None: + from datetime import datetime, timezone + + from pydantic import TypeAdapter, ValidationError + + from litellm.proxy._types import hash_token + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_key_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + user_id, session_key = _session_identity_from_cookie(request) + if not user_id or not session_key or prisma_client is None: + return None + try: + auth: Final = ( + await get_key_object( + hash_token(session_key), + prisma_client, + user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + check_db_only=True, + ) + if session_key.startswith("sk-") + else ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(session_key) + ) + except Exception: + return None + if auth is None or auth.user_id != user_id or auth.blocked or auth.expires is None: + return None + try: + expiration: Final = TypeAdapter(datetime).validate_python(auth.expires) + except ValidationError: + return None + expires: Final = expiration.replace(tzinfo=timezone.utc) if expiration.tzinfo is None else expiration + return user_id if expires > datetime.now(timezone.utc) else None + + async def _byok_session_auth(request: Request) -> UserAPIKeyAuth: """Require the UI session cookie, with the embedded session key re-resolved against the DB so a revoked (logged-out) session cannot diff --git a/tests/unit/enterprise/proxy/test_liteadmin.py b/tests/unit/enterprise/proxy/test_liteadmin.py new file mode 100644 index 00000000000..4d1c3ce7f0a --- /dev/null +++ b/tests/unit/enterprise/proxy/test_liteadmin.py @@ -0,0 +1,244 @@ +from __future__ import annotations + +import json +import re +from typing import Final + +import httpx +import pytest +from fastapi import FastAPI, Request +from fastapi.testclient import TestClient +from litellm_enterprise.proxy.liteadmin import AdminSession, NativeAdminContext, native_admin_context, router +from pydantic import SecretStr, TypeAdapter + +from litellm.proxy._types import LiteLLM_UserTable + +TOKEN: Final = "a" * 43 +PATH: Final = "/liteadmin/slack/connect/" + TOKEN +ORIGIN: Final = "https://gateway.example.com" + + +class Worker: + def __init__(self, *, email: str = "alice@example.com", status: int = 200) -> None: + self.email = email + self.status = status + self.session: object = None + self.role = "proxy_admin" + + def request(self, request: httpx.Request) -> httpx.Response: + assert request.headers["X-LiteLLM-Admin-Agent-Token"] == "s" * 32 + if request.method == "POST": + self.session = json.loads(request.content) + return httpx.Response(self.status, json={"status": "connected"}) + return httpx.Response( + self.status, + json={ + "workspace_id": "Tworkspace", + "slack_user_id": "Ualice", + "email": self.email, + }, + ) + + +def client_for( + worker: Worker, *, role: str | None = None, logged_in: bool = True, email: str = "alice@example.com" +) -> TestClient: + async def session_user(request: Request) -> str | None: + return "alice" if logged_in else None + + async def load_user(user_id: str) -> LiteLLM_UserTable: + return LiteLLM_UserTable(user_id=user_id, user_email=email, user_role=role or worker.role) + + def mint(user: LiteLLM_UserTable) -> AdminSession: + return AdminSession(user_id=user.user_id, credential=SecretStr("personal-session"), expires_at=86400) + + context: Final = NativeAdminContext( + "http://private-worker:10000", + SecretStr("s" * 32), + httpx.AsyncClient(transport=httpx.MockTransport(worker.request)), + session_user, + load_user, + mint, + ) + app: Final = FastAPI() + app.include_router(router) + app.dependency_overrides[native_admin_context] = lambda: context + return TestClient(app, base_url=ORIGIN) + + +def csrf_from(client: TestClient) -> str: + page: Final = client.get(PATH) + assert page.status_code == 200 + match: Final = re.search('name="csrf" value="([^"]+)"', page.text) + assert match is not None + return match[1] + + +def test_connect_uses_existing_login_without_a_hosted_oauth_callback(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PROXY_BASE_URL", ORIGIN) + with client_for(Worker(), logged_in=False) as client: + response: Final = client.get(PATH, follow_redirects=False) + assert response.status_code == 303 + assert ( + response.headers["location"] == ORIGIN + "/sso/key/generate?return_to=%2Fliteadmin%2Fslack%2Fconnect%2F" + TOKEN + ) + assert response.headers["cache-control"] == "no-store" + + +def test_connect_hands_off_personal_session_only_over_private_worker_channel(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PROXY_BASE_URL", ORIGIN) + worker: Final = Worker() + with client_for(worker) as client: + csrf: Final = csrf_from(client) + response: Final = client.post(PATH, data={"csrf": csrf}, headers={"Origin": ORIGIN}) + assert response.status_code == 200 + assert "Account connected" in response.text + assert worker.session == {"user_id": "alice", "credential": "personal-session", "expires_at": 86400.0} + assert "personal-session" not in response.text + assert "Max-Age=0" in response.headers["set-cookie"] + + +@pytest.mark.parametrize("role,email", [("internal_user", "alice@example.com"), ("proxy_admin", "bob@example.com")]) +def test_connect_rejects_nonadmin_and_another_slack_users_link( + monkeypatch: pytest.MonkeyPatch, + role: str, + email: str, +) -> None: + monkeypatch.setenv("PROXY_BASE_URL", ORIGIN) + worker: Final = Worker(email=email) + with client_for(worker, role=role) as client: + response: Final = client.get(PATH) + assert response.status_code == 403 + assert worker.session is None + + +@pytest.mark.parametrize("origin,csrf", [("https://attacker.example", None), ("null", None), (ORIGIN, "b" * 43)]) +def test_connect_requires_same_origin_and_browser_csrf( + monkeypatch: pytest.MonkeyPatch, + origin: str, + csrf: str | None, +) -> None: + monkeypatch.setenv("PROXY_BASE_URL", ORIGIN) + worker: Final = Worker() + with client_for(worker) as client: + valid: Final = csrf_from(client) + response: Final = client.post(PATH, data={"csrf": csrf or valid}, headers={"Origin": origin}) + assert response.status_code == 403 + assert worker.session is None + + +@pytest.mark.parametrize("status,expected", [(410, 410), (403, 403), (500, 503), (302, 503)]) +def test_worker_denial_expiry_and_failure_never_mint_a_session( + monkeypatch: pytest.MonkeyPatch, + status: int, + expected: int, +) -> None: + monkeypatch.setenv("PROXY_BASE_URL", ORIGIN) + worker: Final = Worker(status=status) + with client_for(worker) as client: + response: Final = client.get(PATH) + assert response.status_code == expected + assert worker.session is None + + +def test_csrf_cookie_cannot_cross_links(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PROXY_BASE_URL", ORIGIN) + worker: Final = Worker() + with client_for(worker) as client: + csrf: Final = csrf_from(client) + response: Final = client.post(PATH.replace(TOKEN, "b" * 43), data={"csrf": csrf}, headers={"Origin": ORIGIN}) + assert response.status_code == 403 + assert worker.session is None + + +@pytest.mark.parametrize( + "url,secret,enterprise,database,status", + [ + ("", "s" * 32, True, True, 404), + ("http://worker:10000", "s" * 32, False, True, 403), + ("http://worker:10000", "s" * 32, True, False, 503), + ("file:///etc/passwd", "s" * 32, True, True, 503), + ("https://user:password@worker", "s" * 32, True, True, 503), + ("https://worker/path", "s" * 32, True, True, 503), + ("https://worker", "short", True, True, 503), + ("http://[broken", "s" * 32, True, True, 503), + ("http://worker:broken", "s" * 32, True, True, 503), + ], +) +def test_native_configuration_requires_enterprise_database_and_private_worker_credentials( + url: str, + secret: str, + enterprise: bool, + database: bool, + status: int, +) -> None: + from fastapi import HTTPException + from litellm_enterprise.proxy.liteadmin import validate_native_configuration + + with pytest.raises(HTTPException) as error: + validate_native_configuration(url, secret, enterprise, database) + assert error.value.status_code == status + + +def test_connect_rechecks_admin_permission_after_consent_page(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PROXY_BASE_URL", ORIGIN) + worker: Final = Worker() + with client_for(worker) as client: + csrf: Final = csrf_from(client) + worker.role = "internal_user" + response: Final = client.post(PATH, data={"csrf": csrf}, headers={"Origin": ORIGIN}) + assert response.status_code == 403 + assert worker.session is None + + +def test_consent_escapes_slack_email(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("PROXY_BASE_URL", ORIGIN) + email: Final = '@example.com' + with client_for(Worker(email=email), email=email) as client: + page: Final = client.get(PATH) + assert page.status_code == 200 + assert "