diff --git a/.github/PULL_REQUEST_TEMPLATE/general.md b/.github/PULL_REQUEST_TEMPLATE/general.md index 4f1a18e925e..d766897ab3c 100644 --- a/.github/PULL_REQUEST_TEMPLATE/general.md +++ b/.github/PULL_REQUEST_TEMPLATE/general.md @@ -6,19 +6,12 @@ ## Handwritten 1-3 sentence description - -`` - -`` +[insert handwritten description here]
Details diff --git a/cookbook/skill_scanner/README.md b/cookbook/skill_scanner/README.md new file mode 100644 index 00000000000..2a68e116002 --- /dev/null +++ b/cookbook/skill_scanner/README.md @@ -0,0 +1,60 @@ +# Manual skill security scans + +This standalone adapter runs NVIDIA SkillSpector on a separate machine. It is not installed in the LiteLLM proxy and does not require a shared Helm or Compose deployment + +Start it with Python 3.12 or newer and uv. The script pins SkillSpector and its HTTP dependencies + +```sh +export SKILL_SCANNER_API_KEY="$(openssl rand -hex 32)" +uv run server.py +``` + +For another host, place the loopback service behind an HTTPS reverse proxy that forwards the Authorization header and enforces a 65 MiB request limit. Run it under an unprivileged account with a writable temporary directory and no production files or credentials. The scanner host needs Git and network access to requested sources and optional model providers. Apply network egress restrictions suitable for those sources + +Configure the same scanner secret in the gateway's environment, then add this to LiteLLM's configuration + +```yaml +general_settings: + skill_scanner: + endpoint: https://scanner.example.com/scan + api_key_env: SKILL_SCANNER_API_KEY + use_llm: false + timeout_seconds: 180 +``` + +The endpoint must be reachable from the gateway. A customer's laptop works while it is reachable through their network, VPN or authenticated HTTPS tunnel. `localhost` refers to the gateway machine, not the admin's browser. HTTP is accepted only for loopback development + +Open Skills, Submissions, and Review as a proxy admin. Click **Scan for threats**, inspect the report, then independently approve or reject the proposal. Scanning is optional and never changes approval permissions or blocks publication. The report shows findings, locations, remediation, risk score, scanner version and whether semantic analysis actually ran. Incomplete reports do not establish safety + +Uploaded ZIPs are sent privately to the scanner with their exact stored bytes. Metadata-only edits reuse the matching stored archive. If that archive is missing or changed, submit a new ZIP before scanning. The displayed SHA-256 identifies that archive, and the scan examines its files rather than proposed catalog metadata. Remote URL reports describe the source at scan time and do not freeze its contents. GitHub branch/tag refs and explicit-ref subdirectories are supported. Other ref/subdirectory combinations require a ZIP so the adapter does not silently scan the wrong revision + +The latest scan result is retained with the submission. Running scans have a deadline. If a gateway restarts or the scanner becomes unavailable, the attempt expires and an admin can retry. Retries replace the latest report. Approving or rejecting during a scan does not cancel it or change the decision + +Set `use_llm: true` to request semantic analysis and configure the provider on the scanner machine using [SkillSpector's provider instructions](https://github.com/NVIDIA/skillspector#llm-analysis). Provider credentials stay on that machine. Static analysis can also call OSV for dependency checks + +## HTTP contract + +The gateway sends an authenticated `POST` to the configured endpoint. The body contains form fields `source` (JSON source descriptor), `use_llm` (boolean), and `expires_at` (timezone-aware ISO timestamp), plus `file` (ZIP) when bytes are available. The service returns HTTP 200 with this JSON shape + +```json +{ + "scanner_version": "2.12.0", + "complete": true, + "scan_mode": "static-only", + "risk_score": 70, + "findings": [{ + "rule_id": "example-rule", + "severity": "HIGH", + "message": "Potential data exfiltration", + "file": "SKILL.md", + "start_line": 8, + "explanation": "Instructions request sending local files to an external host", + "remediation": "Remove the external upload instruction" + }], + "limitations": [] +} +``` + +Reports are limited to 1 MiB and 500 findings. Scanner uploads allow 64 MiB because the gateway can expand compressed skill uploads into stored compatibility ZIPs. The gateway rejects malformed reports, refuses redirects and verifies TLS. Only the configured scanner credential is sent. A timeout, HTTP error or incomplete report remains visible to the reviewer + +The bundled adapter runs at most two child scans concurrently. Each child and its descendants are terminated at the gateway deadline, capped at 570 seconds. Keep the gateway and scanner clocks synchronized. Direct adapter clients may omit `expires_at` to use the cap. Child processes never execute submitted skill scripts. Raise the gateway deadline up to 600 seconds when needed for semantic analysis and configure the reverse proxy timeout to match diff --git a/cookbook/skill_scanner/server.py b/cookbook/skill_scanner/server.py new file mode 100644 index 00000000000..44a1cecb00f --- /dev/null +++ b/cookbook/skill_scanner/server.py @@ -0,0 +1,172 @@ +# /// script +# requires-python = ">=3.12" +# dependencies = [ +# "fastapi==0.162.0", "uvicorn==0.54.0", "python-multipart==0.0.33", +# "skillspector @ git+https://github.com/NVIDIA/skillspector.git@62c2fe8622f580ac292d4078c3b71f790194dbbd", +# ] +# /// +import asyncio +import hmac +import json +import os +import signal +import sys +import tempfile +from contextlib import suppress +from datetime import datetime, timezone +from pathlib import Path +from typing import Annotated, Final +from urllib.parse import quote, urlsplit + +from fastapi import Depends, FastAPI, File, Form, HTTPException, UploadFile +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer +from pydantic import AwareDatetime, TypeAdapter, ValidationError + +MAX_UPLOAD: Final = 64 * 1024 * 1024 +MAX_REPORT: Final = 1024 * 1024 +SCAN_TIMEOUT: Final = 570 +SOURCE: Final = TypeAdapter(dict[str, str]) +SLOTS: Final = asyncio.Semaphore(2) +app: Final = FastAPI(title="SkillSpector adapter", docs_url=None, redoc_url=None, openapi_url=None) + + +def authenticate(credentials: Annotated[HTTPAuthorizationCredentials, Depends(HTTPBearer())]) -> None: + token: Final = os.environ.get("SKILL_SCANNER_API_KEY", "") + if not token or not hmac.compare_digest(credentials.credentials, token): + raise HTTPException(401, "Invalid scanner credentials") + + +def source_target(source: dict[str, str]) -> str: + target: Final = ( + "https://github.com/" + source.get("repo", "") if source.get("source") == "github" else source.get("url", "") + ) + parsed: Final = urlsplit(target) + if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password or parsed.fragment: + raise HTTPException(400, "Remote sources require an HTTPS URL without credentials or a fragment") + ref: Final = source.get("ref") + subdir: Final = source.get("path") + if ref or subdir: + if parsed.hostname != "github.com" or len(parsed.path.strip("/").split("/")) != 2 or not ref: + raise HTTPException(400, "Upload a ZIP to scan this source's exact ref and subdirectory") + suffix: Final = "/" + quote(subdir, safe="/") if subdir else "" + return target.removesuffix(".git").rstrip("/") + "/tree/" + quote(ref, safe="/") + suffix + return target + + +async def worker(target: str, output: Path, use_llm: bool) -> None: + from skillspector.mcp_server import run_scan + + result: Final = await run_scan(target, use_llm=use_llm, allow_local_targets=not target.startswith("https://")) + coverage: Final = result.get("analysis_completeness", {}) + findings: Final = result["findings"] + complete: Final = ( + result.get("execution_successful") is True + and coverage.get("is_complete") is True + and (not use_llm or (result.get("llm_available") is True and result.get("llm_used") is True)) + and len(findings) <= 500 + ) + limitations: Final = tuple(str(value)[:2000] for value in coverage.get("limitations", ())) + report: Final = { + "scanner_version": result["version"], + "complete": complete, + "scan_mode": result["scan_mode"], + "risk_score": result["risk_score"], + "findings": [ + { + "rule_id": finding["id"], + "severity": finding["severity"], + "message": finding.get("pattern") or finding.get("finding") or "Security finding", + "file": finding["location"]["file"], + "start_line": finding["location"].get("start_line"), + "explanation": finding.get("explanation"), + "remediation": finding.get("remediation"), + } + for finding in findings[:500] + ], + "limitations": ( + *limitations, + *(("Only the first 500 findings are displayed.",) if len(findings) > 500 else ()), + ), + } + encoded: Final = json.dumps(report).encode() + if len(encoded) > MAX_REPORT: + raise ValueError("Scanner report exceeded the response limit") + output.write_bytes(encoded) + + +async def run_worker(target: str, output: Path, use_llm: bool, expires_at: datetime | None = None) -> dict[str, object]: + timeout: Final = ( + SCAN_TIMEOUT + if expires_at is None + else min(SCAN_TIMEOUT, (expires_at - datetime.now(timezone.utc)).total_seconds()) + ) + if timeout <= 0: + raise HTTPException(504, "SkillSpector scan deadline expired") + process: Final = await asyncio.create_subprocess_exec( + sys.executable, + str(Path(__file__).resolve()), + "--scan", + target, + str(output), + str(use_llm).lower(), + stdout=asyncio.subprocess.DEVNULL, + stderr=asyncio.subprocess.DEVNULL, + start_new_session=True, + env={ + **os.environ, + "HOME": str(output.parent), + "TMPDIR": str(output.parent), + "GIT_CONFIG_GLOBAL": os.devnull, + "GIT_CONFIG_NOSYSTEM": "1", + "GIT_TERMINAL_PROMPT": "0", + "GIT_ALLOW_PROTOCOL": "https", + }, + ) + try: + code: Final = await asyncio.wait_for(process.wait(), timeout) + if code != 0 or not output.is_file() or output.stat().st_size > MAX_REPORT: + raise HTTPException(502, "SkillSpector could not produce a report for this source") + return TypeAdapter(dict[str, object]).validate_json(output.read_bytes()) + except TimeoutError: + raise HTTPException(504, "SkillSpector scan timed out") from None + finally: + with suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() + + +@app.post("/scan", dependencies=[Depends(authenticate)]) +async def scan( + source: Annotated[str, Form(max_length=16384)], + use_llm: Annotated[bool, Form()] = False, + expires_at: Annotated[AwareDatetime | None, Form()] = None, + file: Annotated[UploadFile | None, File()] = None, +) -> dict[str, object]: + if SLOTS.locked(): + raise HTTPException(429, "Scanner is busy; try again shortly") + try: + parsed: Final = SOURCE.validate_json(source) + except ValidationError: + raise HTTPException(400, "Invalid source metadata") from None + async with SLOTS: + with tempfile.TemporaryDirectory(prefix="skill-scan-") as directory: + root: Final = Path(directory) + if file is not None: + content: Final = await file.read(MAX_UPLOAD + 1) + if len(content) > MAX_UPLOAD: + raise HTTPException(413, "Upload exceeds 64 MiB") + archive: Final = root / "skill.zip" + archive.write_bytes(content) + return await run_worker(str(archive), root / "report.json", use_llm, expires_at) + return await run_worker(source_target(parsed), root / "report.json", use_llm, expires_at) + + +if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "--scan": + asyncio.run(worker(sys.argv[2], Path(sys.argv[3]), sys.argv[4] == "true")) + else: + import uvicorn + + if not os.environ.get("SKILL_SCANNER_API_KEY"): + raise SystemExit("Set SKILL_SCANNER_API_KEY before starting the scanner") + uvicorn.run(app, host="127.0.0.1", port=int(os.environ.get("PORT", "8099"))) diff --git a/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py b/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py index a2322d7f650..bb3b712d50b 100644 --- a/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py +++ b/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py @@ -1,9 +1,9 @@ import enum -from typing import Dict, List, Optional from pydantic import Field from litellm.proxy._types import WebhookEvent +from litellm.types.integrations.model_deprecation_email import MODEL_DEPRECATION_EMAIL_EVENT from litellm.types.llms.base import LiteLLMBaseModel @@ -26,7 +26,7 @@ class SendKeyCreatedEmailEvent(WebhookEvent): class SendKeyRotatedEmailEvent(WebhookEvent): virtual_key: str - key_alias: Optional[str] = None + key_alias: str | None = None """ The virtual key that was rotated this will be sk-123xxx, since we will be emailing this to the user to start using the new key @@ -34,23 +34,33 @@ class SendKeyRotatedEmailEvent(WebhookEvent): class EmailEvent(str, enum.Enum): + model_deprecations = MODEL_DEPRECATION_EMAIL_EVENT virtual_key_created = "Virtual Key Created" new_user_invitation = "New User Invitation" virtual_key_rotated = "Virtual Key Rotated" soft_budget_crossed = "Soft Budget Crossed" max_budget_alert = "Max Budget Alert" + class EmailEventSettings(LiteLLMBaseModel): event: EmailEvent enabled: bool + + class EmailEventSettingsUpdateRequest(LiteLLMBaseModel): - settings: List[EmailEventSettings] + settings: list[EmailEventSettings] + + class EmailEventSettingsResponse(LiteLLMBaseModel): - settings: List[EmailEventSettings] + settings: list[EmailEventSettings] + + class DefaultEmailSettings(LiteLLMBaseModel): """Default settings for email events""" - settings: Dict[EmailEvent, bool] = Field( + + settings: dict[EmailEvent, bool] = Field( default_factory=lambda: { + EmailEvent.model_deprecations: False, # Explicit opt-in for owner retirement reminders. EmailEvent.virtual_key_created: True, # On by default EmailEvent.new_user_invitation: True, # On by default EmailEvent.virtual_key_rotated: True, # On by default @@ -58,10 +68,12 @@ class DefaultEmailSettings(LiteLLMBaseModel): EmailEvent.max_budget_alert: True, # On by default } ) - def to_dict(self) -> Dict[str, bool]: + + def to_dict(self) -> dict[str, bool]: """Convert to dictionary with string keys for storage""" return {event.value: enabled for event, enabled in self.settings.items()} + @classmethod - def get_defaults(cls) -> Dict[str, bool]: + def get_defaults(cls) -> dict[str, bool]: """Get the default settings as a dictionary with string keys""" - return cls().to_dict() \ No newline at end of file + return cls().to_dict() diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_skill_version_history/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_skill_version_history/migration.sql new file mode 100644 index 00000000000..ecfa7321d87 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_skill_version_history/migration.sql @@ -0,0 +1,16 @@ +SET LOCAL lock_timeout = '10s'; +CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion" ( + "id" TEXT NOT NULL, + "plugin_id" TEXT NOT NULL, + "version" TEXT, + "description" TEXT, + "manifest_json" TEXT, + "published_at" TIMESTAMP(3), + "published_by" TEXT, + "archive_available" BOOLEAN NOT NULL DEFAULT false, + "content" BYTEA, + CONSTRAINT "LiteLLM_ClaudeCodePluginVersion_pkey" PRIMARY KEY ("id"), + CONSTRAINT "LiteLLM_ClaudeCodePluginVersion_plugin_id_fkey" FOREIGN KEY ("plugin_id") REFERENCES "LiteLLM_ClaudeCodePluginTable"("id") ON DELETE CASCADE ON UPDATE CASCADE +); +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion_plugin_id_version_key" ON "LiteLLM_ClaudeCodePluginVersion"("plugin_id", "version"); +CREATE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginVersion_plugin_id_published_at_idx" ON "LiteLLM_ClaudeCodePluginVersion"("plugin_id", "published_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105010000_skill_submissions/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105010000_skill_submissions/migration.sql new file mode 100644 index 00000000000..e7975258ca9 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105010000_skill_submissions/migration.sql @@ -0,0 +1,27 @@ +SET LOCAL lock_timeout = '10s'; + +CREATE TABLE IF NOT EXISTS "LiteLLM_SkillSubmission" ( + "id" TEXT NOT NULL, + "plugin_name" TEXT NOT NULL, + "operation" TEXT NOT NULL, + "status" TEXT NOT NULL DEFAULT 'pending_review', + "proposal_json" TEXT NOT NULL, + "archive_content" BYTEA, + "submitted_by" TEXT NOT NULL, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "base_plugin_id" TEXT, + "base_updated_at" TIMESTAMP(3), + "reviewed_by" TEXT, + "reviewed_at" TIMESTAMP(3), + "rejection_reason" TEXT, + CONSTRAINT "LiteLLM_SkillSubmission_pkey" PRIMARY KEY ("id") +); + +CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_submitted_by_created_at_id_idx" + ON "LiteLLM_SkillSubmission"("submitted_by", "created_at" DESC, "id" DESC); +CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_submitted_by_status_created_at_id_idx" + ON "LiteLLM_SkillSubmission"("submitted_by", "status", "created_at" DESC, "id" DESC); +CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_status_created_at_id_idx" + ON "LiteLLM_SkillSubmission"("status", "created_at" DESC, "id" DESC); +CREATE INDEX IF NOT EXISTS "LiteLLM_SkillSubmission_created_at_id_idx" + ON "LiteLLM_SkillSubmission"("created_at" DESC, "id" DESC); diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105020000_skill_scanning/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105020000_skill_scanning/migration.sql new file mode 100644 index 00000000000..6b70cca4a9a --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105020000_skill_scanning/migration.sql @@ -0,0 +1,3 @@ +SET LOCAL lock_timeout = '10s'; + +ALTER TABLE "LiteLLM_SkillSubmission" ADD COLUMN IF NOT EXISTS "scan_json" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105094500_add_skill_download_history/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105094500_add_skill_download_history/migration.sql new file mode 100644 index 00000000000..43e26b324a0 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261105094500_add_skill_download_history/migration.sql @@ -0,0 +1,37 @@ +SET LOCAL lock_timeout = '10s'; + +DO $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM pg_attribute + WHERE attrelid = to_regclass('"LiteLLM_ClaudeCodePluginTable"') + AND attname = 'anonymous_download_count' AND NOT attisdropped + ) THEN + ALTER TABLE "LiteLLM_ClaudeCodePluginTable" + ADD COLUMN IF NOT EXISTS "anonymous_download_count" BIGINT NOT NULL DEFAULT 0; + END IF; + + IF to_regclass('"LiteLLM_ClaudeCodePluginDownload"') IS NULL THEN + CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginDownload" ( + "id" BIGSERIAL NOT NULL, + "plugin_id" TEXT NOT NULL, + "downloaded_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "version" TEXT, + "digest" TEXT NOT NULL, + "authenticated" BOOLEAN NOT NULL, + "user_id" TEXT, + "user_email" TEXT, + "key_alias" TEXT, + "team_id" TEXT, + "team_alias" TEXT, + CONSTRAINT "LiteLLM_ClaudeCodePluginDownload_pkey" PRIMARY KEY ("id"), + CONSTRAINT "LiteLLM_ClaudeCodePluginDownload_plugin_id_fkey" FOREIGN KEY ("plugin_id") + REFERENCES "LiteLLM_ClaudeCodePluginTable"("id") ON DELETE CASCADE ON UPDATE CASCADE + ); + END IF; + + IF to_regclass('"LiteLLM_ClaudeCodePluginDownload_plugin_id_id_idx"') IS NULL THEN + CREATE INDEX IF NOT EXISTS "LiteLLM_ClaudeCodePluginDownload_plugin_id_id_idx" + ON "LiteLLM_ClaudeCodePluginDownload"("plugin_id", "id"); + END IF; +END $$; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 78f16d9f1fd..818f2dc4858 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1616,10 +1616,69 @@ model LiteLLM_ClaudeCodePluginTable { created_by String? uploaded_archive LiteLLM_ClaudeCodePluginUpload? + versions LiteLLM_ClaudeCodePluginVersion[] + downloads LiteLLM_ClaudeCodePluginDownload[] + anonymous_download_count BigInt @default(0) @@map("LiteLLM_ClaudeCodePluginTable") } +model LiteLLM_ClaudeCodePluginVersion { + id String @id @default(uuid()) + plugin_id String + version String? + description String? + manifest_json String? + published_at DateTime? + published_by String? + archive_available Boolean @default(false) + content Bytes? + plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_id], references: [id], onDelete: Cascade, onUpdate: Cascade) + + @@unique([plugin_id, version]) + @@index([plugin_id, published_at]) +} + + +model LiteLLM_SkillSubmission { + id String @id @default(uuid()) + plugin_name String + operation String + status String @default("pending_review") + proposal_json String + archive_content Bytes? + submitted_by String + created_at DateTime @default(now()) + base_plugin_id String? + base_updated_at DateTime? + reviewed_by String? + reviewed_at DateTime? + rejection_reason String? + scan_json String? + + @@index([submitted_by, created_at(sort: Desc), id(sort: Desc)]) + @@index([submitted_by, status, created_at(sort: Desc), id(sort: Desc)]) + @@index([status, created_at(sort: Desc), id(sort: Desc)]) + @@index([created_at(sort: Desc), id(sort: Desc)]) +} + +model LiteLLM_ClaudeCodePluginDownload { + id BigInt @id @default(autoincrement()) + plugin_id String + downloaded_at DateTime @default(now()) + version String? + digest String + authenticated Boolean + user_id String? + user_email String? + key_alias String? + team_id String? + team_alias String? + plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_id], references: [id], onDelete: Cascade, onUpdate: Cascade) + + @@index([plugin_id, id]) +} + model LiteLLM_ClaudeCodePluginUpload { plugin_name String @id content Bytes diff --git a/litellm/accounting/application.py b/litellm/accounting/application.py index 1e5f23e2997..c911f4e35ce 100644 --- a/litellm/accounting/application.py +++ b/litellm/accounting/application.py @@ -7,6 +7,7 @@ from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, build_b from litellm.proxy.db.db_span import db_span from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import WindowSpendTransaction from litellm.proxy.db.model_usage_rollup import queue_model_usage_upsert +from litellm.proxy.db.rollup_lock_timeout import apply_rollup_lock_timeout from litellm.repositories.accounting_delivery_repository import DeliveryProjection, LedgerTransaction if TYPE_CHECKING: @@ -25,6 +26,7 @@ async def apply_deliveries( from litellm.proxy.guardrails.usage_tracking import UsageUnitIncrement, UsageUnitKey, upsert_usage_unit_row contributions: Final = project_deliveries(deliveries) + await apply_rollup_lock_timeout(transaction) await writer.apply_accounting_entities( prisma_client, proxy_logging_obj, transaction, entity_transactions(contributions) ) diff --git a/litellm/accounting/snapshot.py b/litellm/accounting/snapshot.py index d5652f79ed0..ab85df08320 100644 --- a/litellm/accounting/snapshot.py +++ b/litellm/accounting/snapshot.py @@ -14,6 +14,7 @@ if TYPE_CHECKING: from litellm.router import Router _FIELDS: Final = TypeAdapter(dict[str, object]) +_JSON_SAFE_INTEGER: Final = 2**53 - 1 def _mapping(value: object) -> Mapping[str, object] | None: @@ -86,6 +87,7 @@ def build_accounting_payload( "metadata": accounting_metadata, }, non_finite_as_null=True, + integer_limit=_JSON_SAFE_INTEGER, ) ) diff --git a/litellm/accounting/worker.py b/litellm/accounting/worker.py index df71f2da31c..be18d8ac6b1 100644 --- a/litellm/accounting/worker.py +++ b/litellm/accounting/worker.py @@ -14,6 +14,7 @@ from litellm.accounting.metrics import record_deliveries from litellm.accounting.projection import Contribution from litellm.repositories.accounting_delivery_repository import ( AccountingDeliveryRepository, + DeliveryBatchResult, DeliveryProjection, LedgerDatabase, LedgerTransaction, @@ -34,6 +35,14 @@ class _AccountingAlerts(Protocol): _JOB: Final = "accounting_delivery" +def _filled(result: DeliveryBatchResult, batch_size: int) -> bool: + return result.applied + len(result.parked) >= batch_size + + +def _progressed(result: DeliveryBatchResult) -> bool: + return bool(result.applied or result.failed) + + async def flush_accounting_deliveries( writer: "DBSpendUpdateWriter", prisma_client: "PrismaClient", proxy_logging_obj: "ProxyLogging" ) -> None: @@ -54,9 +63,9 @@ async def flush_accounting_deliveries( ) repository: Final = AccountingDeliveryRepository(database) started: Final = time.monotonic() - batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "10000")) + batch_size: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "500")) - async def run_batch(*, retry_only: bool) -> bool: + async def run_batch(*, retry_only: bool) -> DeliveryBatchResult: result: Final = await repository.apply_pending(apply, batch_size=batch_size, retry_only=retry_only) if result.applied: await writer.invalidate_accounting_entities(applied, proxy_logging_obj) @@ -68,17 +77,25 @@ async def flush_accounting_deliveries( proxy_logging_obj, message=f"Accounting delivery parked {len(result.parked)} requests. Included spend totals require repair.", ) - return bool(result.applied or result.failed) + return result - await run_batch(retry_only=False) + fresh_deadline: Final = time.monotonic() + 5 + for _ in range(batch_size): + if not _filled(await run_batch(retry_only=False), batch_size) or time.monotonic() >= fresh_deadline: + break retry_deadline: Final = time.monotonic() + 5 for _ in range(batch_size): - if time.monotonic() >= retry_deadline or not await run_batch(retry_only=True): + if time.monotonic() >= retry_deadline or not _progressed(await run_batch(retry_only=True)): + break + retention_days: Final = int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")) + cleanup_deadline: Final = time.monotonic() + 5 + for _ in range(batch_size): + if ( + await repository.cleanup_applied(retention_days=retention_days, batch_size=batch_size * 2) + < batch_size * 2 + or time.monotonic() >= cleanup_deadline + ): break - await repository.cleanup_applied( - retention_days=int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_RETENTION_DAYS", "1")), - batch_size=batch_size * 2, - ) metrics: Final = await repository.metrics() record_deliveries(metrics, time.monotonic() - started) if metrics.batch_count >= int( diff --git a/litellm/constants.py b/litellm/constants.py index 50e090eabaf..8d05115e76e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -28,7 +28,7 @@ GITHUB_COPILOT_PER_USER_AUTH_TYPE: Final = "per_user_oauth" GITHUB_COPILOT_AUTH_TYPE_KEY: Final = "github_copilot_auth_type" GITHUB_COPILOT_HOST_KEY: Final = "github_copilot_host" GITHUB_COPILOT_DEFAULT_HOST: Final = "github.com" -GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS: Final = 60 +GITHUB_COPILOT_USER_SESSION_TTL_SECONDS: Final = 1800 GITHUB_COPILOT_USER_CREDENTIAL_CACHE_TTL_SECONDS: Final = 60 GITHUB_COPILOT_DEVICE_FLOW_CACHE_PREFIX: Final = "github_copilot_device_flow" USER_PROVIDER_CREDENTIAL_CACHE_PREFIX: Final = "user_provider_credential" diff --git a/litellm/integrations/email_templates/model_deprecation_email.py b/litellm/integrations/email_templates/model_deprecation_email.py new file mode 100644 index 00000000000..5d325b2442a --- /dev/null +++ b/litellm/integrations/email_templates/model_deprecation_email.py @@ -0,0 +1,181 @@ +"""Model deprecation emails with provider identity, a migration comparison and next steps.""" + +from collections.abc import Mapping +from html import escape +from types import MappingProxyType +from typing import Final +from urllib.parse import urlencode, urlsplit + +from litellm.types.integrations.model_deprecation_email import ( + DeprecationEmailMilestone, + ModelDeprecationEmail, + RenderedDeprecationEmail, +) +from litellm.types.proxy.model_deprecation import SuccessorModelSource + +_LOGO_BASE_URL: Final = ( + "https://raw.githubusercontent.com/BerriAI/litellm/main/ui/litellm-dashboard/public/assets/logos" +) +_EMAIL_LOGO_URL: Final = f"{_LOGO_BASE_URL}/litellm_logo.png" +# Reuse dashboard brand assets; PNGs also render in email clients that block SVG. +_PROVIDER_BRANDS: Final[Mapping[str, tuple[str, str]]] = MappingProxyType( + { + "gemini": ("Google", "google_pse.png"), + "vertex_ai": ("Google", "google_pse.png"), + "openai": ("OpenAI", "openai.png"), + "anthropic": ("Anthropic", "anthropic.png"), + "azure": ("Microsoft Azure", "microsoft_azure.png"), + "bedrock": ("AWS Bedrock", "aws.png"), + "bedrock_converse": ("AWS Bedrock", "aws.png"), + } +) + + +def _recommendation_source(source: SuccessorModelSource | None) -> str: + if source is SuccessorModelSource.DEPLOYMENT_CONFIG: + return "Gateway configuration" + if source is SuccessorModelSource.MODEL_COST_MAP: + return "LiteLLM model cost map" + return "Not provided" + + +def _http_url(value: str | None) -> str | None: + if value is None: + return None + try: + parsed: Final = urlsplit(value) + except ValueError: + return None + return value if parsed.scheme in ("http", "https") and parsed.netloc and not parsed.username else None + + +def _provider(model: str, fallback: str | None = None) -> str: + prefix, separator, _ = model.partition("/") + return prefix if separator else fallback or "" + + +def _provider_name(provider: str) -> str: + brand: Final = _PROVIDER_BRANDS.get(provider) + if brand: + return brand[0] + return provider.replace("_", " ") or "The provider" + + +def _model_html(model: str, provider: str) -> str: + brand: Final = _PROVIDER_BRANDS.get(provider) + logo: Final = ( + f'{escape(brand[0], quote=True)}' + if brand + else "" + ) + return ( + f'{logo}' + f'
{escape(_provider_name(provider))}
' + f"
{escape(model)}
" + ) + + +def render_model_deprecation_email(email: ModelDeprecationEmail, *, base_url: str) -> RenderedDeprecationEmail: + gateway: Final = _http_url(base_url) + if gateway is None: + raise ValueError("Model deprecation emails require a public HTTP(S) PROXY_BASE_URL") + urgent: Final = email.milestone is DeprecationEmailMilestone.ONE_DAY + days: Final = int(email.milestone) + timing: Final = "tomorrow" if urgent else "in 30 days" + provider: Final = _provider(email.usage_model, email.model.litellm_provider) + provider_name: Final = _provider_name(provider) + subject: Final = f"[LiteLLM] Model being deprecated by {provider_name} {timing}: {email.usage_model}" + deadline: Final = email.model.deprecation_date.strftime("%d %b %Y") + announcement: Final = f"{provider_name} is deprecating this model on {deadline} (UTC)." + preheader: Final = f"{announcement} Review the model change and migrate your affected keys." + query: Final = urlencode( + { + "tab": "models", + "from": email.start_date.isoformat(), + "to": email.end_date.isoformat(), + } + ) + review_url: Final = f"{gateway.rstrip('/')}/ui/usage?{query}" + replacement: Final = email.model.successor_model + after: Final = ( + _model_html(replacement, _provider(replacement, provider)) + if replacement + else '

No recommended model listed. Choose a supported replacement.

' + ) + source: Final = _recommendation_source(email.model.successor_model_source) + notice: Final = _http_url(email.model.deprecation_source) + source_link: Final = ( + f'Read the provider notice ↗' + if notice + else "" + ) + team_names: Final = ", ".join(email.team_names[:3]) or "Personal keys" + more_teams: Final = f" + {len(email.team_names) - 3} more" if len(email.team_names) > 3 else "" + day_label: Final = "DAY" if urgent else "DAYS" + label: Final = "FINAL REMINDER" if urgent else "DEPRECATION NOTICE" + test_step: Final = ( + "Test the recommended model with your workload." + if replacement + else "Choose a supported replacement and test it with your workload." + ) + switch_step: Final = ( + f"Switch your app or gateway routing before {deadline}. Then check that requests to the old model reach zero." + ) + html: Final = f""" + +{escape(subject)} + + +
{escape(preheader)}
+
+ + + + + +
+ +
LiteLLM/ MODEL DEPRECATIONS
+ +
{label}{days:02d} {day_label} UNTIL DEPRECATION
+

Model being deprecated

+

{escape(announcement)}

+
+ + +
+
BEFORE · MODEL TO REPLACE
+{_model_html(email.usage_model, provider)} +		 +
AFTER · RECOMMENDED MODEL
+{after} +
+

Recommendation source: {escape(source)}
{source_link}

+

Your affected usage

+ + +
{email.key_count:,}
Affected keys
{email.requests:,}
Requests · last 30 days
+

{escape(team_names + more_teams)}
Your keys and teams you administer.

+

What to do

+
    +
  1. Review the affected keys in your usage dashboard.
    2. +
  2. {escape(test_step)}
    3. +
  3. {escape(switch_step)}
+
+Review affected usage ↗
+
+Your gateway admin enabled these emails. Reminder schedule: 30 days and 1 day before deprecation.

LiteLLM · AI Gateway
""" + text: Final = ( + f"{subject}\n\n{announcement}\n\n" + f"BEFORE - MODEL TO REPLACE\n{email.usage_model}\n\n" + f"AFTER - RECOMMENDED MODEL\n{replacement or 'No recommended model listed. Choose a supported replacement.'}\n" + f"Recommendation source: {source}\nProvider notice: {notice or 'Not provided'}\n\n" + f"YOUR AFFECTED USAGE\n{email.key_count} keys | {email.requests} requests in the last 30 days\n" + f"{team_names}{more_teams}\nYour keys and teams you administer.\n\n" + f"WHAT TO DO\n1. Review the affected keys in your usage dashboard.\n2. {test_step}\n3. {switch_step}\n\n" + f"Review affected usage: {review_url}\n\n" + "Your gateway admin enabled these emails. Reminder schedule: 30 days and 1 day before deprecation." + ) + return RenderedDeprecationEmail(subject=subject, html=html, text=text) diff --git a/litellm/integrations/model_deprecation_email.py b/litellm/integrations/model_deprecation_email.py new file mode 100644 index 00000000000..4093a108dcc --- /dev/null +++ b/litellm/integrations/model_deprecation_email.py @@ -0,0 +1,267 @@ +"""Two owner-scoped retirement reminders, with durable per-recipient delivery receipts.""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +import os +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import date, datetime, timedelta, timezone +from enum import Enum +from itertools import chain +from types import MappingProxyType +from typing import TYPE_CHECKING, Final + +from pydantic import AliasPath, BaseModel, ConfigDict, Field, TypeAdapter + +from litellm.integrations.email_templates.model_deprecation_email import ( + render_model_deprecation_email, +) +from litellm.models.team import LiteLLM_TeamTable +from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations +from litellm.repositories.config_repository import ConfigRepository +from litellm.repositories.table_repositories import DailyUserSpendRepository +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import ( + VerificationTokenRepository, +) +from litellm.types.integrations.model_deprecation_email import ( + MODEL_DEPRECATION_EMAIL_EVENT, + DeprecationEmailMilestone, + ModelDeprecationEmail, +) +from litellm.types.proxy.model_deprecation import ( + DEPRECATION_IDLE_POLL_SECONDS, + ModelDeprecationInfo, +) + +if TYPE_CHECKING: + from litellm import Router + from litellm.proxy.utils import PrismaClient + +_RECEIPT_PREFIX: Final = "model_deprecation_email:" +_PAGE_SIZE: Final = 100 + + +class _OptIn(BaseModel): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + enabled: bool = Field( + False, + strict=True, + validation_alias=AliasPath("email_settings", MODEL_DEPRECATION_EMAIL_EVENT), + ) + + +class _RequestCount(BaseModel): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + api_requests: int + + +class _UsedKey(BaseModel): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + api_key: str + metrics: _RequestCount = Field(alias="_sum") + + +_USED_KEYS: Final = TypeAdapter(tuple[_UsedKey, ...]) + + +class _DeliveryStatus(str, Enum): + SENDING = "sending" + SENT = "sent" + + +class _Receipt(BaseModel): + model_config = ConfigDict(frozen=True, hide_input_in_errors=True) + revision: int = 0 + status: _DeliveryStatus = _DeliveryStatus.SENDING + lease_until: datetime + + +@dataclass(frozen=True, slots=True) +class _Impact: + key_count: int = 0 + requests: int = 0 + teams: frozenset[str] = frozenset() + + +def _admin_ids(team: LiteLLM_TeamTable | None) -> frozenset[str]: + if team is None: + return frozenset() + return frozenset( + member.user_id for member in (team.members_with_roles or ()) if member.role == "admin" and member.user_id + ) + + +def _recipient_emails( + user_id: str | None, team: LiteLLM_TeamTable | None, addresses: Mapping[str, str] +) -> frozenset[str]: + owner_ids: Final[frozenset[str]] = frozenset((user_id,)) if user_id else frozenset() + resolved: Final = frozenset(addresses[owner] for owner in owner_ids | _admin_ids(team) if owner in addresses) + if team is None: + return resolved + return resolved | frozenset( + member.user_email.strip() + for member in (team.members_with_roles or ()) + if member.role == "admin" + and member.user_email + and member.user_email.strip() + and member.user_id not in addresses + ) + + +def _usage_model(model: ModelDeprecationInfo) -> str: + name: Final = model.litellm_model or model.model_name + provider: Final = model.litellm_provider + if provider and not name.startswith(f"{provider}/"): + return f"{provider}/{name}" + return name + + +def _receipt_key(email: ModelDeprecationEmail) -> str: + identity: Final = json.dumps( + [ + email.usage_model, + email.model.deprecation_date.isoformat(), + int(email.milestone), + email.recipient_email.casefold(), + ] + ) + return _RECEIPT_PREFIX + hashlib.sha256(identity.encode()).hexdigest() + + +class ModelDeprecationEmailSender: + def __init__(self, prisma_client: PrismaClient) -> None: + self.prisma: Final = prisma_client + self.config: Final = ConfigRepository(prisma_client, use_writer=True) + + async def _impacts(self, model: str, start: date, end: date) -> Mapping[str, _Impact]: + impacts: Final[dict[str, _Impact]] = {} # mutable-ok: combines per-recipient counts across bounded usage pages + offset: int = 0 # rebind-ok: each page advances the stable key grouping cursor + while True: + raw = await DailyUserSpendRepository(self.prisma).table.group_by( + by=["api_key"], + where={ + "model": model, + "date": {"gte": start.isoformat(), "lte": end.isoformat()}, + "api_requests": {"gt": 0}, + }, + sum={"api_requests": True}, + order={"api_key": "asc"}, + skip=offset, + take=_PAGE_SIZE, + ) + usage = _USED_KEYS.validate_python(raw) + if not usage: + return MappingProxyType(impacts) + keys = await VerificationTokenRepository(self.prisma).find_many_in("token", (row.api_key for row in usage)) + teams = await TeamRepository(self.prisma).find_many_in( + "team_id", frozenset(key.team_id for key in keys if key.team_id) + ) + team_by_id = {team.team_id: team for team in teams} + ids = frozenset(key.user_id for key in keys if key.user_id) | frozenset( + chain.from_iterable(_admin_ids(team) for team in teams) + ) + users = await UserRepository(self.prisma).find_many_in("user_id", ids) + addresses = { + user.user_id: user.user_email.strip() for user in users if user.user_email and user.user_email.strip() + } + counts = {row.api_key: row.metrics.api_requests for row in usage} + for key in keys: + team = team_by_id.get(key.team_id) if key.team_id else None + if key.token is None: + continue + recipients = _recipient_emails(key.user_id, team, addresses) + for address in recipients: + previous = impacts.get(address, _Impact()) + names: frozenset[str] = frozenset((team.team_alias or team.team_id,)) if team else frozenset() + impacts[address] = _Impact( + key_count=previous.key_count + 1, + requests=previous.requests + counts[key.token], + teams=previous.teams | names, + ) + if len(usage) < _PAGE_SIZE: + return MappingProxyType(impacts) + offset += _PAGE_SIZE + + async def _deliver(self, email: ModelDeprecationEmail, now: datetime, base_url: str) -> bool: + from litellm.proxy.utils import send_email + + key: Final = _receipt_key(email) + record: Final = await self.config.get_param(key) + previous: Final = _Receipt.model_validate(record.param_value) if record else None + if previous and (previous.status is _DeliveryStatus.SENT or previous.lease_until > now): + return False + revision: Final = previous.revision if previous else 0 + claim: Final = _Receipt(revision=revision + 1, lease_until=now + timedelta(minutes=5)) + if not await self.config.set_param_if_revision(key, claim.model_dump(mode="json"), revision): + return False + content: Final = render_model_deprecation_email(email, base_url=base_url) + if not await send_email( + receiver_email=email.recipient_email, + subject=content.subject, + html=content.html, + text=content.text, + message_id=f"", + ): + return False + sent: Final = claim.model_copy(update={"revision": claim.revision + 1, "status": _DeliveryStatus.SENT}) + if not await self.config.set_param_if_revision(key, sent.model_dump(mode="json"), claim.revision): + raise RuntimeError("Could not record the delivered model deprecation reminder") + return True + + async def run(self, llm_router: Router, *, now: datetime | None = None) -> int: + current: Final = now or datetime.now(timezone.utc) + settings: Final = await self.config.get_param("general_settings") + if settings is None or not _OptIn.model_validate(settings.param_value).enabled: + return 0 + base_url: Final = os.getenv("PROXY_BASE_URL", "") + if not base_url: + return 0 + await self.config.delete_params_older_than(_RECEIPT_PREFIX, current - timedelta(days=60)) + snapshot: Final = collect_model_deprecations(llm_router=llm_router, today=current.date()) + start: Final = current.date() - timedelta(days=29) + seen: Final[set[tuple[str, date]]] = set() # mutable-ok: deduplicates aliases in this pass + sent: int = 0 # rebind-ok: counts successfully delivered reminders in this pass + for model in snapshot.imminent: + if model.days_until_deprecation not in tuple(DeprecationEmailMilestone): + continue + name = _usage_model(model) + identity = (name, model.deprecation_date) + if identity in seen: + continue + seen.add(identity) + impacts = await self._impacts(name, start, current.date()) + for address, impact in sorted(impacts.items()): + email = ModelDeprecationEmail( + model=model, + usage_model=name, + recipient_email=address, + milestone=DeprecationEmailMilestone(model.days_until_deprecation), + key_count=impact.key_count, + requests=impact.requests, + team_names=tuple(sorted(impact.teams)), + start_date=start, + end_date=current.date(), + ) + if await self._deliver(email, now or datetime.now(timezone.utc), base_url): + sent += 1 + return sent + + +async def run_scheduled_model_deprecation_emails() -> None: + """Poll the email opt-in independently; a large owner mailing never delays Slack checks.""" + from litellm._logging import verbose_proxy_logger + + while True: + try: + from litellm.proxy.proxy_server import llm_router, prisma_client + + if llm_router is not None and prisma_client is not None: + await ModelDeprecationEmailSender(prisma_client).run(llm_router) + except Exception: # noqa: BLE001 # transient SMTP/DB failures must not terminate the reminder loop + verbose_proxy_logger.exception("Error sending model deprecation email reminders") + await asyncio.sleep(DEPRECATION_IDLE_POLL_SECONDS) diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index 5b81fd9a331..0f98a8fce0e 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -341,6 +341,9 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_SpendLogDelivery", "LiteLLM_ClaudeCodePluginTable", "LiteLLM_ClaudeCodePluginUpload", + "LiteLLM_ClaudeCodePluginVersion", + "LiteLLM_SkillSubmission", + "LiteLLM_ClaudeCodePluginDownload", "LiteLLM_MemoryTable", "LiteLLM_AdaptiveRouterState", "LiteLLM_AdaptiveRouterSession", diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index 909e1222ae6..ea26d549326 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -19,6 +19,12 @@ def _finite_or_none(value: float | None) -> float | None: return None if isinstance(value, float) and not math.isfinite(value) else value +def _bounded(value: float | None, integer_limit: int | None) -> float | None: + if integer_limit is None or isinstance(value, bool) or not isinstance(value, int): + return value + return max(-integer_limit, min(integer_limit, value)) + + def safe_json_structure( data: object, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, @@ -26,12 +32,14 @@ def safe_json_structure( key: str | None = None, *, non_finite_as_null: bool = False, + integer_limit: int | None = None, ) -> object: """ Rebuild data out of JSON-native pieces while detecting circular references. If a circular reference is detected then a marker string is returned. NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors. With non_finite_as_null, NaN and infinite floats become null, since jsonb rejects them. + With integer_limit, integers are clamped to that magnitude. value_transform, when given, is applied to every string leaf (and to the str() fallback for non-serializable objects) with the mapping key the leaf @@ -51,7 +59,8 @@ def safe_json_structure( cleaned = obj.replace("\x00", "") if "\x00" in obj else obj return _transform(key, cleaned) if isinstance(obj, (int, float, bool, type(None))): - return _finite_or_none(obj) if non_finite_as_null else obj + bounded: Final = _bounded(obj, integer_limit) + return _finite_or_none(bounded) if non_finite_as_null else bounded # Check for circular reference. if id(obj) in seen: return "CircularReference Detected" diff --git a/litellm/llms/github_copilot/authenticator.py b/litellm/llms/github_copilot/authenticator.py index ee0e7af7e9b..3af3b50ed8f 100644 --- a/litellm/llms/github_copilot/authenticator.py +++ b/litellm/llms/github_copilot/authenticator.py @@ -1,62 +1,30 @@ +import hashlib import json import os import time -from dataclasses import dataclass -from datetime import datetime -from typing import Any, Final, TypeVar +from collections.abc import Mapping +from typing import Final, Protocol, cast # noqa: TID251 # pins untyped clients and caches to local Protocols import httpx -from pydantic import BaseModel, ConfigDict, TypeAdapter, with_config +from pydantic import ConfigDict, TypeAdapter, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import GITHUB_COPILOT_USER_SESSION_TTL_SECONDS from litellm.litellm_core_utils.asyncify import can_block_current_thread from litellm.llms.custom_httpx.http_handler import get_httpx_client from .common_utils import ( GetAccessTokenError, - GetAPIKeyError, GetDeviceCodeError, - RefreshAPIKeyError, ) # Constants (default values — overridable via environment variables at call time) DEFAULT_GITHUB_CLIENT_ID: Final = "Iv1.b507a08c87ecfe98" DEFAULT_GITHUB_DEVICE_CODE_URL: Final = "https://github.com/login/device/code" DEFAULT_GITHUB_ACCESS_TOKEN_URL: Final = "https://github.com/login/oauth/access_token" -DEFAULT_GITHUB_API_KEY_URL: Final = "https://api.github.com/copilot_internal/v2/token" - - -class CopilotEndpoints(BaseModel): - model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True) - - api: str | None = None - - -class CopilotCachedKey(BaseModel): - model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True) - - token: str - expires_at: float - endpoints: CopilotEndpoints | None = None - - -class CopilotCachedEndpoints(BaseModel): - model_config = ConfigDict(frozen=True, strict=True, extra="ignore", hide_input_in_errors=True) - - endpoints: CopilotEndpoints | None = None - - -ApiKeyFileView = TypeVar("ApiKeyFileView", CopilotCachedKey, CopilotCachedEndpoints) - - -@dataclass(frozen=True, slots=True) -class _UnsavedApiKey: - access_token: str - key: CopilotCachedKey - - -_UNSAVED_API_KEYS: Final[dict[str, _UnsavedApiKey]] = {} # mutable-ok: refreshed keys api-key.json could not hold +DEFAULT_GITHUB_COPILOT_USER_URL: Final = "https://api.github.com/copilot_internal/user" def _read_token_file(path: str) -> str: @@ -78,11 +46,66 @@ class _AccessTokenPoll(TypedDict): error: ReadOnly[NotRequired[object]] -_JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(strict=True)) +_JSON_OBJECT: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter( + Mapping[str, object], config=ConfigDict(strict=True) +) _DEVICE_CODE: Final = TypeAdapter(_DeviceCode) _ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPoll) +class _CopilotAPIBaseCache(Protocol): + def get_cache(self, key: str) -> object: ... + + def set_cache(self, key: str, value: str, ttl: float | None = None) -> object: ... + + +class _SyncGetClient(Protocol): + def get(self, url: str, *, headers: Mapping[str, str]) -> httpx.Response: ... + + +_COPILOT_API_BASE_CACHE: Final[_CopilotAPIBaseCache] = cast( # cast-ok: [LIT006] InMemoryCache methods lack type stubs + _CopilotAPIBaseCache, + InMemoryCache( + max_size_in_memory=1000, + default_ttl=GITHUB_COPILOT_USER_SESSION_TTL_SECONDS, + ), +) + + +def _copilot_api_base_for_token(access_token: str) -> str | None: + cache_key: Final = hashlib.sha256(access_token.encode()).hexdigest() + cached_api_base: Final[object] = _COPILOT_API_BASE_CACHE.get_cache(cache_key) + if isinstance(cached_api_base, str): + return cached_api_base + + try: + response: Final = cast( # cast-ok: [LIT006] HTTPHandler.get has no return annotation + _SyncGetClient, + get_httpx_client(), + ).get( + DEFAULT_GITHUB_COPILOT_USER_URL, + headers=github_api_headers(access_token), + ) + if not response.is_success: + return None + payload: Final[Mapping[str, object]] = _JSON_OBJECT.validate_python(response.json()) + endpoints: Final[Mapping[str, object]] = _JSON_OBJECT.validate_python(payload.get("endpoints")) + api_base: Final[object] = endpoints.get("api") + if not isinstance(api_base, str): + return None + from .github_host import github_host_endpoints + from .per_user_auth import validated_copilot_api_base + + validated: Final = validated_copilot_api_base(api_base, github_host_endpoints("github.com")) + if validated != api_base.rstrip("/"): + return None + _COPILOT_API_BASE_CACHE.set_cache(cache_key, validated) + except Exception: + return None + else: + return validated + + def github_api_headers( access_token: str | None = None, ) -> dict[str, str]: # mutable-ok: returned straight into httpx handlers whose headers params require dict @@ -111,7 +134,6 @@ class Authenticator: self.token_dir, os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_FILE", "access-token"), ) - self.api_key_file = os.path.join(self.token_dir, os.getenv("GITHUB_COPILOT_API_KEY_FILE", "api-key.json")) self._ensure_token_dir() def get_access_token(self) -> str: @@ -155,7 +177,7 @@ class Authenticator: except OSError: verbose_logger.error("Error saving access token to file") return access_token - except (GetDeviceCodeError, GetAccessTokenError, RefreshAPIKeyError) as e: + except (GetDeviceCodeError, GetAccessTokenError) as e: verbose_logger.warning("Failed attempt %s: %s", attempt + 1, e) continue @@ -165,69 +187,15 @@ class Authenticator: ) def get_api_key(self) -> str: - """ - Get the API key, refreshing if necessary. - - Returns: - str: The GitHub Copilot API key. - - Raises: - GetAPIKeyError: If unable to obtain an API key. - """ - unsaved: Final = self._unsaved_api_key() - if unsaved is not None: - return unsaved.token - cached: Final = self._read_api_key_file(CopilotCachedKey) - if cached is not None and cached.expires_at > datetime.now().timestamp(): - return cached.token - if cached is not None: - verbose_logger.warning("API key expired, refreshing") - - try: - api_key_info: Final = self._refresh_api_key() - except RefreshAPIKeyError as e: - raise GetAPIKeyError( - message=f"Failed to refresh API key: {e}", - status_code=401, - ) from e - serialized: Final = json.dumps(api_key_info) - try: - with open(self.api_key_file, "w") as f: - f.write(serialized) - except OSError as e: - verbose_logger.warning( - "Keeping the refreshed API key in memory, it could not be saved to %s: %s", self.api_key_file, e - ) - self._keep_unsaved_api_key(serialized) - token: Final = api_key_info.get("token") - if not token: - raise GetAPIKeyError( - message="API key response missing token", - status_code=401, - ) - return token + """Return the stored GitHub access token, loading it with device flow when needed.""" + return self.get_access_token() def get_api_base(self) -> str | None: - """ - Get the API endpoint of the current API key. - - Returns: - Optional[str]: The GitHub Copilot API endpoint, or None if not found. - """ - cached: Final = self._unsaved_api_key() or self._read_api_key_file(CopilotCachedEndpoints) - if cached is None or cached.endpoints is None: - return None - return cached.endpoints.api - - def _read_api_key_file(self, view: type[ApiKeyFileView]) -> ApiKeyFileView | None: - try: - return view.model_validate_json(_read_token_file(self.api_key_file)) - except OSError: - verbose_logger.warning("No API key file found or error opening file") - return None - except ValueError as e: - verbose_logger.warning("Ignoring %s, not a Copilot API key file: %s", self.api_key_file, e) + """Return the Copilot API endpoint for a stored token without starting device flow.""" + access_token: Final = self._stored_access_token() + if access_token is None: return None + return _copilot_api_base_for_token(access_token) def _stored_access_token(self) -> str | None: try: @@ -236,59 +204,6 @@ class Authenticator: return None return access_token or None - def _unsaved_api_key(self) -> CopilotCachedKey | None: - unsaved: Final = _UNSAVED_API_KEYS.get(self.api_key_file) - if unsaved is None or unsaved.key.expires_at <= time.time(): - return None - return unsaved.key if unsaved.access_token == self._stored_access_token() else None - - def _keep_unsaved_api_key(self, serialized: str) -> None: - access_token: Final = self._stored_access_token() - if access_token is None: - return - try: - key: Final = CopilotCachedKey.model_validate_json(serialized) - except ValueError: - return - _UNSAVED_API_KEYS[self.api_key_file] = _UnsavedApiKey(access_token=access_token, key=key) - - def _refresh_api_key(self) -> dict[str, Any]: - """ - Refresh the API key using the access token. - - Returns: - Dict[str, Any]: The API key information including token and expiration. - - Raises: - RefreshAPIKeyError: If unable to refresh the API key. - """ - access_token: Final = self.get_access_token() - headers: Final = self._get_github_headers(access_token) - api_key_url: Final = os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL) - - max_retries: Final = 3 - for attempt in range(max_retries): - try: - sync_client = get_httpx_client() - response = sync_client.get(api_key_url, headers=headers) - response.raise_for_status() - - response_json = response.json() - - if "token" in response_json: - return response_json - else: - verbose_logger.warning("API key response missing token: %s", response_json) - except httpx.HTTPStatusError as e: - verbose_logger.error("HTTP error refreshing API key (attempt %s/%s): %s", attempt + 1, max_retries, e) - except Exception as e: - verbose_logger.error("Unexpected error refreshing API key: %s", e) - - raise RefreshAPIKeyError( - message="Failed to refresh API key after maximum retries", - status_code=401, - ) - def _ensure_token_dir(self) -> None: """Ensure the token directory exists.""" if not os.path.exists(self.token_dir): diff --git a/litellm/llms/github_copilot/chat/transformation.py b/litellm/llms/github_copilot/chat/transformation.py index 83bc8df23a6..0285bde2a87 100644 --- a/litellm/llms/github_copilot/chat/transformation.py +++ b/litellm/llms/github_copilot/chat/transformation.py @@ -16,7 +16,7 @@ from litellm.types.utils import ModelResponse from ..authenticator import Authenticator from ..common_utils import ( DEFAULT_GITHUB_COPILOT_API_BASE, - GetAPIKeyError, + GetAccessTokenError, get_copilot_default_headers, ) from ..per_user_auth import require_github_copilot_user_session @@ -56,7 +56,7 @@ class GithubCopilotConfig(OpenAIConfig): dynamic_api_base: Final = self.api_base_without_login(api_base) try: dynamic_api_key: Final = self.authenticator.get_api_key() - except GetAPIKeyError as e: + except GetAccessTokenError as e: raise AuthenticationError( model=model, llm_provider=custom_llm_provider, @@ -104,7 +104,7 @@ class GithubCopilotConfig(OpenAIConfig): return get_copilot_default_headers(session_token) try: return get_copilot_default_headers(self.authenticator.get_api_key()) - except GetAPIKeyError: + except GetAccessTokenError: return {} def validate_environment( diff --git a/litellm/llms/github_copilot/common_utils.py b/litellm/llms/github_copilot/common_utils.py index 515b7bf4fb5..fcadb0cb2ef 100644 --- a/litellm/llms/github_copilot/common_utils.py +++ b/litellm/llms/github_copilot/common_utils.py @@ -46,18 +46,6 @@ class GetAccessTokenError(GithubCopilotError): pass -class APIKeyExpiredError(GithubCopilotError): - pass - - -class RefreshAPIKeyError(GithubCopilotError): - pass - - -class GetAPIKeyError(GithubCopilotError): - pass - - def get_copilot_default_headers( api_key: str, ) -> dict[str, str]: # mutable-ok: callers merge caller headers into the returned dict diff --git a/litellm/llms/github_copilot/embedding/transformation.py b/litellm/llms/github_copilot/embedding/transformation.py index 4f308b01159..740397f0419 100644 --- a/litellm/llms/github_copilot/embedding/transformation.py +++ b/litellm/llms/github_copilot/embedding/transformation.py @@ -30,7 +30,7 @@ from litellm.utils import convert_to_model_response_object from ..authenticator import Authenticator from ..common_utils import ( DEFAULT_GITHUB_COPILOT_API_BASE, - GetAPIKeyError, + GetAccessTokenError, get_copilot_default_headers, pin_session_authorization, ) @@ -100,7 +100,7 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): return merged_headers - except GetAPIKeyError as e: + except GetAccessTokenError as e: raise AuthenticationError( model=model, llm_provider="github_copilot", diff --git a/litellm/llms/github_copilot/github_host.py b/litellm/llms/github_copilot/github_host.py index cfe122ebac4..09a5a79575a 100644 --- a/litellm/llms/github_copilot/github_host.py +++ b/litellm/llms/github_copilot/github_host.py @@ -7,8 +7,8 @@ from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST from .authenticator import ( DEFAULT_GITHUB_ACCESS_TOKEN_URL, - DEFAULT_GITHUB_API_KEY_URL, DEFAULT_GITHUB_CLIENT_ID, + DEFAULT_GITHUB_COPILOT_USER_URL, DEFAULT_GITHUB_DEVICE_CODE_URL, ) @@ -20,11 +20,9 @@ class GithubHostEndpoints: host: str device_code_url: str access_token_url: str - api_key_url: str copilot_user_url: str user_url: str client_id: str - exchanges_copilot_token: bool copilot_api_hosts: frozenset[str] copilot_api_host_suffix: str | None default_copilot_api_base: str @@ -51,17 +49,14 @@ def github_host_endpoints(host: str) -> GithubHostEndpoints: if normalized is None: raise ValueError("host must be github.com or a single-subdomain GHE.com hostname") client_id: Final = os.getenv("GITHUB_COPILOT_CLIENT_ID", DEFAULT_GITHUB_CLIENT_ID) - exchanges_copilot_token: Final = client_id == DEFAULT_GITHUB_CLIENT_ID if normalized == GITHUB_COPILOT_DEFAULT_HOST: return GithubHostEndpoints( host=normalized, device_code_url=os.getenv("GITHUB_COPILOT_DEVICE_CODE_URL", DEFAULT_GITHUB_DEVICE_CODE_URL), access_token_url=os.getenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", DEFAULT_GITHUB_ACCESS_TOKEN_URL), - api_key_url=os.getenv("GITHUB_COPILOT_API_KEY_URL", DEFAULT_GITHUB_API_KEY_URL), - copilot_user_url="https://api.github.com/copilot_internal/user", + copilot_user_url=DEFAULT_GITHUB_COPILOT_USER_URL, user_url="https://api.github.com/user", client_id=client_id, - exchanges_copilot_token=exchanges_copilot_token, copilot_api_hosts=frozenset({"githubcopilot.com"}), copilot_api_host_suffix=".githubcopilot.com", default_copilot_api_base="https://api.githubcopilot.com", @@ -74,11 +69,9 @@ def github_host_endpoints(host: str) -> GithubHostEndpoints: host=normalized, device_code_url=f"https://{normalized}/login/device/code", access_token_url=f"https://{normalized}/login/oauth/access_token", - api_key_url=f"https://{api_host}/copilot_internal/v2/token", copilot_user_url=f"https://{api_host}/copilot_internal/user", user_url=f"https://{api_host}/user", client_id=client_id, - exchanges_copilot_token=exchanges_copilot_token, copilot_api_hosts=frozenset({copilot_api_host}), copilot_api_host_suffix=None, default_copilot_api_base=f"https://{copilot_api_host}", diff --git a/litellm/llms/github_copilot/messages/transformation.py b/litellm/llms/github_copilot/messages/transformation.py index bc2b68cd02b..bed855ca178 100644 --- a/litellm/llms/github_copilot/messages/transformation.py +++ b/litellm/llms/github_copilot/messages/transformation.py @@ -10,7 +10,7 @@ from litellm.llms.anthropic.messages.native.transformation import ( from ..authenticator import Authenticator from ..common_utils import ( DEFAULT_GITHUB_COPILOT_API_BASE, - GetAPIKeyError, + GetAccessTokenError, get_copilot_default_headers, pin_session_authorization, ) @@ -84,7 +84,7 @@ class GithubCopilotAnthropicMessagesConfig(AnthropicMessagesConfig): dynamic_api_key: Final = ( user_session.token if user_session is not None else self.authenticator.get_api_key() ) - except GetAPIKeyError as e: + except GetAccessTokenError as e: raise AuthenticationError( model=model, llm_provider="github_copilot", diff --git a/litellm/llms/github_copilot/per_user_auth.py b/litellm/llms/github_copilot/per_user_auth.py index e210414ce8f..54c64c48084 100644 --- a/litellm/llms/github_copilot/per_user_auth.py +++ b/litellm/llms/github_copilot/per_user_auth.py @@ -1,7 +1,6 @@ -"""Per-user GitHub Copilot OAuth: exchange a caller's stored GitHub token for a -short-lived Copilot token, plus the device-flow helpers the proxy endpoints use -to establish that connection. The shared device login (Authenticator) is never -touched on the per-user path. +"""Per-user GitHub Copilot OAuth sessions and the device-flow helpers the proxy +endpoints use to establish those connections. The shared device login +(Authenticator) is never touched on the per-user path. """ import asyncio @@ -28,7 +27,7 @@ from litellm.constants import ( GITHUB_COPILOT_AUTH_TYPE_KEY, GITHUB_COPILOT_HOST_KEY, GITHUB_COPILOT_PER_USER_AUTH_TYPE, - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS, + GITHUB_COPILOT_USER_SESSION_TTL_SECONDS, ) from litellm.exceptions import ( APIConnectionError, @@ -39,6 +38,7 @@ from litellm.exceptions import ( ) from .authenticator import github_api_headers +from .common_utils import get_copilot_default_headers from .github_host import GithubHostEndpoints, github_host_endpoints, normalized_github_host GITHUB_COPILOT_USER_SESSION_KWARG_KEY: Final = "github_copilot_user_session" @@ -46,7 +46,6 @@ _LLM_PROVIDER: Final = "github_copilot" _DEVICE_FLOW_SCOPE: Final = "read:user" _SESSION_CACHE_MAX_SIZE: Final = 1000 _SYNC_LOCK_COUNT: Final = 64 -_ACCOUNT_LOOKUP_TTL_SECONDS: Final = 1800 class _SessionCache(Protocol): @@ -88,19 +87,19 @@ _SESSION_CACHE: Final = cast( # cast-ok: pins the untyped module cache to _Sess _SessionCache, InMemoryCache( max_size_in_memory=_SESSION_CACHE_MAX_SIZE, - default_ttl=GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS, + default_ttl=GITHUB_COPILOT_USER_SESSION_TTL_SECONDS, ), ) -_EXCHANGE_LOCKS: Final = tuple(threading.Lock() for _ in range(_SYNC_LOCK_COUNT)) +_RESOLVE_LOCKS: Final = tuple(threading.Lock() for _ in range(_SYNC_LOCK_COUNT)) _IN_FLIGHT: Final[ # mutable-ok: single-flight dedup registry mutated under its own locking dict[str, tuple[str, "asyncio.Future[GithubCopilotUserSession]"]] ] = {} @with_config(ConfigDict(extra="allow", strict=True)) -class _CopilotTokenPayload(TypedDict): - token: ReadOnly[str] - expires_at: ReadOnly[NotRequired[object]] +class _CopilotUserPayload(TypedDict): + access_type_sku: ReadOnly[NotRequired[object]] + chat_enabled: ReadOnly[NotRequired[object]] endpoints: ReadOnly[NotRequired[object]] @@ -126,20 +125,12 @@ class _CopilotEndpointsPayload(TypedDict): api: ReadOnly[NotRequired[object]] -@with_config(ConfigDict(extra="allow", strict=True)) -class _CopilotAccountPayload(TypedDict): - access_type_sku: ReadOnly[NotRequired[object]] - chat_enabled: ReadOnly[NotRequired[object]] - endpoints: ReadOnly[NotRequired[object]] - - @with_config(ConfigDict(extra="allow", strict=True)) class _GithubUserPayload(TypedDict): login: ReadOnly[str] -_COPILOT_TOKEN: Final = TypeAdapter(_CopilotTokenPayload) -_COPILOT_ACCOUNT: Final = TypeAdapter(_CopilotAccountPayload) +_COPILOT_USER: Final = TypeAdapter(_CopilotUserPayload) _DEVICE_CODE: Final = TypeAdapter(_DeviceCodePayload) _ACCESS_TOKEN_POLL: Final = TypeAdapter(_AccessTokenPollPayload) _GITHUB_USER: Final = TypeAdapter(_GithubUserPayload) @@ -237,13 +228,13 @@ def _cached_session(cache_key: str, github_token: str) -> GithubCopilotUserSessi return None if not hmac.compare_digest(cached_github_token, github_token): return None - if expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS <= time.time(): + if expires_at <= time.time(): return None return GithubCopilotUserSession(token=token, api_base=api_base) def _cache_session(cache_key: str, github_token: str, session: GithubCopilotUserSession, expires_at: float) -> None: - ttl: Final = expires_at - GITHUB_COPILOT_USER_TOKEN_SAFETY_MARGIN_SECONDS - time.time() + ttl: Final = expires_at - time.time() if ttl <= 0: return _SESSION_CACHE.set_cache(cache_key, (github_token, session.token, session.api_base, expires_at), ttl=ttl) @@ -272,57 +263,32 @@ def _session_from_response( raise _reject_connection(credential_name) if response.status_code == 429: raise CallerCredentialRateLimitError( - message="GitHub rate limited the GitHub Copilot token exchange", + message="GitHub rate limited the GitHub Copilot user lookup", llm_provider=_LLM_PROVIDER, model="", ) if not response.is_success: raise ServiceUnavailableError( - message=f"GitHub Copilot token exchange failed with status {response.status_code}", + message=f"GitHub Copilot user lookup failed with status {response.status_code}", llm_provider=_LLM_PROVIDER, model="", ) - if not endpoints.exchanges_copilot_token: - return _session_from_account(response, github_token, credential_name, cache_key, endpoints) try: - payload: Final = _COPILOT_TOKEN.validate_python(response.json()) + payload: Final = _COPILOT_USER.validate_python(response.json()) except Exception as e: raise ServiceUnavailableError( - message="GitHub Copilot token exchange returned an invalid response", + message="GitHub Copilot user lookup returned an invalid response", llm_provider=_LLM_PROVIDER, model="", ) from e - api_base: Final = validated_copilot_api_base(_endpoints_api(payload.get("endpoints")), endpoints) - session: Final = GithubCopilotUserSession(token=payload["token"], api_base=api_base) - expires_at: Final = _expires_at_seconds(payload.get("expires_at")) - if expires_at > 0: - _cache_session(cache_key, github_token, session, expires_at) - return session - - -def _session_from_account( - response: httpx.Response, - github_token: str, - credential_name: str, - cache_key: str, - endpoints: GithubHostEndpoints, -) -> GithubCopilotUserSession: - try: - payload: Final = _COPILOT_ACCOUNT.validate_python(response.json()) - except Exception as e: - raise ServiceUnavailableError( - message="GitHub Copilot account lookup returned an invalid response", - llm_provider=_LLM_PROVIDER, - model="", - ) from e if payload.get("chat_enabled") is False or payload.get("access_type_sku") == "no_access": _SESSION_CACHE.delete_cache(cache_key) raise _reject_connection(credential_name) api_base: Final = validated_copilot_api_base(_endpoints_api(payload.get("endpoints")), endpoints) session: Final = GithubCopilotUserSession(token=github_token, api_base=api_base) - _cache_session(cache_key, github_token, session, time.time() + _ACCOUNT_LOOKUP_TTL_SECONDS) + _cache_session(cache_key, github_token, session, time.time() + GITHUB_COPILOT_USER_SESSION_TTL_SECONDS) return session @@ -332,10 +298,6 @@ def _endpoints_api(payload_endpoints: object) -> object: return _COPILOT_ENDPOINTS.validate_python(payload_endpoints).get("api") -def _session_lookup_url(endpoints: GithubHostEndpoints) -> str: - return endpoints.api_key_url if endpoints.exchanges_copilot_token else endpoints.copilot_user_url - - def _to_int(raw: object) -> int: if isinstance(raw, bool): return 0 @@ -351,20 +313,7 @@ def _to_int(raw: object) -> int: return 0 -def _expires_at_seconds(raw: object) -> float: - if isinstance(raw, bool): - return 0.0 - if isinstance(raw, (int, float)): - return float(raw) - if isinstance(raw, str): - try: - return float(raw) - except ValueError: - return 0.0 - return 0.0 - - -def _do_exchange( +def _do_resolve( github_token: str, credential_name: str, cache_key: str, @@ -377,12 +326,12 @@ def _do_exchange( _SyncGetClient, litellm.module_level_client ) response: Final = client.get( - _session_lookup_url(endpoints), + endpoints.copilot_user_url, headers=github_api_headers(github_token), ) except Exception as e: raise APIConnectionError( - message="GitHub Copilot token exchange request failed", + message="GitHub Copilot user lookup request failed", llm_provider=_LLM_PROVIDER, model="", ) from e @@ -398,7 +347,7 @@ def _new_async_github_client() -> _AsyncGitHubClient: )(_LLM_PROVIDER) -async def _do_aexchange( +async def _do_aresolve( github_token: str, credential_name: str, cache_key: str, @@ -407,19 +356,19 @@ async def _do_aexchange( client: Final = _new_async_github_client() try: response: Final = await client.get( - _session_lookup_url(endpoints), + endpoints.copilot_user_url, headers=github_api_headers(github_token), ) except Exception as e: raise APIConnectionError( - message="GitHub Copilot token exchange request failed", + message="GitHub Copilot user lookup request failed", llm_provider=_LLM_PROVIDER, model="", ) from e return _session_from_response(response, github_token, credential_name, cache_key, endpoints) -def exchange_github_token( +def resolve_copilot_user_session( user_id: str, github_token: str, credential_name: str, @@ -429,15 +378,15 @@ def exchange_github_token( cached: Final = _cached_session(cache_key, github_token) if cached is not None: return cached - lock: Final = _EXCHANGE_LOCKS[hash(cache_key) % _SYNC_LOCK_COUNT] + lock: Final = _RESOLVE_LOCKS[hash(cache_key) % _SYNC_LOCK_COUNT] with lock: rechecked: Final = _cached_session(cache_key, github_token) if rechecked is not None: return rechecked - return _do_exchange(github_token, credential_name, cache_key, endpoints) + return _do_resolve(github_token, credential_name, cache_key, endpoints) -async def aexchange_github_token( +async def aresolve_copilot_user_session( user_id: str, github_token: str, credential_name: str, @@ -452,14 +401,14 @@ async def aexchange_github_token( if in_flight_entry is not None: in_flight_token, in_flight = in_flight_entry if not hmac.compare_digest(in_flight_token, github_token): - return await _do_aexchange(github_token, credential_name, cache_key, endpoints) + return await _do_aresolve(github_token, credential_name, cache_key, endpoints) if not in_flight.done() and in_flight.get_loop() is loop: return await asyncio.shield(in_flight) future: Final = loop.create_future() entry: Final = (github_token, future) _IN_FLIGHT[cache_key] = entry try: - session: Final = await _do_aexchange(github_token, credential_name, cache_key, endpoints) + session: Final = await _do_aresolve(github_token, credential_name, cache_key, endpoints) except BaseException as e: if not future.done(): future.set_exception(e) @@ -610,12 +559,41 @@ async def acheck_copilot_seat( credential_name: str, endpoints: GithubHostEndpoints, ) -> GithubCopilotUserSession: - return await aexchange_github_token( + session: Final = await aresolve_copilot_user_session( user_id=user_id, github_token=github_token, credential_name=credential_name, endpoints=endpoints, ) + client: Final = _new_async_github_client() + models_url: Final = f"{session.api_base}/models" + try: + response: Final = await client.get( + models_url, + headers=get_copilot_default_headers(session.token), + ) + except Exception as e: + raise APIConnectionError( + message="GitHub Copilot models check request failed", + llm_provider=_LLM_PROVIDER, + model="", + ) from e + if response.status_code in (401, 403, 404): + _SESSION_CACHE.delete_cache(_session_cache_key(user_id, credential_name, endpoints.host)) + raise _reject_connection(credential_name) + if response.status_code == 429: + raise CallerCredentialRateLimitError( + message="GitHub rate limited the GitHub Copilot models check", + llm_provider=_LLM_PROVIDER, + model="", + ) + if not response.is_success: + raise ServiceUnavailableError( + message=f"GitHub Copilot models check failed with status {response.status_code}", + llm_provider=_LLM_PROVIDER, + model="", + ) + return session def github_copilot_auth_mode(litellm_credential_name: object, auth_type_value: object) -> bool: @@ -804,7 +782,7 @@ def attach_github_copilot_user_session( credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else "" endpoints: Final = github_copilot_host_endpoints(credential_name) user_id, github_token = _caller_connection(kwargs, credential_name, endpoints) - session: Final = exchange_github_token( + session: Final = resolve_copilot_user_session( user_id=user_id, github_token=github_token, credential_name=credential_name, @@ -826,7 +804,7 @@ async def aattach_github_copilot_user_session( credential_name: Final = credential_name_obj if isinstance(credential_name_obj, str) else "" endpoints: Final = github_copilot_host_endpoints(credential_name) user_id, github_token = _caller_connection(kwargs, credential_name, endpoints) - session: Final = await aexchange_github_token( + session: Final = await aresolve_copilot_user_session( user_id=user_id, github_token=github_token, credential_name=credential_name, diff --git a/litellm/llms/github_copilot/responses/transformation.py b/litellm/llms/github_copilot/responses/transformation.py index b87a3d543e7..f337d4ae0d3 100644 --- a/litellm/llms/github_copilot/responses/transformation.py +++ b/litellm/llms/github_copilot/responses/transformation.py @@ -38,7 +38,7 @@ from litellm.utils import cached_get_model_info_helper from ..authenticator import Authenticator from ..common_utils import ( DEFAULT_GITHUB_COPILOT_API_BASE, - GetAPIKeyError, + GetAccessTokenError, get_copilot_default_headers, pin_session_authorization, ) @@ -256,7 +256,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): return merged_headers - except GetAPIKeyError as e: + except GetAccessTokenError as e: raise AuthenticationError( model=model, llm_provider="github_copilot", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 76fd85d001b..aba1957f03c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -32320,7 +32320,9 @@ "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", "successor_model": "gemini/gemini-3.8-flash", "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, - "output_cost_per_token_above_200k_tokens_batches": 7.5e-06 + "input_cost_per_token_above_200k_tokens_flex": 1.25e-06, + "output_cost_per_token_above_200k_tokens_batches": 7.5e-06, + "output_cost_per_token_above_200k_tokens_flex": 7.5e-06 }, "gemini/gemini-2.5-computer-use-preview-10-2025": { "deprecation_date": "2026-07-28", @@ -32899,7 +32901,9 @@ "output_cost_per_token_flex": 6e-06, "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, "input_cost_per_token_above_200k_tokens_batches": 2e-06, - "output_cost_per_token_above_200k_tokens_batches": 9e-06 + "input_cost_per_token_above_200k_tokens_flex": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, + "output_cost_per_token_above_200k_tokens_flex": 9e-06 }, "gemini/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -32965,7 +32969,9 @@ "output_cost_per_token_flex": 6e-06, "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, "input_cost_per_token_above_200k_tokens_batches": 2e-06, - "output_cost_per_token_above_200k_tokens_batches": 9e-06 + "input_cost_per_token_above_200k_tokens_flex": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, + "output_cost_per_token_above_200k_tokens_flex": 9e-06 }, "gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, @@ -60260,6 +60266,7 @@ "rpm": 10, "gemini_audio_only_live": true, "input_cost_per_second": 8.33333333333e-05, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "supports_response_schema": false, "supports_reasoning": true }, @@ -72140,6 +72147,7 @@ "supports_audio_input": true, "tpm": 250000, "rpm": 10, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -72161,6 +72169,7 @@ "supports_audio_input": true, "tpm": 250000, "rpm": 10, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -83130,6 +83139,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_prompt_caching": true, "supports_function_calling": true, "supports_response_schema": true, @@ -83146,6 +83160,11 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_prompt_caching": true, "supports_function_calling": true, "supports_response_schema": true, @@ -83155,6 +83174,7 @@ "litellm_provider": "empiriolabs", "mode": "chat", "max_input_tokens": 202745, + "max_output_tokens": 128000, "input_cost_per_token": 8.25e-07, "output_cost_per_token": 3.301e-06, "input_cost_per_token_above_32k_tokens": 1.1e-06, @@ -83178,6 +83198,11 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_function_calling": true, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/glm-5-2" @@ -83230,6 +83255,11 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_function_calling": true, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/glm-5-3" @@ -83245,6 +83275,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_function_calling": true, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/glm-5-3-flash" @@ -83721,6 +83756,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b" }, @@ -83737,6 +83777,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-27b" }, @@ -83753,6 +83798,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b" }, @@ -83769,6 +83819,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b" }, @@ -83830,6 +83885,10 @@ "max_output_tokens": 65536, "input_cost_per_token": 2.9e-08, "output_cost_per_token": 2.87e-07, + "input_cost_per_token_above_128k_tokens": 1.15e-07, + "output_cost_per_token_above_128k_tokens": 1.147e-06, + "input_cost_per_token_above_256k_tokens": 1.72e-07, + "output_cost_per_token_above_256k_tokens": 1.72e-06, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, @@ -83852,6 +83911,11 @@ "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, "supports_audio_input": true, "source": "https://empiriolabs.ai/models/qwen3-5-omni-flash" }, @@ -83865,6 +83929,11 @@ "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, "supports_response_schema": true, "supports_audio_input": true, "source": "https://empiriolabs.ai/models/qwen3-5-omni-plus" @@ -83924,6 +83993,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-6-27b" }, @@ -84421,7 +84495,7 @@ "max_output_tokens": 131072, "input_cost_per_token": 1e-06, "output_cost_per_token": 2.7e-06, - "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost": 1e-07, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, @@ -85310,6 +85384,24 @@ "output_cost_per_token": 6.09e-07, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/XiaomiMiMo/MiMo-V2.6-Flash": { + "cache_read_input_token_cost": 2.24e-09, + "input_cost_per_token": 1.12e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 2.24e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/XiaomiMiMo/MiMo-V2.6-Pro": { + "cache_read_input_token_cost": 2.88e-09, + "input_cost_per_token": 3.48e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 6.96e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/deepseek-ai/DeepSeek-Prover-V2-671B": { "input_cost_per_token": 5e-07, "litellm_provider": "gmi", @@ -85434,6 +85526,15 @@ "output_cost_per_token": 3.168e-06, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/deepseek-ai/DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 3.6e-09, + "input_cost_per_token": 1.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/deepseek-ai/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 1.4e-08, "input_cost_per_token": 4.4e-07, @@ -85443,6 +85544,24 @@ "output_cost_per_token": 1.32e-06, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/deepseek-ai/deepseek-v4.1-flash-1017": { + "cache_read_input_token_cost": 3.6e-09, + "input_cost_per_token": 1.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/gmicloud/Ripple-65B-Preview": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/google/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, @@ -85759,6 +85878,14 @@ "output_cost_per_token": 2.5e-07, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/openai/gpt-oss-20b": { + "input_cost_per_token": 4e-08, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/stepfun-ai/Step-3.7-Flash": { "cache_read_input_token_cost": 4e-08, "input_cost_per_token": 2e-07, @@ -85803,6 +85930,15 @@ "output_cost_per_token": 6e-07, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/tencent/hy4": { + "cache_read_input_token_cost": 2.1e-08, + "input_cost_per_token": 4.17e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.2505e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/tencent/hy4-preview": { "cache_read_input_token_cost": 4.2e-08, "input_cost_per_token": 8.34e-07, @@ -86176,5 +86312,41 @@ "output_cost_per_token": 6e-07, "source": "https://api.together.ai/v1/models", "supports_prompt_caching": true + }, + "vertex_ai/gemini-omni-1.2-flash": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 131072, + "max_output_tokens": 260640, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1beta/interactions" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "video" + ], + "supports_audio_input": false, + "supports_audio_output": false, + "supports_computer_use": false, + "supports_function_calling": false, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_system_messages": false, + "supports_tool_choice": false, + "supports_url_context": false, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": false } } diff --git a/litellm/models/plugin_upload.py b/litellm/models/plugin_upload.py index 41fbf901a03..bd5e85c79cc 100644 --- a/litellm/models/plugin_upload.py +++ b/litellm/models/plugin_upload.py @@ -1,6 +1,23 @@ -from collections.abc import Mapping +from collections.abc import AsyncIterator, Mapping from dataclasses import dataclass -from typing import Protocol +from datetime import datetime +from enum import Enum +from types import MappingProxyType +from typing import Literal, Protocol + + +def has_agent_skill_metadata(manifest: Mapping[str, object]) -> bool: + return isinstance(manifest.get("agent_skill"), dict) + + +def uploaded_plugin_source(base_url: str, name: str, digest: str) -> Mapping[str, str]: + return MappingProxyType( + { + "source": "archive", + "url": f"{base_url.rstrip('/')}/claude-code/plugins/{name}/archive/{digest}.zip", + "sha256": digest, + } + ) @dataclass(frozen=True, slots=True) @@ -19,11 +36,130 @@ class PluginArchiveMetadata: enabled: bool source: Mapping[str, str] manifest_json: str | None + id: str | None = None + version: str | None = None + description: str | None = None + has_agent_skill_metadata: bool = False + updated_at: datetime | None = None + + +@dataclass(frozen=True, slots=True) +class PluginPublication: + name: str + version: str | None + description: str | None + manifest: Mapping[str, object] + created_by: str | None + content: bytes | None = None + + +@dataclass(frozen=True, slots=True) +class PluginPublicationBase: + id: str + updated_at: datetime | None + + +@dataclass(frozen=True, slots=True) +class PublishedPlugin: + id: str + name: str + version: str + enabled: bool + updated_at: datetime + + +@dataclass(frozen=True, slots=True) +class PluginPublicationError: + kind: Literal["invalid_version", "conflict", "not_found"] + message: str + + +@dataclass(frozen=True, slots=True) +class PluginVersion: + id: str + version: str | None + description: str | None + manifest: Mapping[str, object] + published_at: datetime | None + published_by: str | None + archive_available: bool + + +@dataclass(frozen=True, slots=True) +class PluginSnapshot: + id: str + name: str + version: str | None + description: str | None + manifest: Mapping[str, object] + enabled: bool + created_at: datetime | None + updated_at: datetime | None + created_by: str | None + + +@dataclass(frozen=True, slots=True) +class PluginDownloader: + authenticated: bool + user_id: str | None + user_email: str | None + key_alias: str | None + team_id: str | None + team_alias: str | None + + +class PluginDownloadResult(Enum): + RECORDED = "recorded" + SKILL_NOT_FOUND = "skill_not_found" + + +@dataclass(frozen=True, slots=True) +class PluginDownloadEvent: + id: str + downloaded_at: datetime + version: str | None + digest: str + authenticated: bool + user_id: str | None + user_email: str | None + key_alias: str | None + team_id: str | None + team_alias: str | None + + +@dataclass(frozen=True, slots=True) +class PluginDownloadsResponse: + total_downloads: int + anonymous_downloads: int + downloads: tuple[PluginDownloadEvent, ...] + next_cursor: str | None class PluginUploadStore(Protocol): async def create(self, upload: PluginUpload) -> str | None: ... + async def get_plugin(self, name: str) -> PluginSnapshot | None: ... + + async def publish( + self, publication: PluginPublication, expected: PluginPublicationBase | None = None + ) -> PublishedPlugin | PluginPublicationError: ... + + async def versions(self, metadata: PluginArchiveMetadata, limit: int, offset: int) -> tuple[PluginVersion, ...]: ... + + async def version(self, metadata: PluginArchiveMetadata, version_id: str) -> PluginVersion | None: ... + + async def version_content(self, metadata: PluginArchiveMetadata, version_id: str) -> bytes | None: ... + async def get_metadata(self, name: str) -> PluginArchiveMetadata | None: ... async def get_content(self, metadata: PluginArchiveMetadata) -> bytes | None: ... + + def public_metadata(self) -> AsyncIterator[PluginArchiveMetadata]: ... + + async def record_download( + self, metadata: PluginArchiveMetadata, digest: str, downloader: PluginDownloader + ) -> PluginDownloadResult: ... + + async def downloads(self, name: str, limit: int, cursor: int | None) -> PluginDownloadsResponse | None: ... + + async def download_counts(self, plugin_ids: tuple[str, ...]) -> Mapping[str, int]: ... diff --git a/litellm/models/skill_submission.py b/litellm/models/skill_submission.py new file mode 100644 index 00000000000..09d722fa69f --- /dev/null +++ b/litellm/models/skill_submission.py @@ -0,0 +1,110 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from datetime import datetime +from typing import Literal, Protocol, TypeAlias + +from litellm.models.plugin_upload import ( + PluginPublication, + PluginPublicationBase, + PluginPublicationError, + PublishedPlugin, +) +from litellm.types.proxy.skill_scanning import SkillScan + +SubmissionStatus: TypeAlias = Literal["pending_review", "approved", "rejected"] +SubmissionOperation: TypeAlias = Literal["create", "update"] + + +@dataclass(frozen=True, slots=True) +class SkillProposal: + name: str + version: str | None + description: str | None + manifest: Mapping[str, object] + + +@dataclass(frozen=True, slots=True) +class SkillSubmission: + id: str + plugin_name: str + operation: SubmissionOperation + status: SubmissionStatus + proposal: SkillProposal + has_archive: bool + submitted_by: str + created_at: datetime + base_plugin_id: str | None + base_updated_at: datetime | None + reviewed_by: str | None + reviewed_at: datetime | None + rejection_reason: str | None + + +@dataclass(frozen=True, slots=True) +class SubmissionPage: + submissions: tuple[SkillSubmission, ...] + has_more: bool + + +@dataclass(frozen=True, slots=True) +class SubmissionLimitExceeded: + message: str + + +@dataclass(frozen=True, slots=True) +class SubmissionMissing: + pass + + +@dataclass(frozen=True, slots=True) +class SubmissionConflict: + message: str + + +@dataclass(frozen=True, slots=True) +class SubmissionReviewed: + submission: SkillSubmission + plugin: PublishedPlugin | None + published_source: Mapping[str, str] | None = None + + +SubmissionReviewResult: TypeAlias = SubmissionReviewed | SubmissionMissing | SubmissionConflict | PluginPublicationError + + +@dataclass(frozen=True, slots=True) +class SkillScanInput: + scan: SkillScan + source: Mapping[str, str] + content: bytes | None + + +class SkillSubmissionStore(Protocol): + async def get_scan(self, submission_id: str) -> SkillScan | SubmissionMissing | None: ... + + async def start_scan( + self, submission_id: str, scan: SkillScan + ) -> SkillScanInput | SubmissionMissing | SubmissionConflict: ... + + async def finish_scan(self, submission_id: str, scan: SkillScan) -> None: ... + + async def submit( + self, upload: PluginPublication, submitted_by: str, base: PluginPublicationBase | None + ) -> SkillSubmission | SubmissionLimitExceeded: ... + + async def list( + self, submitted_by: str | None, status: SubmissionStatus | None, limit: int, offset: int + ) -> SubmissionPage: ... + + async def get(self, submission_id: str, submitted_by: str | None) -> SkillSubmission | None: ... + + async def get_archive(self, submission_id: str, submitted_by: str | None) -> bytes | None: ... + + async def review( + self, + submission_id: str, + reviewer: str, + approve: bool, + reason: str | None, + *, + publication_base_url: str | None = None, + ) -> SubmissionReviewResult: ... diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index 2f613396413..01c22dc6725 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -48,6 +48,15 @@ def _mount_app(prefix: str, attr_name: str = "app") -> Callable[["FastAPI", obje return _register +def _include_skill_hub(app: "FastAPI", module: object) -> None: + from litellm.proxy.discovery_endpoints.skill_hub_submissions import router + from litellm.proxy.discovery_endpoints.skill_scanning import router as scan_router + + _include_router()(app, module) + app.include_router(router) + app.include_router(scan_router) + + @dataclass(frozen=True) class LazyFeature: name: str @@ -283,6 +292,12 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( module_path="litellm.proxy.decisions_endpoints.endpoints", path_prefixes=("/v1/decisions", "/decisions", "/v1/systemone", "/systemone", "/v1/extraction", "/extraction"), ), + LazyFeature( + name="skill_hub", + module_path="litellm.proxy.discovery_endpoints.skill_hub_endpoints", + path_prefixes=("/skill-hub",), + register_fn=_include_skill_hub, + ), LazyFeature( name="claude_code_marketplace", module_path="litellm.proxy.anthropic_endpoints.claude_code_endpoints", diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index ab85b6f8842..7d9e230f8ef 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -7541,6 +7541,145 @@ "title": "PluginAuthor", "type": "object" }, + "PluginDownloadEvent": { + "properties": { + "authenticated": { + "title": "Authenticated", + "type": "boolean" + }, + "digest": { + "title": "Digest", + "type": "string" + }, + "downloaded_at": { + "format": "date-time", + "title": "Downloaded At", + "type": "string" + }, + "id": { + "title": "Id", + "type": "string" + }, + "key_alias": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key Alias" + }, + "team_alias": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team Alias" + }, + "team_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team Id" + }, + "user_email": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Email" + }, + "user_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Id" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Version" + } + }, + "required": [ + "id", + "downloaded_at", + "version", + "digest", + "authenticated", + "user_id", + "user_email", + "key_alias", + "team_id", + "team_alias" + ], + "title": "PluginDownloadEvent", + "type": "object" + }, + "PluginDownloadsResponse": { + "properties": { + "anonymous_downloads": { + "title": "Anonymous Downloads", + "type": "integer" + }, + "downloads": { + "items": { + "$ref": "#/components/schemas/PluginDownloadEvent" + }, + "title": "Downloads", + "type": "array" + }, + "next_cursor": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Next Cursor" + }, + "total_downloads": { + "title": "Total Downloads", + "type": "integer" + } + }, + "required": [ + "total_downloads", + "anonymous_downloads", + "downloads", + "next_cursor" + ], + "title": "PluginDownloadsResponse", + "type": "object" + }, "PluginListItem": { "description": "Plugin item in list responses.", "properties": { @@ -7912,7 +8051,7 @@ } ], "default": "1.0.0", - "description": "Semantic version", + "description": "Semantic version, defaults to 1.0.0 on creation", "title": "Version" } }, @@ -7950,7 +8089,7 @@ "type": "object" }, "UpdatePluginRequest": { - "description": "Request body for replacing an existing plugin.\n\nThe plugin name is the resource identity and is supplied as the path\nparameter, so it cannot be changed here. This is a full replace: omitted\nfields reset to their defaults, so version is cleared rather than\ndefaulting to the create-time \"1.0.0\".", + "description": "Request body for replacing an existing plugin.\n\nThe plugin name is the resource identity and is supplied as the path\nparameter, so it cannot be changed here. This is a full replace: omitted\nfields reset to their defaults, except version which automatically increments\nthe published patch version when omitted or null.", "properties": { "author": { "anyOf": [ @@ -7999,6 +8138,19 @@ "description": "Skill domain (e.g., 'Productivity')", "title": "Domain" }, + "expected_updated_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Published updated_at read by the editor. A stale value returns 409", + "title": "Expected Updated At" + }, "homepage": { "anyOf": [ { @@ -8072,7 +8224,7 @@ "type": "null" } ], - "description": "Semantic version; cleared if omitted", + "description": "Explicit higher SemVer, or omit/null for an automatic patch bump", "title": "Version" } }, @@ -8203,7 +8355,7 @@ } ], "default": "1.0.0", - "description": "Semantic version", + "description": "Semantic version, defaults to 1.0.0 on creation", "title": "Version" } }, @@ -8535,7 +8687,7 @@ ] }, "put": { - "description": "Update an existing plugin in the LiteLLM marketplace.\n\nThe plugin is identified by its name in the path, which is the resource\nidentity and cannot be changed here. This is a full replace, not a merge:\nthe manifest is rebuilt from the request body, so any optional field left\nout is reset to its default (e.g. an omitted version is cleared, not kept).\nSend the full desired state.\n\nReturns 404 if no plugin with the given name exists; use\nPOST /claude-code/plugins to create a new plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n - plugin_name: Name of the plugin to update (path parameter)\n - source: Plugin source reference (github, url, git-subdir, or archive format)\n - version: Semantic version (optional)\n - description: Plugin description (optional)\n - author: Author information (optional)\n - homepage: Plugin homepage URL (optional)\n - keywords: Search keywords (optional)\n - category: Plugin category (optional)\n - installation_preference: 'available', 'auto_install', or 'required' (optional)\n\nReturns:\n Update status (action is always \"updated\") and plugin information.\n\nExample:\n ```bash\n curl -X PUT http://localhost:4000/claude-code/plugins/my-plugin \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n \"version\": \"2.0.0\",\n \"description\": \"My awesome plugin\"\n }'\n ```", + "description": "Update an existing plugin in the LiteLLM marketplace.\n\nThe plugin is identified by its name in the path, which is the resource\nidentity and cannot be changed here. This is a full replace, not a merge:\nthe manifest is rebuilt from the request body, so any optional field left\nout is reset to its default, except version which automatically increments\nthe published patch version when omitted or null.\nSend the full desired state.\n\nReturns 404 if no plugin with the given name exists; use\nPOST /claude-code/plugins to create a new plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n - plugin_name: Name of the plugin to update (path parameter)\n - source: Plugin source reference (github, url, git-subdir, or archive format)\n - version: Semantic version (optional)\n - description: Plugin description (optional)\n - author: Author information (optional)\n - homepage: Plugin homepage URL (optional)\n - keywords: Search keywords (optional)\n - category: Plugin category (optional)\n - installation_preference: 'available', 'auto_install', or 'required' (optional)\n\nReturns:\n Update status (action is always \"updated\") and plugin information.\n\nExample:\n ```bash\n curl -X PUT http://localhost:4000/claude-code/plugins/my-plugin \\\n -H \"Authorization: Bearer sk-...\" \\\n -H \"Content-Type: application/json\" \\\n -d '{\n \"source\": {\"source\": \"github\", \"repo\": \"org/my-plugin\"},\n \"version\": \"2.0.0\",\n \"description\": \"My awesome plugin\"\n }'\n ```", "operationId": "update_plugin_claude_code_plugins__plugin_name__put", "parameters": [ { @@ -8690,6 +8842,87 @@ ] } }, + "/claude-code/plugins/{plugin_name}/downloads": { + "get": { + "description": "Admin-only ZIP download totals and authenticated downloader history.\n\nAnonymous downloads contribute only to totals. Counts include repeat requests and client prefetches,\nstarting when tracking is deployed. External sources and installation completion are not observed.", + "operationId": "get_plugin_downloads_claude_code_plugins__plugin_name__downloads_get", + "parameters": [ + { + "in": "path", + "name": "plugin_name", + "required": true, + "schema": { + "title": "Plugin Name", + "type": "string" + } + }, + { + "description": "Maximum number of download events", + "in": "query", + "name": "limit", + "required": false, + "schema": { + "default": 25, + "description": "Maximum number of download events", + "maximum": 100, + "minimum": 1, + "title": "Limit", + "type": "integer" + } + }, + { + "description": "Previous next_cursor", + "in": "query", + "name": "cursor", + "required": false, + "schema": { + "anyOf": [ + { + "pattern": "^[1-9][0-9]{0,18}$", + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Previous next_cursor", + "title": "Cursor" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/PluginDownloadsResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Plugin Downloads", + "tags": [ + "claude_code_marketplace" + ] + } + }, "/claude-code/plugins/{plugin_name}/enable": { "post": { "description": "Enable a disabled plugin.\n\nRequires a proxy admin API key.\n\nParameters:\n - plugin_name: The name of the plugin to enable", @@ -58849,6 +59082,2622 @@ } } }, + "skill_hub": { + "components": { + "schemas": { + "Body_replace_skill_archive_skill_hub_skills__name__upload_put": { + "properties": { + "expected_updated_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Expected Updated At" + }, + "file": { + "contentMediaType": "application/octet-stream", + "description": "Replacement Agent Skills ZIP with the same declared skill name", + "title": "File", + "type": "string" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Version" + } + }, + "required": [ + "file" + ], + "title": "Body_replace_skill_archive_skill_hub_skills__name__upload_put", + "type": "object" + }, + "Body_upload_skill_skill_hub_skills_upload_post": { + "properties": { + "file": { + "contentMediaType": "application/octet-stream", + "description": "Agent Skills ZIP with SKILL.md and optional supporting files", + "title": "File", + "type": "string" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Version" + } + }, + "required": [ + "file" + ], + "title": "Body_upload_skill_skill_hub_skills_upload_post", + "type": "object" + }, + "HTTPValidationError": { + "properties": { + "detail": { + "items": { + "$ref": "#/components/schemas/ValidationError" + }, + "title": "Detail", + "type": "array" + } + }, + "title": "HTTPValidationError", + "type": "object" + }, + "ImportSkillRequest": { + "properties": { + "path": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Repository-relative skill directory", + "title": "Path" + }, + "ref": { + "default": "HEAD", + "description": "Git commit, tag or branch to import", + "title": "Ref", + "type": "string" + }, + "skill": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Skill name to select when a repository contains several skills", + "title": "Skill" + }, + "url": { + "description": "Public skills.sh skill URL or GitHub repository URL", + "title": "Url", + "type": "string" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Semantic version, or omit/null to start at 1.0.0", + "title": "Version" + } + }, + "required": [ + "url" + ], + "title": "ImportSkillRequest", + "type": "object" + }, + "PluginDownloadEvent": { + "properties": { + "authenticated": { + "title": "Authenticated", + "type": "boolean" + }, + "digest": { + "title": "Digest", + "type": "string" + }, + "downloaded_at": { + "format": "date-time", + "title": "Downloaded At", + "type": "string" + }, + "id": { + "title": "Id", + "type": "string" + }, + "key_alias": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Key Alias" + }, + "team_alias": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team Alias" + }, + "team_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Team Id" + }, + "user_email": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Email" + }, + "user_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Id" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Version" + } + }, + "required": [ + "id", + "downloaded_at", + "version", + "digest", + "authenticated", + "user_id", + "user_email", + "key_alias", + "team_id", + "team_alias" + ], + "title": "PluginDownloadEvent", + "type": "object" + }, + "PluginDownloadsResponse": { + "properties": { + "anonymous_downloads": { + "title": "Anonymous Downloads", + "type": "integer" + }, + "downloads": { + "items": { + "$ref": "#/components/schemas/PluginDownloadEvent" + }, + "title": "Downloads", + "type": "array" + }, + "next_cursor": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Next Cursor" + }, + "total_downloads": { + "title": "Total Downloads", + "type": "integer" + } + }, + "required": [ + "total_downloads", + "anonymous_downloads", + "downloads", + "next_cursor" + ], + "title": "PluginDownloadsResponse", + "type": "object" + }, + "SkillHubAgentSkill": { + "properties": { + "description": { + "title": "Description", + "type": "string" + }, + "digest": { + "title": "Digest", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "origin": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Origin" + } + }, + "required": [ + "name", + "description", + "digest", + "origin" + ], + "title": "SkillHubAgentSkill", + "type": "object" + }, + "SkillHubList": { + "properties": { + "count": { + "title": "Count", + "type": "integer" + }, + "next_offset": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Next Offset" + }, + "skills": { + "items": { + "$ref": "#/components/schemas/SkillHubSkill" + }, + "title": "Skills", + "type": "array" + } + }, + "required": [ + "skills", + "count", + "next_offset" + ], + "title": "SkillHubList", + "type": "object" + }, + "SkillHubMetadata": { + "additionalProperties": false, + "properties": { + "description": { + "anyOf": [ + { + "maxLength": 1024, + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "domain": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Domain" + }, + "expected_updated_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Published updated_at read by the editor. A stale value returns 409", + "title": "Expected Updated At" + }, + "keywords": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Keywords" + }, + "namespace": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Namespace" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Explicit higher SemVer, or omit/null for an automatic patch bump", + "title": "Version" + } + }, + "title": "SkillHubMetadata", + "type": "object" + }, + "SkillHubPendingSkill": { + "properties": { + "action": { + "const": "submitted", + "default": "submitted", + "title": "Action", + "type": "string" + }, + "proposal": { + "$ref": "#/components/schemas/SkillHubProposal" + }, + "status": { + "const": "pending_review", + "default": "pending_review", + "title": "Status", + "type": "string" + }, + "submission_id": { + "title": "Submission Id", + "type": "string" + } + }, + "required": [ + "submission_id", + "proposal" + ], + "title": "SkillHubPendingSkill", + "type": "object" + }, + "SkillHubProposal": { + "properties": { + "agent_skill": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubAgentSkill" + }, + { + "type": "null" + } + ] + }, + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "domain": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Domain" + }, + "keywords": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Keywords" + }, + "name": { + "title": "Name", + "type": "string" + }, + "namespace": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Namespace" + }, + "source": { + "additionalProperties": { + "type": "string" + }, + "title": "Source", + "type": "object" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Version" + } + }, + "required": [ + "name", + "description", + "version", + "source" + ], + "title": "SkillHubProposal", + "type": "object" + }, + "SkillHubPublication": { + "additionalProperties": false, + "properties": {}, + "title": "SkillHubPublication", + "type": "object" + }, + "SkillHubRejection": { + "additionalProperties": false, + "properties": { + "reason": { + "anyOf": [ + { + "maxLength": 2000, + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Reason" + } + }, + "title": "SkillHubRejection", + "type": "object" + }, + "SkillHubReviewResponse": { + "properties": { + "action": { + "enum": [ + "approved", + "rejected" + ], + "title": "Action", + "type": "string" + }, + "skill": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubSkill" + }, + { + "type": "null" + } + ] + }, + "status": { + "const": "success", + "default": "success", + "title": "Status", + "type": "string" + }, + "submission": { + "$ref": "#/components/schemas/SkillHubSubmission" + } + }, + "required": [ + "action", + "submission", + "skill" + ], + "title": "SkillHubReviewResponse", + "type": "object" + }, + "SkillHubSkill": { + "properties": { + "content_url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Content Url" + }, + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "domain": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Domain" + }, + "download_count": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "description": "Gateway ZIP downloads, included in proxy-admin catalog listings", + "title": "Download Count" + }, + "download_url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Download Url" + }, + "id": { + "title": "Id", + "type": "string" + }, + "keywords": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Keywords" + }, + "name": { + "title": "Name", + "type": "string" + }, + "namespace": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Namespace" + }, + "source": { + "additionalProperties": { + "type": "string" + }, + "title": "Source", + "type": "object" + }, + "updated_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Updated At" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Version" + }, + "visibility": { + "enum": [ + "public", + "private" + ], + "title": "Visibility", + "type": "string" + } + }, + "required": [ + "id", + "name", + "visibility", + "source" + ], + "title": "SkillHubSkill", + "type": "object" + }, + "SkillHubSubmission": { + "properties": { + "base_skill_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Base Skill Id" + }, + "base_updated_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Base Updated At" + }, + "created_at": { + "format": "date-time", + "title": "Created At", + "type": "string" + }, + "has_archive": { + "title": "Has Archive", + "type": "boolean" + }, + "id": { + "title": "Id", + "type": "string" + }, + "operation": { + "enum": [ + "create", + "update" + ], + "title": "Operation", + "type": "string" + }, + "proposal": { + "$ref": "#/components/schemas/SkillHubProposal" + }, + "rejection_reason": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Rejection Reason" + }, + "reviewed_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Reviewed At" + }, + "reviewed_by": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Reviewed By" + }, + "skill_name": { + "title": "Skill Name", + "type": "string" + }, + "status": { + "enum": [ + "pending_review", + "approved", + "rejected" + ], + "title": "Status", + "type": "string" + }, + "submitted_by": { + "title": "Submitted By", + "type": "string" + } + }, + "required": [ + "id", + "skill_name", + "operation", + "status", + "proposal", + "has_archive", + "submitted_by", + "created_at", + "base_skill_id", + "base_updated_at", + "reviewed_by", + "reviewed_at", + "rejection_reason" + ], + "title": "SkillHubSubmission", + "type": "object" + }, + "SkillHubSubmissionDetail": { + "properties": { + "current_skill": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubSkill" + }, + { + "type": "null" + } + ] + }, + "submission": { + "$ref": "#/components/schemas/SkillHubSubmission" + } + }, + "required": [ + "submission", + "current_skill" + ], + "title": "SkillHubSubmissionDetail", + "type": "object" + }, + "SkillHubSubmissionList": { + "properties": { + "count": { + "title": "Count", + "type": "integer" + }, + "has_more": { + "title": "Has More", + "type": "boolean" + }, + "submissions": { + "items": { + "$ref": "#/components/schemas/SkillHubSubmission" + }, + "title": "Submissions", + "type": "array" + } + }, + "required": [ + "submissions", + "count", + "has_more" + ], + "title": "SkillHubSubmissionList", + "type": "object" + }, + "SkillHubVersion": { + "properties": { + "content_kind": { + "enum": [ + "stored_archive", + "source_reference" + ], + "title": "Content Kind", + "type": "string" + }, + "description": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Description" + }, + "download_url": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Download Url" + }, + "id": { + "title": "Id", + "type": "string" + }, + "manifest": { + "additionalProperties": true, + "title": "Manifest", + "type": "object" + }, + "published_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Published At" + }, + "published_by": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Published By" + }, + "version": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Version" + } + }, + "required": [ + "id", + "version", + "description", + "manifest", + "published_at", + "published_by", + "content_kind", + "download_url" + ], + "title": "SkillHubVersion", + "type": "object" + }, + "SkillHubVersions": { + "properties": { + "next_offset": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Next Offset" + }, + "versions": { + "items": { + "$ref": "#/components/schemas/SkillHubVersion" + }, + "title": "Versions", + "type": "array" + } + }, + "required": [ + "versions", + "next_offset" + ], + "title": "SkillHubVersions", + "type": "object" + }, + "SkillScan": { + "properties": { + "archive_sha256": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Archive Sha256" + }, + "completed_at": { + "anyOf": [ + { + "format": "date-time", + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Completed At" + }, + "error": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Error" + }, + "expires_at": { + "format": "date-time", + "title": "Expires At", + "type": "string" + }, + "id": { + "title": "Id", + "type": "string" + }, + "report": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillScanReport" + }, + { + "type": "null" + } + ] + }, + "requested_by": { + "title": "Requested By", + "type": "string" + }, + "started_at": { + "format": "date-time", + "title": "Started At", + "type": "string" + }, + "status": { + "enum": [ + "running", + "completed", + "failed" + ], + "title": "Status", + "type": "string" + } + }, + "required": [ + "id", + "status", + "requested_by", + "started_at", + "expires_at" + ], + "title": "SkillScan", + "type": "object" + }, + "SkillScanFinding": { + "properties": { + "explanation": { + "anyOf": [ + { + "maxLength": 8000, + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Explanation" + }, + "file": { + "maxLength": 2048, + "title": "File", + "type": "string" + }, + "message": { + "maxLength": 8000, + "title": "Message", + "type": "string" + }, + "remediation": { + "anyOf": [ + { + "maxLength": 8000, + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Remediation" + }, + "rule_id": { + "maxLength": 200, + "title": "Rule Id", + "type": "string" + }, + "severity": { + "enum": [ + "LOW", + "MEDIUM", + "HIGH", + "CRITICAL" + ], + "title": "Severity", + "type": "string" + }, + "start_line": { + "anyOf": [ + { + "minimum": 1.0, + "type": "integer" + }, + { + "type": "null" + } + ], + "title": "Start Line" + } + }, + "required": [ + "rule_id", + "severity", + "message", + "file" + ], + "title": "SkillScanFinding", + "type": "object" + }, + "SkillScanReport": { + "properties": { + "complete": { + "title": "Complete", + "type": "boolean" + }, + "findings": { + "items": { + "$ref": "#/components/schemas/SkillScanFinding" + }, + "maxItems": 500, + "title": "Findings", + "type": "array" + }, + "limitations": { + "default": [], + "items": { + "maxLength": 2000, + "type": "string" + }, + "maxItems": 100, + "title": "Limitations", + "type": "array" + }, + "risk_score": { + "maximum": 100.0, + "minimum": 0.0, + "title": "Risk Score", + "type": "integer" + }, + "scan_mode": { + "enum": [ + "static-only", + "static+llm" + ], + "title": "Scan Mode", + "type": "string" + }, + "scanner_version": { + "maxLength": 100, + "minLength": 1, + "title": "Scanner Version", + "type": "string" + } + }, + "required": [ + "scanner_version", + "complete", + "scan_mode", + "risk_score", + "findings" + ], + "title": "SkillScanReport", + "type": "object" + }, + "SkillScanRequest": { + "additionalProperties": false, + "properties": {}, + "title": "SkillScanRequest", + "type": "object" + }, + "SkillScanResponse": { + "properties": { + "configured": { + "title": "Configured", + "type": "boolean" + }, + "scan": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillScan" + }, + { + "type": "null" + } + ] + } + }, + "required": [ + "configured", + "scan" + ], + "title": "SkillScanResponse", + "type": "object" + }, + "ValidationError": { + "properties": { + "ctx": { + "title": "Context", + "type": "object" + }, + "input": { + "title": "Input" + }, + "loc": { + "items": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "integer" + } + ] + }, + "title": "Location", + "type": "array" + }, + "msg": { + "title": "Message", + "type": "string" + }, + "type": { + "title": "Error Type", + "type": "string" + } + }, + "required": [ + "loc", + "msg", + "type" + ], + "title": "ValidationError", + "type": "object" + } + } + }, + "paths": { + "/skill-hub/skills": { + "get": { + "description": "Browse the caller's public and granted private skills.", + "operationId": "list_skills_skill_hub_skills_get", + "parameters": [ + { + "in": "query", + "name": "query", + "required": false, + "schema": { + "default": "", + "title": "Query", + "type": "string" + } + }, + { + "in": "query", + "name": "offset", + "required": false, + "schema": { + "default": 0, + "minimum": 0, + "title": "Offset", + "type": "integer" + } + }, + { + "in": "query", + "name": "limit", + "required": false, + "schema": { + "default": 50, + "maximum": 100, + "minimum": 1, + "title": "Limit", + "type": "integer" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubList" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "List Skills", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/import": { + "post": { + "description": "Import a public skills.sh or GitHub skill as a portable snapshot.", + "operationId": "import_skill_skill_hub_skills_import_post", + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ImportSkillRequest" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubSkill" + }, + { + "$ref": "#/components/schemas/SkillHubPendingSkill" + } + ], + "title": "Response Import Skill Skill Hub Skills Import Post" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Import Skill", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/upload": { + "post": { + "description": "Publish a portable skill while retaining Claude marketplace compatibility.", + "operationId": "upload_skill_skill_hub_skills_upload_post", + "requestBody": { + "content": { + "multipart/form-data": { + "schema": { + "$ref": "#/components/schemas/Body_upload_skill_skill_hub_skills_upload_post" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubSkill" + }, + { + "$ref": "#/components/schemas/SkillHubPendingSkill" + } + ], + "title": "Response Upload Skill Skill Hub Skills Upload Post" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Upload Skill", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/{name}": { + "delete": { + "description": "Delete a skill and its managed archive.", + "operationId": "remove_skill_skill_hub_skills__name__delete", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + } + ], + "responses": { + "204": { + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Remove Skill", + "tags": [ + "skill_hub" + ] + }, + "get": { + "description": "Read one visible skill and its portable download URL.", + "operationId": "get_skill_skill_hub_skills__name__get", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubSkill" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Skill", + "tags": [ + "skill_hub" + ] + }, + "put": { + "description": "Replace registry metadata, preserving files. Members submit an immutable proposal for admin review.", + "operationId": "replace_skill_metadata_skill_hub_skills__name__put", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubMetadata" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubSkill" + }, + { + "$ref": "#/components/schemas/SkillHubPendingSkill" + } + ], + "title": "Response Replace Skill Metadata Skill Hub Skills Name Put" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Replace Skill Metadata", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/{name}/archive": { + "get": { + "description": "Download the complete portable skill, with SKILL.md at the ZIP root.", + "operationId": "download_skill_skill_hub_skills__name__archive_get", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + }, + { + "in": "query", + "name": "digest", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Digest" + } + } + ], + "responses": { + "200": { + "content": { + "application/zip": { + "schema": { + "format": "binary", + "type": "string" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Download Skill", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/{name}/content": { + "get": { + "description": "Read the original SKILL.md without downloading supporting files.", + "operationId": "skill_content_skill_hub_skills__name__content_get", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "text/markdown": {} + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Skill Content", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/{name}/publish": { + "post": { + "description": "Make a skill public using the existing Hub publication policy.", + "operationId": "publish_skill_skill_hub_skills__name__publish_post", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubPublication" + }, + { + "type": "null" + } + ], + "title": " Body" + } + } + } + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubSkill" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Publish Skill", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/{name}/unpublish": { + "post": { + "description": "Remove a skill from public discovery while retaining explicit grants.", + "operationId": "unpublish_skill_skill_hub_skills__name__unpublish_post", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubPublication" + }, + { + "type": "null" + } + ], + "title": " Body" + } + } + } + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubSkill" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Unpublish Skill", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/{name}/upload": { + "put": { + "description": "Replace skill files, retaining registry fields. Members propose an update against the original editor timestamp.", + "operationId": "replace_skill_archive_skill_hub_skills__name__upload_put", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "multipart/form-data": { + "schema": { + "$ref": "#/components/schemas/Body_replace_skill_archive_skill_hub_skills__name__upload_put" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubSkill" + }, + { + "$ref": "#/components/schemas/SkillHubPendingSkill" + } + ], + "title": "Response Replace Skill Archive Skill Hub Skills Name Upload Put" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Replace Skill Archive", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/{name}/versions": { + "get": { + "description": "Read published metadata history. External references are not retained artifacts.", + "operationId": "list_skill_versions_skill_hub_skills__name__versions_get", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + }, + { + "in": "query", + "name": "limit", + "required": false, + "schema": { + "default": 20, + "maximum": 100, + "minimum": 1, + "title": "Limit", + "type": "integer" + } + }, + { + "in": "query", + "name": "offset", + "required": false, + "schema": { + "default": 0, + "minimum": 0, + "title": "Offset", + "type": "integer" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubVersions" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "List Skill Versions", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/{name}/versions/{version_id}": { + "get": { + "description": "Read a retained publication under the current skill's visibility policy.", + "operationId": "get_skill_version_skill_hub_skills__name__versions__version_id__get", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + }, + { + "in": "path", + "name": "version_id", + "required": true, + "schema": { + "title": "Version Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubVersion" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Skill Version", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/{name}/versions/{version_id}/archive": { + "get": { + "description": "Download a retained publication as a portable ZIP with SKILL.md at its root.", + "operationId": "download_skill_version_skill_hub_skills__name__versions__version_id__archive_get", + "parameters": [ + { + "in": "path", + "name": "name", + "required": true, + "schema": { + "title": "Name", + "type": "string" + } + }, + { + "in": "path", + "name": "version_id", + "required": true, + "schema": { + "title": "Version Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/zip": { + "schema": { + "format": "binary", + "type": "string" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Download Skill Version", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/skills/{plugin_name}/downloads": { + "get": { + "description": "Admin-only ZIP download totals and authenticated downloader history.\n\nAnonymous downloads contribute only to totals. Counts include repeat requests and client prefetches,\nstarting when tracking is deployed. External sources and installation completion are not observed.", + "operationId": "get_plugin_downloads_skill_hub_skills__plugin_name__downloads_get", + "parameters": [ + { + "in": "path", + "name": "plugin_name", + "required": true, + "schema": { + "title": "Plugin Name", + "type": "string" + } + }, + { + "description": "Maximum number of download events", + "in": "query", + "name": "limit", + "required": false, + "schema": { + "default": 25, + "description": "Maximum number of download events", + "maximum": 100, + "minimum": 1, + "title": "Limit", + "type": "integer" + } + }, + { + "description": "Previous next_cursor", + "in": "query", + "name": "cursor", + "required": false, + "schema": { + "anyOf": [ + { + "pattern": "^[1-9][0-9]{0,18}$", + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Previous next_cursor", + "title": "Cursor" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/PluginDownloadsResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Plugin Downloads", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/submissions": { + "get": { + "description": "Read a page of owned proposals. Proxy admins and admin viewers can read every owner's proposals.", + "operationId": "list_submissions_skill_hub_submissions_get", + "parameters": [ + { + "in": "query", + "name": "status", + "required": false, + "schema": { + "anyOf": [ + { + "enum": [ + "pending_review", + "approved", + "rejected" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Status" + } + }, + { + "in": "query", + "name": "limit", + "required": false, + "schema": { + "default": 20, + "maximum": 100, + "minimum": 1, + "title": "Limit", + "type": "integer" + } + }, + { + "in": "query", + "name": "offset", + "required": false, + "schema": { + "default": 0, + "minimum": 0, + "title": "Offset", + "type": "integer" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubSubmissionList" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "List Submissions", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/submissions/{submission_id}": { + "get": { + "description": "Read an immutable proposal and its currently visible published skill.", + "operationId": "get_submission_skill_hub_submissions__submission_id__get", + "parameters": [ + { + "in": "path", + "name": "submission_id", + "required": true, + "schema": { + "title": "Submission Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubSubmissionDetail" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Submission", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/submissions/{submission_id}/approve": { + "put": { + "description": "Publish a pending proposal and record its approval atomically. Only proxy admins may review.", + "operationId": "approve_submission_skill_hub_submissions__submission_id__approve_put", + "parameters": [ + { + "in": "path", + "name": "submission_id", + "required": true, + "schema": { + "title": "Submission Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubReviewResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Approve Submission", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/submissions/{submission_id}/archive": { + "get": { + "description": "Download an owned proposal as a private portable ZIP. Legacy nonportable archives return 404.", + "operationId": "download_submission_skill_hub_submissions__submission_id__archive_get", + "parameters": [ + { + "in": "path", + "name": "submission_id", + "required": true, + "schema": { + "title": "Submission Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/zip": { + "schema": { + "format": "binary", + "type": "string" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Download Submission", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/submissions/{submission_id}/reject": { + "put": { + "description": "Reject a pending proposal without changing the published skill. Only proxy admins may review.", + "operationId": "reject_submission_skill_hub_submissions__submission_id__reject_put", + "parameters": [ + { + "in": "path", + "name": "submission_id", + "required": true, + "schema": { + "title": "Submission Id", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillHubRejection" + }, + { + "type": "null" + } + ], + "title": "Body" + } + } + } + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillHubReviewResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Reject Submission", + "tags": [ + "skill_hub" + ] + } + }, + "/skill-hub/submissions/{submission_id}/scan": { + "get": { + "description": "Read the latest manual security scan. Only proxy admins may read reports.", + "operationId": "get_skill_scan_skill_hub_submissions__submission_id__scan_get", + "parameters": [ + { + "in": "path", + "name": "submission_id", + "required": true, + "schema": { + "title": "Submission Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillScanResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Skill Scan", + "tags": [ + "skill_hub" + ] + }, + "post": { + "description": "Start an optional scan of a pending submission. Scans never change approval decisions.", + "operationId": "start_skill_scan_skill_hub_submissions__submission_id__scan_post", + "parameters": [ + { + "in": "path", + "name": "submission_id", + "required": true, + "schema": { + "title": "Submission Id", + "type": "string" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "anyOf": [ + { + "$ref": "#/components/schemas/SkillScanRequest" + }, + { + "type": "null" + } + ], + "title": " Request" + } + } + } + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/SkillScanResponse" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Start Skill Scan", + "tags": [ + "skill_hub" + ] + } + } + } + }, "tools": { "components": { "schemas": { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 1864c3ae82e..d3ac264518e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -88,6 +88,7 @@ from litellm.types.proxy.carried_budget_state import ( ) from litellm.types.proxy.control_plane_endpoints import WorkerRegistryEntry from litellm.types.proxy.request_tag_mappings import RequestTagMapping +from litellm.types.proxy.skill_scanning import SkillScannerConfig from litellm.types.proxy.spend_capture_rate import SpendCaptureRateCheckSettings from litellm.types.router import AllowedModelRegion, RouterErrors, UpdateRouterConfig from litellm.types.router_weights import validate_router_settings_dict @@ -449,6 +450,28 @@ class LiteLLMRoutes(enum.Enum): "/v1/traces/{trace_id}/spans/{span_id}/error", ) + skill_hub_routes = ( + "/skill-hub/skills", + "/skill-hub/skills/upload", + "/skill-hub/skills/import", + "/skill-hub/skills/{name}", + "/skill-hub/skills/{name}/archive", + "/skill-hub/skills/{name}/content", + "/skill-hub/skills/{plugin_name}/downloads", + "/skill-hub/skills/{name}/versions", + "/skill-hub/skills/{name}/versions/{version_id}", + "/skill-hub/skills/{name}/versions/{version_id}/archive", + "/skill-hub/skills/{name}/publish", + "/skill-hub/skills/{name}/unpublish", + "/skill-hub/skills/{name}/upload", + "/skill-hub/submissions", + "/skill-hub/submissions/{submission_id}", + "/skill-hub/submissions/{submission_id}/archive", + "/skill-hub/submissions/{submission_id}/approve", + "/skill-hub/submissions/{submission_id}/reject", + "/skill-hub/submissions/{submission_id}/scan", + ) + litellm_native_routes = [ "/rag/ingest", "/v1/rag/ingest", @@ -468,6 +491,7 @@ class LiteLLMRoutes(enum.Enum): "/claude-code/plugins", "/claude-code/plugins/upload", "/claude-code/plugins/{plugin_name}/archive/{digest}.zip", + "/claude-code/plugins/{plugin_name}/downloads", "/claude-code/plugins/{plugin_name}", ] @@ -698,6 +722,7 @@ class LiteLLMRoutes(enum.Enum): + key_management_routes + mcp_management_routes + list(agent_management_routes) + + list(skill_hub_routes) ) spend_tracking_routes = [ @@ -940,6 +965,7 @@ class LiteLLMRoutes(enum.Enum): # Agent registry - reads are role-scoped and writes are proxy-admin-gated # inside agent_endpoints/endpoints.py *agent_management_routes, + *skill_hub_routes, ] # routes that manage their own allowed/disallowed logic ## Org Admin Routes ## @@ -2481,8 +2507,8 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): """ use_spend_logs_ledger: bool = Field( - False, - description="Apply request accounting from stored SpendLogs deliveries. Enable only after all accounting writers are upgraded and ledger activation gates pass.", + True, + description="Apply request accounting from stored SpendLogs deliveries. Set to false to roll back to the previous in-memory and Redis spend queues.", ) request_tag_mappings: Sequence[RequestTagMapping] | None = Field( @@ -3019,6 +3045,9 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="If True, forwards client headers (e.g. Authorization) to the LLM API. Required for Claude Code with Max subscription.", ) + skill_scanner: SkillScannerConfig | None = Field( + None, description="Opt-in external scanner for manual admin review of skill submissions" + ) mcp_required_fields: list[str] | None = Field( None, description="List of MCP server fields that must be filled in for a submission to pass standards checks (e.g. ['description', 'source_url', 'alias']).", diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index b8c90fe7e93..26cb0a2f3e6 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -27,14 +27,22 @@ from datetime import datetime, timezone from types import MappingProxyType from typing import Annotated, Final, Protocol, TypedDict, get_args -from fastapi import APIRouter, Depends, File, Form, HTTPException, Request, UploadFile +from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile from fastapi.responses import JSONResponse, Response from pydantic import Json, TypeAdapter -from typing_extensions import ReadOnly +from typing_extensions import ReadOnly, assert_never from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.url_utils import is_https_url_with_host -from litellm.models.plugin_upload import PluginUpload +from litellm.models.plugin_upload import ( + PluginDownloader, + PluginDownloadResult, + PluginDownloadsResponse, + PluginPublication, + PluginPublicationBase, + PluginPublicationError, + has_agent_skill_metadata, +) from litellm.proxy._types import CommonProxyErrors, ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import ( SkillVisibility, @@ -125,7 +133,23 @@ async def get_prisma_client() -> object: _get_prisma_client: Final = get_prisma_client -async def _header_key_caller(request: Request) -> UserAPIKeyAuth: +def plugin_archive_url(request: Request, name: str, digest: str) -> str: + from litellm.proxy.utils import get_custom_url + + return get_custom_url( + request_base_url=str(request.base_url), route=f"claude-code/plugins/{name}/archive/{digest}.zip" + ) + + +def plugin_archive_source_matches( + request: Request, name: str, source: Mapping[str, str], digest: str, *, has_agent_skill_metadata: bool +) -> bool: + if has_agent_skill_metadata: + return source.get("source") == "archive" and source.get("sha256") == digest + return source == {"source": "archive", "url": plugin_archive_url(request, name, digest), "sha256": digest} + + +async def header_key_caller(request: Request) -> UserAPIKeyAuth: try: caller: Final = await user_api_key_auth( request=request, @@ -154,7 +178,7 @@ async def _header_key_caller(request: Request) -> UserAPIKeyAuth: async def _marketplace_caller(request: Request, key: str | None) -> UserAPIKeyAuth | None: if marketplace_auth_required(): - return await _header_key_caller(request) + return await header_key_caller(request) if key: return await user_api_key_auth(request=request, api_key=f"Bearer {key}") return None @@ -173,8 +197,6 @@ async def upload_plugin( metadata: Annotated[Json[UploadPluginRequest], Form(description="JSON skill metadata without a source URL")], ) -> RegisterPluginResponse: """Create a skill and its uploaded ZIP atomically. Duplicate names return 409.""" - from litellm.proxy.utils import get_custom_url - _require_proxy_admin(user_api_key_dict) if uploads is None: raise _error_response(500, CommonProxyErrors.db_not_connected_error.value) @@ -189,16 +211,13 @@ async def upload_plugin( digest: Final = hashlib.sha256(archive).hexdigest() source: Final = { "source": "archive", - "url": get_custom_url( - request_base_url=str(request.base_url), - route=f"claude-code/plugins/{metadata.name}/archive/{digest}.zip", - ), + "url": plugin_archive_url(request, metadata.name, digest), "sha256": digest, } spec: Final = RegisterPluginRequest.model_validate({**metadata.model_dump(), "source": source}) _validate_installation_preference(spec) - plugin_id: Final = await uploads.create( - PluginUpload( + published: Final = await uploads.publish( + PluginPublication( name=spec.name, version=spec.version, description=spec.description, @@ -207,15 +226,15 @@ async def upload_plugin( created_by=user_api_key_dict.user_id, ) ) - if plugin_id is None: - raise _name_conflict_error(spec.name) + if isinstance(published, PluginPublicationError): + raise _publication_error(published) return RegisterPluginResponse( status="success", action="created", plugin=PluginResponse( - id=plugin_id, + id=published.id, name=spec.name, - version=spec.version, + version=published.version, description=spec.description, source=source, installation_preference=spec.installation_preference, @@ -224,6 +243,27 @@ async def upload_plugin( ) +def require_download_recorded(result: PluginDownloadResult) -> None: + match result: + case PluginDownloadResult.RECORDED: + return + case PluginDownloadResult.SKILL_NOT_FOUND: + raise _error_response(404, "Skill archive not found") + case _: + assert_never(result) + + +def plugin_downloader(caller: UserAPIKeyAuth | None) -> PluginDownloader: + return PluginDownloader( + authenticated=caller is not None, + user_id=caller.user_id if caller is not None else None, + user_email=caller.user_email if caller is not None else None, + key_alias=caller.key_alias if caller is not None else None, + team_id=caller.team_id if caller is not None else None, + team_alias=caller.team_alias if caller is not None else None, + ) + + @router.get( "/claude-code/plugins/{plugin_name}/archive/{digest}.zip", tags=["Claude Code Marketplace"], @@ -237,10 +277,8 @@ async def download_plugin_archive( uploads: PluginUploadRepositoryDep, ) -> Response: """Download an uploaded ZIP under the same visibility policy as its skill.""" - from litellm.proxy.utils import get_custom_url - caller: Final = ( - await _header_key_caller(request) + await header_key_caller(request) if marketplace_key_header_present(request) else await _marketplace_caller(request, None) ) @@ -248,18 +286,18 @@ async def download_plugin_archive( if uploads is None: raise _error_response(500, CommonProxyErrors.db_not_connected_error.value) archive: Final = await uploads.get_metadata(plugin_name) - expected_url: Final = get_custom_url( - request_base_url=str(request.base_url), route=f"claude-code/plugins/{plugin_name}/archive/{digest}.zip" - ) if ( archive is None or not visibility.allows(archive) - or archive.source != {"source": "archive", "url": expected_url, "sha256": digest} + or not plugin_archive_source_matches( + request, plugin_name, archive.source, digest, has_agent_skill_metadata=archive.has_agent_skill_metadata + ) ): raise _error_response(404, "Skill archive not found") content: Final = await uploads.get_content(archive) if content is None: raise _error_response(404, "Skill archive not found") + require_download_recorded(await uploads.record_download(archive, digest, plugin_downloader(caller))) return Response( content=content, media_type="application/zip", @@ -271,6 +309,35 @@ async def download_plugin_archive( ) +@router.get( + "/claude-code/plugins/{plugin_name}/downloads", + tags=["Claude Code Marketplace"], + response_model=PluginDownloadsResponse, +) +async def get_plugin_downloads( + plugin_name: str, + uploads: PluginUploadRepositoryDep, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + limit: Annotated[int, Query(ge=1, le=100, description="Maximum number of download events")] = 25, + cursor: Annotated[str | None, Query(pattern=r"^[1-9][0-9]{0,18}$", description="Previous next_cursor")] = None, +) -> PluginDownloadsResponse: + """Admin-only ZIP download totals and authenticated downloader history. + + Anonymous downloads contribute only to totals. Counts include repeat requests and client prefetches, + starting when tracking is deployed. External sources and installation completion are not observed. + """ + _require_proxy_admin(user_api_key_dict) + if uploads is None: + raise _error_response(500, CommonProxyErrors.db_not_connected_error.value) + before: Final = int(cursor) if cursor is not None else None + if before is not None and before > 9223372036854775807: + raise _error_response(422, "Invalid download cursor") + result: Final = await uploads.downloads(plugin_name, limit, before) + if result is None: + raise _error_response(404, "Plugin not found") + return result + + @router.get( "/claude-code/marketplace.json", tags=["Claude Code Marketplace"], @@ -464,10 +531,14 @@ def _error_response(status_code: int, message: str) -> HTTPException: return HTTPException(status_code=status_code, detail={"error": message}) -def _name_conflict_error(name: str) -> HTTPException: - return _error_response( - 409, f"A skill named '{name}' already exists. Update the existing skill instead of adding it again." - ) +def _publication_error(error: PluginPublicationError) -> HTTPException: + match error.kind: + case "invalid_version": + return _error_response(400, error.message) + case "not_found": + return _error_response(404, error.message) + case "conflict": + return _error_response(409, error.message) def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: @@ -490,6 +561,7 @@ def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: ) async def register_plugin( request: RegisterPluginRequest, + uploads: PluginUploadRepositoryDep, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -532,12 +604,11 @@ async def register_plugin( }' ``` """ - from prisma.errors import UniqueViolationError - try: _require_proxy_admin(user_api_key_dict) - prisma_client: Final = await get_prisma_client() + if uploads is None: + raise _error_response(500, CommonProxyErrors.db_not_connected_error.value) if not re.match(r"^[a-z0-9-]+$", request.name): raise HTTPException( @@ -548,30 +619,19 @@ async def register_plugin( _validate_plugin_source(request.source) _validate_installation_preference(request) - existing: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.find_unique( - where={"name": request.name} - ) - if existing: - raise _name_conflict_error(request.name) - manifest: Final[Mapping[str, object]] = _build_plugin_manifest(request.name, request) - try: - plugin: Final[_PluginRecord] = await ClaudeCodePluginRepository(prisma_client).table.create( - data={ - "name": request.name, - "version": request.version, - "description": request.description, - "manifest_json": json.dumps(manifest), - "files_json": "{}", - "enabled": True, - "created_at": datetime.now(timezone.utc), - "updated_at": datetime.now(timezone.utc), - "created_by": user_api_key_dict.user_id, - } + plugin: Final = await uploads.publish( + PluginPublication( + name=request.name, + version=request.version, + description=request.description, + manifest=manifest, + created_by=user_api_key_dict.user_id, ) - except UniqueViolationError: - raise _name_conflict_error(request.name) + ) + if isinstance(plugin, PluginPublicationError): + raise _publication_error(plugin) verbose_proxy_logger.info("Plugin %s created successfully", request.name) @@ -582,7 +642,7 @@ async def register_plugin( id=plugin.id, name=plugin.name, version=plugin.version, - description=plugin.description, + description=request.description, source=request.source, installation_preference=request.installation_preference, enabled=plugin.enabled, @@ -608,7 +668,7 @@ async def register_plugin( async def list_plugins( enabled_only: bool = False, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): +) -> ListPluginsResponse: """ List all plugins in the marketplace. @@ -640,6 +700,7 @@ async def list_plugins( version=p.version, description=p.description, source=manifest.get("source", {}), + has_agent_skill_metadata=has_agent_skill_metadata(MANIFEST_FIELDS.validate_python(manifest)), author=manifest.get("author"), homepage=manifest.get("homepage"), keywords=manifest.get("keywords"), @@ -746,6 +807,7 @@ async def get_plugin( async def update_plugin( plugin_name: str, request: UpdatePluginRequest, + uploads: PluginUploadRepositoryDep, user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], ): """ @@ -754,7 +816,8 @@ async def update_plugin( The plugin is identified by its name in the path, which is the resource identity and cannot be changed here. This is a full replace, not a merge: the manifest is rebuilt from the request body, so any optional field left - out is reset to its default (e.g. an omitted version is cleared, not kept). + out is reset to its default, except version which automatically increments + the published patch version when omitted or null. Send the full desired state. Returns 404 if no plugin with the given name exists; use @@ -792,6 +855,8 @@ async def update_plugin( try: _require_proxy_admin(user_api_key_dict) + if uploads is None: + raise _error_response(500, CommonProxyErrors.db_not_connected_error.value) prisma_client: Final = await get_prisma_client() @@ -799,26 +864,35 @@ async def update_plugin( existing: Final[_PluginRecord | None] = await repository.table.find_unique(where={"name": plugin_name}) if not existing: raise _error_response(404, f"Plugin '{plugin_name}' not found") + if request.expected_updated_at is not None and request.expected_updated_at != existing.updated_at: + raise _error_response(409, "The published skill changed. Reload it before publishing an update") current_manifest: Final = MANIFEST_FIELDS.validate_json(existing.manifest_json or "{}") if request.source != current_manifest.get("source"): _validate_plugin_source(request.source) _validate_installation_preference(request) - manifest: Final[Mapping[str, object]] = _build_plugin_manifest(plugin_name, request) + manifest: Final[Mapping[str, object]] = { + **_build_plugin_manifest(plugin_name, request), + **( + {"agent_skill": current_manifest["agent_skill"]} + if "agent_skill" in current_manifest and request.source == current_manifest.get("source") + else {} + ), + } - plugin: Final[_PluginRecord | None] = await repository.table.update( - where={"name": plugin_name}, - data={ - "version": request.version, - "description": request.description, - "manifest_json": json.dumps(manifest), - "files_json": "{}", - "updated_at": datetime.now(timezone.utc), - }, + plugin: Final = await uploads.publish( + PluginPublication( + name=plugin_name, + version=request.version, + description=request.description, + manifest=manifest, + created_by=user_api_key_dict.user_id, + ), + PluginPublicationBase(existing.id, existing.updated_at), ) - if plugin is None: - raise _error_response(404, f"Plugin '{plugin_name}' not found") + if isinstance(plugin, PluginPublicationError): + raise _publication_error(plugin) verbose_proxy_logger.info("Plugin %s updated successfully", plugin_name) @@ -829,7 +903,7 @@ async def update_plugin( id=plugin.id, name=plugin.name, version=plugin.version, - description=plugin.description, + description=request.description, source=request.source, installation_preference=request.installation_preference, enabled=plugin.enabled, diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 59fd00cec43..9c433af487e 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -909,13 +909,14 @@ class ResetBudgetJob: db_span("reset_budget_cascade", "LiteLLM_BudgetTable"), budget_cascade_unit_of_work(self._new_batch) as uow, ): + uow.row_locks.queue_row_locks(cascade.budget_ids) + _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) uow.team_memberships.queue_total_spend_repair(cascade.budget_ids) _queue_budget_linked_resets(uow.team_memberships, cascade) _queue_budget_linked_resets(uow.keys, cascade, extra=_LINKED_KEYS_WHERE) - _queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE) + _queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE) _queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE) - _queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE) for budget_id, budget_reset_at in cascade.budget_resets: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) diff --git a/litellm/proxy/config_resolvers/alerting.py b/litellm/proxy/config_resolvers/alerting.py index 4de7197f88b..1d7fbe0fffa 100644 --- a/litellm/proxy/config_resolvers/alerting.py +++ b/litellm/proxy/config_resolvers/alerting.py @@ -11,6 +11,7 @@ from typing import Final from litellm.proxy.config_resolvers._descriptors import FieldDescriptor EMAIL_DESCRIPTORS: Final[tuple[FieldDescriptor, ...]] = ( + FieldDescriptor("PROXY_BASE_URL", "PROXY_BASE_URL", "PROXY_BASE_URL"), FieldDescriptor("SMTP_HOST", "SMTP_HOST", "SMTP_HOST"), FieldDescriptor("SMTP_PORT", "SMTP_PORT", "SMTP_PORT", default="587"), FieldDescriptor("SMTP_TLS", "SMTP_TLS", "SMTP_TLS", default="True"), diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 80f09df6c86..ff57596232f 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -328,13 +328,13 @@ _SpendTableName = Literal[ "agent_list_transactions", ] _SPEND_TABLE_COMMIT_ORDER: Final[tuple[_SpendTableName, ...]] = ( - "user_list_transactions", - "end_user_list_transactions", - "key_list_transactions", - "team_list_transactions", - "team_member_list_transactions", "org_list_transactions", "org_member_list_transactions", + "team_list_transactions", + "user_list_transactions", + "end_user_list_transactions", + "team_member_list_transactions", + "key_list_transactions", "project_list_transactions", "tag_list_transactions", "model_access_group_list_transactions", @@ -563,6 +563,24 @@ SET spend = "LiteLLM_TeamMembership".spend + EXCLUDED.spend, """ +def _org_members_in_key_order(spend_by_member_key: Mapping[str, float]) -> tuple[tuple[tuple[str, str], float], ...]: + """Sorted by the decoded (organization_id, user_id), the order every other writer locks these rows in, since + the percent-encoded key does not sort the same way.""" + decoded: Final = ( + ((unquote(key.split("::")[1]), unquote(key.split("::")[3])), cost) for key, cost in spend_by_member_key.items() + ) + return tuple(sorted(decoded)) + + +def _charged_team_ids(updates: DBSpendUpdateTransactions) -> tuple[str, ...]: + members: Final = updates.get("team_member_list_transactions") or {} + return tuple( + sorted( + {*(updates.get("team_list_transactions") or {}), *(key.split("::")[1] for key in members)}, + ) + ) + + async def _write_team_member_spend(transaction: _SpendTransaction, spend_by_member_key: Mapping[str, float]) -> None: # key is "team_id::::user_id::"; locks are taken in sorted team_id order like the team endpoints rows: Final = sorted((key.split("::")[1], key.split("::")[3], cost) for key, cost in spend_by_member_key.items()) @@ -712,7 +730,7 @@ class DBSpendUpdateWriter: proxy_settings: Final = cast( # cast-ok: proxy settings use string keys Mapping[str, object], general_settings ) - use_ledger: Final = proxy_settings.get("use_spend_logs_ledger") is True + use_ledger: Final = proxy_settings.get("use_spend_logs_ledger", True) is not False request_fields: Final = _SPEND_METADATA_ADAPTER.validate_python(kwargs or {}) if kwargs is not None: kwargs[ACCOUNTING_LEDGER_KWARG] = use_ledger # rebind-ok: spend counters leave ledger windows to the worker @@ -863,8 +881,12 @@ class DBSpendUpdateWriter: payload=payload, prisma_client=prisma_client, disable_spend_logs=disable_spend_logs ) if "accounting_payload" in payload: + from litellm.proxy.utils import accounting_only_spend_log + await self._insert_spend_log_to_db( - payload={**payload, "accounting_log_disabled": disable_spend_logs}, + payload=accounting_only_spend_log(payload) + if disable_spend_logs + else {**payload, "accounting_log_disabled": False}, prisma_client=prisma_client, llm_router=llm_router, ) @@ -2294,6 +2316,10 @@ class DBSpendUpdateWriter: transaction: _SpendTransaction, updates: DBSpendUpdateTransactions, ) -> None: + """Every team advisory lock is taken before the first row lock, the order the team roster writers use, so + applying a batch of charges cannot deadlock with a concurrent roster change.""" + for team_id in _charged_team_ids(updates): + _ = await transaction.execute_raw(_TEAM_ADVISORY_LOCK_SQL, team_id) await self._commit_spend_updates_to_db( prisma_client, 0, proxy_logging_obj, updates, outer_transaction=transaction ) @@ -2313,6 +2339,104 @@ class DBSpendUpdateWriter: """ from litellm.proxy.utils import ProxyUpdateSpend + ### UPDATE ORG TABLE ### + org_list_transactions: Final = db_spend_update_transactions["org_list_transactions"] + verbose_proxy_logger.debug("Org Spend transactions: %s", org_list_transactions) + if org_list_transactions is not None and len(org_list_transactions.keys()) > 0: + for i in range(n_retry_times + 1): + start_time = time.time() + try: + async with _spend_update_tx( + prisma_client, "LiteLLM_OrganizationTable", outer_transaction=outer_transaction + ) as transaction: + async with transaction.batch_() as batcher: + # Sort by org_id for consistent lock ordering across pods to prevent deadlocks. + for org_id, response_cost in sorted(org_list_transactions.items()): + batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists + where={"organization_id": org_id}, + data={"spend": {"increment": response_cost}}, + ) + break + except Exception as e: + if outer_transaction is not None: + raise + await self.handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + if on_table_committed is not None: + on_table_committed("org_list_transactions") + + org_member_list_transactions: Final = db_spend_update_transactions.get("org_member_list_transactions") + verbose_proxy_logger.debug("Org Membership Spend transactions: %s", org_member_list_transactions) + if org_member_list_transactions is not None and len(org_member_list_transactions.keys()) > 0: + for i in range(n_retry_times + 1): + start_time = time.time() + try: + async with ( + _spend_update_tx( + prisma_client, "LiteLLM_OrganizationMembership", outer_transaction=outer_transaction + ) as transaction, + transaction.batch_() as batcher, + ): + for (org_id, member_id), response_cost in _org_members_in_key_order( + org_member_list_transactions + ): + batcher.litellm_organizationmembership.update_many( + where={"organization_id": org_id, "user_id": member_id}, + data={"spend": {"increment": response_cost}}, + ) + break + except Exception as e: + if outer_transaction is not None: + raise + await self.handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + if on_table_committed is not None: + on_table_committed("org_member_list_transactions") + + ### UPDATE TEAM TABLE ### + team_list_transactions: Final = db_spend_update_transactions["team_list_transactions"] + verbose_proxy_logger.debug("Team Spend transactions: %s", team_list_transactions) + if team_list_transactions is not None and len(team_list_transactions.keys()) > 0: + for i in range(n_retry_times + 1): + start_time = time.time() + try: + async with _spend_update_tx( + prisma_client, "LiteLLM_TeamTable", outer_transaction=outer_transaction + ) as transaction: + async with transaction.batch_() as batcher: + # Sort by team_id for consistent lock ordering across pods to prevent deadlocks. + for team_id, response_cost in sorted(team_list_transactions.items()): + verbose_proxy_logger.debug( + "Updating spend for team id=%s by %s", team_id, response_cost + ) + batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists + where={"team_id": team_id}, + data={"spend": {"increment": response_cost}}, + ) + break + except Exception as e: + if outer_transaction is not None: + raise + await self.handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) + if on_table_committed is not None: + on_table_committed("team_list_transactions") + ### UPDATE USER TABLE ### user_list_transactions: Final = db_spend_update_transactions["user_list_transactions"] verbose_proxy_logger.debug("User Spend transactions: %s", user_list_transactions) @@ -2359,75 +2483,6 @@ class DBSpendUpdateWriter: ) if on_table_committed is not None: on_table_committed("end_user_list_transactions") - ### UPDATE KEY TABLE ### - key_list_transactions: Final = db_spend_update_transactions["key_list_transactions"] - verbose_proxy_logger.debug("KEY Spend transactions: %s", key_list_transactions) - if key_list_transactions is not None and len(key_list_transactions.keys()) > 0: - for i in range(n_retry_times + 1): - start_time = time.time() - try: - async with _spend_update_tx( - prisma_client, "LiteLLM_VerificationToken", outer_transaction=outer_transaction - ) as transaction: - async with transaction.batch_() as batcher: - # Sort by token for consistent lock ordering across pods to prevent deadlocks. - for token, response_cost in sorted(key_list_transactions.items()): - spend_increment: _SpendIncrement = {"increment": response_cost} - batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists - where={"token": token}, - data={ - "spend": spend_increment, - "total_spend": spend_increment, - "last_active": datetime.now(timezone.utc), - }, - ) - break - except Exception as e: - if outer_transaction is not None: - raise - await self.handle_spend_update_failure( - e=e, - attempt=i, - n_retry_times=n_retry_times, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - if on_table_committed is not None: - on_table_committed("key_list_transactions") - - ### UPDATE TEAM TABLE ### - team_list_transactions: Final = db_spend_update_transactions["team_list_transactions"] - verbose_proxy_logger.debug("Team Spend transactions: %s", team_list_transactions) - if team_list_transactions is not None and len(team_list_transactions.keys()) > 0: - for i in range(n_retry_times + 1): - start_time = time.time() - try: - async with _spend_update_tx( - prisma_client, "LiteLLM_TeamTable", outer_transaction=outer_transaction - ) as transaction: - async with transaction.batch_() as batcher: - # Sort by team_id for consistent lock ordering across pods to prevent deadlocks. - for team_id, response_cost in sorted(team_list_transactions.items()): - verbose_proxy_logger.debug( - "Updating spend for team id=%s by %s", team_id, response_cost - ) - batcher.litellm_teamtable.update_many( # 'update_many' prevents error from being raised if no row exists - where={"team_id": team_id}, - data={"spend": {"increment": response_cost}}, - ) - break - except Exception as e: - if outer_transaction is not None: - raise - await self.handle_spend_update_failure( - e=e, - attempt=i, - n_retry_times=n_retry_times, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - if on_table_committed is not None: - on_table_committed("team_list_transactions") ### UPDATE TEAM Membership TABLE with spend ### team_member_list_transactions: Final = db_spend_update_transactions["team_member_list_transactions"] @@ -2478,22 +2533,27 @@ class DBSpendUpdateWriter: elif on_table_committed is not None: on_table_committed("team_member_list_transactions") - ### UPDATE ORG TABLE ### - org_list_transactions: Final = db_spend_update_transactions["org_list_transactions"] - verbose_proxy_logger.debug("Org Spend transactions: %s", org_list_transactions) - if org_list_transactions is not None and len(org_list_transactions.keys()) > 0: + ### UPDATE KEY TABLE ### + key_list_transactions: Final = db_spend_update_transactions["key_list_transactions"] + verbose_proxy_logger.debug("KEY Spend transactions: %s", key_list_transactions) + if key_list_transactions is not None and len(key_list_transactions.keys()) > 0: for i in range(n_retry_times + 1): start_time = time.time() try: async with _spend_update_tx( - prisma_client, "LiteLLM_OrganizationTable", outer_transaction=outer_transaction + prisma_client, "LiteLLM_VerificationToken", outer_transaction=outer_transaction ) as transaction: async with transaction.batch_() as batcher: - # Sort by org_id for consistent lock ordering across pods to prevent deadlocks. - for org_id, response_cost in sorted(org_list_transactions.items()): - batcher.litellm_organizationtable.update_many( # 'update_many' prevents error from being raised if no row exists - where={"organization_id": org_id}, - data={"spend": {"increment": response_cost}}, + # Sort by token for consistent lock ordering across pods to prevent deadlocks. + for token, response_cost in sorted(key_list_transactions.items()): + spend_increment: _SpendIncrement = {"increment": response_cost} + batcher.litellm_verificationtoken.update_many( # 'update_many' prevents error from being raised if no row exists + where={"token": token}, + data={ + "spend": spend_increment, + "total_spend": spend_increment, + "last_active": datetime.now(timezone.utc), + }, ) break except Exception as e: @@ -2507,39 +2567,7 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) if on_table_committed is not None: - on_table_committed("org_list_transactions") - - org_member_list_transactions: Final = db_spend_update_transactions.get("org_member_list_transactions") - verbose_proxy_logger.debug("Org Membership Spend transactions: %s", org_member_list_transactions) - if org_member_list_transactions is not None and len(org_member_list_transactions.keys()) > 0: - for i in range(n_retry_times + 1): - start_time = time.time() - try: - async with ( - _spend_update_tx( - prisma_client, "LiteLLM_OrganizationMembership", outer_transaction=outer_transaction - ) as transaction, - transaction.batch_() as batcher, - ): - for key, response_cost in sorted(org_member_list_transactions.items()): - _, quoted_org_id, _, quoted_user_id = key.split("::") - batcher.litellm_organizationmembership.update_many( - where={"organization_id": unquote(quoted_org_id), "user_id": unquote(quoted_user_id)}, - data={"spend": {"increment": response_cost}}, - ) - break - except Exception as e: - if outer_transaction is not None: - raise - await self.handle_spend_update_failure( - e=e, - attempt=i, - n_retry_times=n_retry_times, - start_time=start_time, - proxy_logging_obj=proxy_logging_obj, - ) - if on_table_committed is not None: - on_table_committed("org_member_list_transactions") + on_table_committed("key_list_transactions") ### UPDATE PROJECT TABLE ### project_list_transactions: Final = db_spend_update_transactions.get("project_list_transactions") diff --git a/litellm/proxy/deps/repositories.py b/litellm/proxy/deps/repositories.py index 8f929119dee..323cf57dfec 100644 --- a/litellm/proxy/deps/repositories.py +++ b/litellm/proxy/deps/repositories.py @@ -5,11 +5,13 @@ from fastapi import Depends from litellm.models.deployment_warning import DeploymentWarningIgnoreStore from litellm.models.gateway_activity import GatewayActivityStore from litellm.models.plugin_upload import PluginUploadStore +from litellm.models.skill_submission import SkillSubmissionStore from litellm.models.user_banner import UserBannerStore from litellm.proxy.deps.database import DatabaseDep from litellm.repositories.deployment_warning_ignore_repository import DeploymentWarningIgnoreRepository from litellm.repositories.gateway_activity_repository import GatewayActivityRepository from litellm.repositories.plugin_upload_repository import PluginUploadRepository +from litellm.repositories.skill_submission_repository import SkillSubmissionRepository from litellm.repositories.user_banner_repository import UserBannerRepository @@ -43,3 +45,12 @@ async def get_plugin_upload_repository(database: DatabaseDep) -> PluginUploadSto PluginUploadRepositoryDep: TypeAlias = Annotated[PluginUploadStore | None, Depends(get_plugin_upload_repository)] + + +async def get_skill_submission_repository(database: DatabaseDep) -> SkillSubmissionStore | None: + return None if database is None else SkillSubmissionRepository(database.prisma) + + +SkillSubmissionRepositoryDep: TypeAlias = Annotated[ + SkillSubmissionStore | None, Depends(get_skill_submission_repository) +] diff --git a/litellm/proxy/discovery_endpoints/agent_skills_archive.py b/litellm/proxy/discovery_endpoints/agent_skills_archive.py index dd78896725f..e45ca7e10f2 100644 --- a/litellm/proxy/discovery_endpoints/agent_skills_archive.py +++ b/litellm/proxy/discovery_endpoints/agent_skills_archive.py @@ -11,26 +11,63 @@ import io import lzma import re import stat +import struct import zipfile import zlib -from collections.abc import Mapping +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Final +from typing import ClassVar, Final import yaml -from pydantic import TypeAdapter +from pydantic import ConfigDict, TypeAdapter, ValidationError MAX_ARCHIVE_UNPACKED_BYTES: Final = 50 * 1024 * 1024 MAX_ARCHIVE_ENTRIES: Final = 1000 +MAX_SKILL_FILE_BYTES: Final = 10 * 1024 * 1024 +MAX_SKILL_MANIFEST_BYTES: Final = 1024 * 1024 +SKILL_ARCHIVE_ERRORS: Final = ( + zipfile.BadZipFile, + OSError, + RuntimeError, + EOFError, + lzma.LZMAError, + zlib.error, + struct.error, + UnicodeError, +) SKILL_MANIFEST_FILENAME: Final = "SKILL.md" _ZIP_ENTRY_TIMESTAMP: Final = (1980, 1, 1, 0, 0, 0) _ZIP_ENTRY_MODE: Final = 0o644 _FRONTMATTER_PATTERN: Final = re.compile(r"^---\s*\n(.*?)\n---\s*(?:\n|$)", re.DOTALL) +_AGENT_SKILL_FRONTMATTER_PATTERN: Final = re.compile(r"\A---[ \t]*\r?\n(.*?)\r?\n---[ \t]*(?:\r?\n|\Z)", re.DOTALL) +_AGENT_SKILL_YAML_SCALARS: Final = MappingProxyType( + { + "tag:yaml.org,2002:null": re.compile(r"^(?:null|Null|NULL|~|)$"), + "tag:yaml.org,2002:bool": re.compile(r"^(?:true|True|TRUE|false|False|FALSE)$"), + "tag:yaml.org,2002:int": re.compile(r"^(?:[-+]?[0-9]+|0o[0-7]+|0x[0-9a-fA-F]+)$"), + "tag:yaml.org,2002:float": re.compile( + r"^(?:[-+]?(?:\.[0-9]+|[0-9]+(?:\.[0-9]*)?)(?:[eE][-+]?[0-9]+)?" + r"|[-+]?\.(?:inf|Inf|INF)|\.(?:nan|NaN|NAN))$" + ), + } +) +_MAX_FRONTMATTER_DEPTH: Final = 32 _WINDOWS_DRIVE_PATTERN: Final = re.compile(r"^[A-Za-z]:") _EMPTY_FRONTMATTER: Final[Mapping[str, object]] = MappingProxyType({}) _DECODED_YAML: Final = TypeAdapter(object) +_YAML_EVENTS: Final = TypeAdapter(Iterable[object]) +_YAML_NODES: Final = TypeAdapter(tuple[yaml.Node, ...], config=ConfigDict(arbitrary_types_allowed=True)) +_YAML_PAIRS: Final = TypeAdapter( + tuple[tuple[yaml.Node, yaml.Node], ...], config=ConfigDict(arbitrary_types_allowed=True) +) + + +class _AgentSkillFrontmatterLoader(yaml.BaseLoader): + yaml_implicit_resolvers: ClassVar[Mapping[str | None, Sequence[tuple[str, re.Pattern[str]]]]] = { # pyright: ignore[reportIncompatibleVariableOverride] # PyYAML resolver tables are class state + None: list(_AGENT_SKILL_YAML_SCALARS.items()) + } @dataclass(frozen=True, slots=True) @@ -39,37 +76,100 @@ class SkillArchive: digest: str declared_name: str | None declared_description: str | None + manifest_content: bytes def read_skill_archive_members( - stored_content: bytes, *, for_upload: bool = False + stored_content: bytes, *, for_upload: bool = False, skill_path: str | None = None ) -> tuple[tuple[str, bytes, int], ...] | None: try: with zipfile.ZipFile(io.BytesIO(stored_content)) as uploaded: - return _flattened_members(uploaded, for_upload=for_upload) - except (zipfile.BadZipFile, OSError, RuntimeError, EOFError, lzma.LZMAError, zlib.error): + return _flattened_members(uploaded, for_upload=for_upload, skill_path=skill_path) + except SKILL_ARCHIVE_ERRORS: return None +def read_skill_archive_entry(uploaded: zipfile.ZipFile, info: zipfile.ZipInfo) -> bytes: + with uploaded.open(info) as member: + return member.read(info.file_size + 1) + + def build_skill_archive(stored_content: bytes) -> SkillArchive | None: """Return the installable archive for an upload, or None when it holds no root SKILL.md.""" members: Final = read_skill_archive_members(stored_content) if members is None: return None - frontmatter: Final = _manifest_frontmatter( - next(data for name, data, _ in members if name == SKILL_MANIFEST_FILENAME) - ) + manifest_content: Final = next(data for name, data, _ in members if name == SKILL_MANIFEST_FILENAME) + frontmatter: Final = _manifest_frontmatter(manifest_content) content: Final = pack_skill_archive(members) return SkillArchive( content=content, digest=f"sha256:{hashlib.sha256(content).hexdigest()}", declared_name=_frontmatter_text(frontmatter, "name"), declared_description=_frontmatter_text(frontmatter, "description"), + manifest_content=manifest_content, ) -def _flattened_members(uploaded: zipfile.ZipFile, *, for_upload: bool) -> tuple[tuple[str, bytes, int], ...] | None: +def build_agent_skill_archive( + content: bytes, *, skill_path: str | None = None, plugin_name: str | None = None +) -> SkillArchive | None: + members: Final = read_skill_archive_members(content, for_upload=True, skill_path=skill_path) + selected: Final = ( + members + if members is not None or plugin_name is None + else read_skill_archive_members(content, for_upload=True, skill_path=f"skills/{plugin_name}") + ) + if selected is None: + return None + manifest_content: Final = next(data for path, data, _ in selected if path == SKILL_MANIFEST_FILENAME) + manifest: Final = agent_skill_frontmatter(manifest_content) + name: Final = _frontmatter_text(manifest, "name") + description: Final = _frontmatter_text(manifest, "description") + if ( + name is None + or not re.fullmatch(r"[a-z0-9]+(?:-[a-z0-9]+)*", name) + or len(name) > 64 + or description is None + or len(description) > 1024 + ): + return None + archive: Final = pack_skill_archive(selected) + return SkillArchive( + content=archive, + digest=f"sha256:{hashlib.sha256(archive).hexdigest()}", + declared_name=name, + declared_description=description, + manifest_content=manifest_content, + ) + + +def _flattened_members( + uploaded: zipfile.ZipFile, *, for_upload: bool, skill_path: str | None +) -> tuple[tuple[str, bytes, int], ...] | None: + kept: Final = skill_archive_entries(uploaded, for_upload=for_upload) + if kept is None: + return None + prefix: Final = f"{skill_path}/" if skill_path is not None else _common_root_prefix(tuple(path for _, path in kept)) + flattened: Final = tuple((info, path[len(prefix) :]) for info, path in kept if path.startswith(prefix)) + names: Final = frozenset(name for _, name in flattened) + if ( + SKILL_MANIFEST_FILENAME not in names + or len(names) != len(flattened) + or (for_upload and any(not safe_skill_archive_entry(info) for info, _ in flattened)) + ): + return None + + return tuple( + (name, read_skill_archive_entry(uploaded, info), _entry_mode(info, preserve_executable_bits=for_upload)) + for info, name in sorted(flattened, key=lambda member: member[1]) + ) + + +def skill_archive_entries( + uploaded: zipfile.ZipFile, *, for_upload: bool +) -> tuple[tuple[zipfile.ZipInfo, str], ...] | None: infos: Final = tuple(info for info in uploaded.infolist() if for_upload or not info.is_dir()) if not infos or len(infos) > MAX_ARCHIVE_ENTRIES: return None @@ -81,26 +181,13 @@ def _flattened_members(uploaded: zipfile.ZipFile, *, for_upload: bool) -> tuple[ ) if any(path is None for _, path in normalized): return None - if for_upload and ( - len(frozenset(path for _, path in normalized)) != len(normalized) - or any(not _safe_upload_entry(info) for info in infos) - ): - return None - - kept: Final = tuple( - (info, path) - for info, path in normalized - if path is not None and not info.is_dir() and (not for_upload or not _is_macos_metadata(path)) - ) - prefix: Final = _common_root_prefix(tuple(path for _, path in kept)) - flattened: Final = tuple((info, path[len(prefix) :]) for info, path in kept) - names: Final = frozenset(name for _, name in flattened) - if SKILL_MANIFEST_FILENAME not in names or len(names) != len(flattened): + if for_upload and len(frozenset(path for _, path in normalized)) != len(normalized): return None return tuple( - (name, uploaded.read(info), _entry_mode(info, preserve_executable_bits=for_upload)) - for info, name in sorted(flattened, key=lambda member: member[1]) + (info, path) + for info, path in normalized + if path is not None and not info.is_dir() and (not for_upload or not _is_macos_metadata(path)) ) @@ -113,9 +200,11 @@ def _entry_mode(info: zipfile.ZipInfo, *, preserve_executable_bits: bool) -> int return _ZIP_ENTRY_MODE | (mode & 0o111) -def _safe_upload_entry(info: zipfile.ZipInfo) -> bool: - return not info.flag_bits & 1 and ( - info.create_system != 3 or stat.S_IFMT(info.external_attr >> 16) in (0, stat.S_IFREG, stat.S_IFDIR) +def safe_skill_archive_entry(info: zipfile.ZipInfo) -> bool: + return ( + info.file_size <= MAX_SKILL_FILE_BYTES + and not info.flag_bits & 1 + and (info.create_system != 3 or stat.S_IFMT(info.external_attr >> 16) in (0, stat.S_IFREG, stat.S_IFDIR)) ) @@ -154,6 +243,65 @@ def _manifest_frontmatter(manifest: bytes) -> Mapping[str, object]: return parsed +def agent_skill_frontmatter(manifest: bytes) -> Mapping[str, str]: + if len(manifest) > MAX_SKILL_MANIFEST_BYTES: + return MappingProxyType({}) + try: + match: Final = _AGENT_SKILL_FRONTMATTER_PATTERN.match(manifest.decode("utf-8-sig")) + if match is None: + return MappingProxyType({}) + frontmatter: Final = match.group(1) + events: Final = _YAML_EVENTS.validate_python(yaml.parse(frontmatter, Loader=yaml.BaseLoader)) # pyright: ignore[reportUnknownMemberType] # PyYAML parse stubs omit their event iterator return type + if any(isinstance(event, yaml.AliasEvent) for event in events): + return MappingProxyType({}) + node: Final = _DECODED_YAML.validate_python(yaml.compose(frontmatter, Loader=_AgentSkillFrontmatterLoader)) # pyright: ignore[reportUnknownMemberType] # PyYAML compose stubs omit their return type + except (yaml.YAMLError, UnicodeError, RecursionError, ValidationError): + return MappingProxyType({}) + if not isinstance(node, yaml.MappingNode) or not _valid_frontmatter_node(node): + return MappingProxyType({}) + pairs: Final = _YAML_PAIRS.validate_python(node.value) + fields: Final = { + key_text: value_text + for key, value in pairs + if (key_text := _frontmatter_string(key)) is not None and (value_text := _frontmatter_string(value)) is not None + } + if "name" not in fields or "description" not in fields: + return MappingProxyType({}) + return MappingProxyType(fields) + + +def _frontmatter_string(node: yaml.Node) -> str | None: + if not isinstance(node, yaml.ScalarNode) or node.tag != "tag:yaml.org,2002:str": + return None + value: Final = _DECODED_YAML.validate_python(node.value) + return value if isinstance(value, str) else None + + +def _valid_frontmatter_node(node: yaml.Node, depth: int = 0) -> bool: + if depth > _MAX_FRONTMATTER_DEPTH: + return False + if isinstance(node, yaml.ScalarNode): + pattern: Final = _AGENT_SKILL_YAML_SCALARS.get(node.tag) + scalar: Final = _DECODED_YAML.validate_python(node.value) + return isinstance(scalar, str) and ( + node.tag == "tag:yaml.org,2002:str" or (pattern is not None and pattern.fullmatch(scalar) is not None) + ) + if isinstance(node, yaml.SequenceNode): + values: Final = _YAML_NODES.validate_python(node.value) + return node.tag == "tag:yaml.org,2002:seq" and all( + _valid_frontmatter_node(value, depth + 1) for value in values + ) + if not isinstance(node, yaml.MappingNode) or node.tag != "tag:yaml.org,2002:map": + return False + pairs: Final = _YAML_PAIRS.validate_python(node.value) + keys: Final = tuple(_frontmatter_string(key) for key, _ in pairs) + return ( + None not in keys + and len(frozenset(keys)) == len(keys) + and all(_valid_frontmatter_node(value, depth + 1) for _, value in pairs) + ) + + def _frontmatter_text(frontmatter: Mapping[str, object], key: str) -> str | None: value: Final = frontmatter.get(key) if not isinstance(value, str): diff --git a/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py b/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py index 830f7b0334e..3b589b1a4f0 100644 --- a/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py +++ b/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py @@ -9,7 +9,7 @@ is enabled, which publishes every stored skill to anyone who can reach the proxy import asyncio import re from collections.abc import Sequence -from itertools import groupby +from itertools import count from types import MappingProxyType from typing import Final @@ -19,6 +19,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.in_memory_cache import InMemoryCache from litellm.models.skills import LiteLLM_SkillsTable +from litellm.proxy.deps.repositories import PluginUploadRepositoryDep from litellm.proxy.discovery_endpoints.agent_skills_archive import SkillArchive, build_skill_archive from litellm.types.proxy.discovery_endpoints.agent_skills_endpoints import ( MAX_SKILL_DESCRIPTION_LENGTH, @@ -83,28 +84,34 @@ async def stored_skill(skill_id: str) -> LiteLLM_SkillsTable | None: ) async def agent_skills_index( request: Request, + uploads: PluginUploadRepositoryDep, skills: Sequence[LiteLLM_SkillsTable] = Depends(stored_skills), ) -> AgentSkillsIndex: - """Agent Skills v0.2.0 discovery index over every skill stored on this proxy.""" + """Discover public Hub skills and the separately enabled hosted skill store.""" + from litellm.proxy.discovery_endpoints.skill_hub_endpoints import public_hub_entries from litellm.proxy.utils import get_custom_url - installable: Final = await _installable(skills) + installable: Final = await _installable(skills[:MAX_INDEXED_SKILLS]) names: Final = _deduplicated(tuple(_base_name(skill, archive) for skill, archive in installable)) - - return AgentSkillsIndex( - skills=tuple( - AgentSkillsIndexEntry( - name=name, - type="archive", - description=_description(skill, archive, name), - url=get_custom_url( - request_base_url=str(request.base_url), - route=f"v1/skills/{skill.skill_id}/archive", - ), - digest=archive.digest, - ) - for (skill, archive), name in zip(installable, names, strict=True) + legacy: Final = tuple( + AgentSkillsIndexEntry( + name=name, + type="archive", + description=_description(skill, archive, name), + url=get_custom_url( + request_base_url=str(request.base_url), + route=f"v1/skills/{skill.skill_id}/archive", + ), + digest=archive.digest, ) + for (skill, archive), name in zip(installable, names, strict=True) + ) + remaining: Final = MAX_INDEXED_SKILLS - len(legacy) + hub: Final = (await public_hub_entries(request, uploads))[:remaining] if remaining else () + hub_names: Final = _deduplicated(tuple(entry.name for entry in hub), reserved=frozenset(names)) + return AgentSkillsIndex( + skills=legacy + + tuple(entry.model_copy(update={"name": name}) for entry, name in zip(hub, hub_names, strict=True)) ) @@ -153,7 +160,7 @@ async def _archive_for(skill: LiteLLM_SkillsTable) -> SkillArchive | None: ) return None - if cache_key is not None and len(archive.content) <= MAX_CACHED_ARCHIVE_BYTES: + if cache_key is not None and len(archive.content) + len(archive.manifest_content) <= MAX_CACHED_ARCHIVE_BYTES: _ARCHIVE_CACHE.set_cache(cache_key, archive) return archive @@ -172,19 +179,25 @@ def _slugify(raw: str | None) -> str | None: return _NON_SLUG_PATTERN.sub("-", raw.lower()).strip("-")[:MAX_SKILL_NAME_LENGTH].rstrip("-") or None -def _entry_name(entry: tuple[int, str]) -> str: - return entry[1] +def _deduplicated(names: Sequence[str], reserved: frozenset[str] = frozenset()) -> tuple[str, ...]: + natural_names: Final = frozenset(names) | reserved + ordinals: Final = MappingProxyType({name: count(1) for name in frozenset(names)}) + used: Final[set[str]] = set(reserved) # mutable-ok: bounded local allocation avoids rebuilding the set + def allocate(name: str) -> str: + selected: Final = ( + name + if name not in used + else next( + candidate + for candidate in (_with_ordinal(name, ordinal) for ordinal in ordinals[name]) + if candidate not in natural_names and candidate not in used + ) + ) + used.add(selected) + return selected -def _deduplicated(names: Sequence[str]) -> tuple[str, ...]: - ordinals: Final = MappingProxyType( - { - position: ordinal - for _, duplicates in groupby(sorted(enumerate(names), key=_entry_name), key=_entry_name) - for ordinal, (position, _) in enumerate(duplicates) - } - ) - return tuple(_with_ordinal(name, ordinals[position]) for position, name in enumerate(names)) + return tuple(allocate(name) for name in names) def _with_ordinal(name: str, ordinal: int) -> str: diff --git a/litellm/proxy/discovery_endpoints/skill_hub_endpoints.py b/litellm/proxy/discovery_endpoints/skill_hub_endpoints.py new file mode 100644 index 00000000000..93d1ace435d --- /dev/null +++ b/litellm/proxy/discovery_endpoints/skill_hub_endpoints.py @@ -0,0 +1,682 @@ +import asyncio +import hashlib +from collections.abc import Mapping +from dataclasses import dataclass, replace +from datetime import datetime +from typing import Annotated, Final +from urllib.parse import quote + +import httpx +from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, Response, UploadFile + +import litellm +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared HTTP factory has untyped optional params +) +from litellm.models.plugin_upload import ( + PluginArchiveMetadata, + PluginDownloader, + PluginDownloadsResponse, + PluginPublication, + PluginSnapshot, + PluginUploadStore, + PluginVersion, + PublishedPlugin, + has_agent_skill_metadata, +) +from litellm.models.skill_submission import SkillSubmission, SkillSubmissionStore +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import ( + MANIFEST_FIELDS, + delete_plugin, + disable_plugin, + enable_plugin, + get_plugin_downloads, + header_key_caller, + list_plugins, + plugin_archive_source_matches, + plugin_archive_url, + plugin_downloader, + require_download_recorded, +) +from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import resolve_skill_visibility +from litellm.proxy.anthropic_endpoints.claude_code_endpoints.skill_upload_archive import build_uploaded_skill_archive +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.claude_code_marketplace_auth import ( + marketplace_auth_required, + marketplace_key_header_present, +) +from litellm.proxy.common_utils.resource_ownership import is_proxy_admin +from litellm.proxy.deps.repositories import PluginUploadRepositoryDep, SkillSubmissionRepositoryDep +from litellm.proxy.discovery_endpoints.agent_skills_archive import ( + SkillArchive, + build_agent_skill_archive, +) +from litellm.proxy.discovery_endpoints.skill_hub_mutations import ( + SkillMutationError, + archive_writer, + write_skill, + write_target, +) +from litellm.proxy.discovery_endpoints.skill_import import ( + MAX_IMPORT_BYTES, + SkillImportError, + fetch_skill_source, + parse_skill_source, + select_imported_skill, +) +from litellm.types.proxy.claude_code_endpoints import PluginListItem +from litellm.types.proxy.discovery_endpoints.agent_skills_endpoints import AgentSkillsIndexEntry +from litellm.types.proxy.skill_hub import ( + ImportSkillRequest, + SkillHubList, + SkillHubMetadata, + SkillHubMutationResponse, + SkillHubPendingSkill, + SkillHubProposal, + SkillHubPublication, + SkillHubSkill, + SkillHubVersion, + SkillHubVersions, +) + +router: Final = APIRouter(prefix="/skill-hub/skills", tags=["Skills Hub"]) +_EDITABLE_METADATA: Final = frozenset(SkillHubMetadata.model_fields) +router.add_api_route( + "/{plugin_name}/downloads", get_plugin_downloads, methods=["GET"], response_model=PluginDownloadsResponse +) + + +@dataclass(frozen=True, slots=True) +class VisibleSkillArchive: + archive: SkillArchive + metadata: PluginArchiveMetadata + downloader: PluginDownloader + + +def _url(request: Request, route: str) -> str: + from litellm.proxy.utils import get_custom_url + + return get_custom_url(request_base_url=str(request.base_url), route=route.lstrip("/")) + + +def _download_url(request: Request, name: str) -> str: + return _url(request, f"skill-hub/skills/{quote(name, safe='')}/archive") + + +def _skill_response(request: Request, plugin: PluginListItem, download_count: int | None = None) -> SkillHubSkill: + local_archive: Final = plugin_archive_source_matches( + request, + plugin.name, + plugin.source, + plugin.source.get("sha256", ""), + has_agent_skill_metadata=plugin.has_agent_skill_metadata, + ) + return SkillHubSkill( + id=plugin.id, + name=plugin.name, + description=plugin.description, + version=plugin.version, + domain=plugin.domain, + namespace=plugin.namespace, + keywords=None if plugin.keywords is None else tuple(plugin.keywords), + updated_at=None if plugin.updated_at is None else datetime.fromisoformat(plugin.updated_at), + visibility="public" if plugin.enabled else "private", + source={"source": "archive", "url": _download_url(request, plugin.name)} if local_archive else plugin.source, + download_url=_download_url(request, plugin.name) if local_archive else None, + content_url=_url(request, f"skill-hub/skills/{quote(plugin.name, safe='')}/content") if local_archive else None, + download_count=download_count, + ) + + +@router.get("", response_model=SkillHubList) +async def list_skills( + request: Request, + uploads: PluginUploadRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + query: str = "", + offset: Annotated[int, Query(ge=0)] = 0, + limit: Annotated[int, Query(ge=1, le=100)] = 50, +) -> SkillHubList: + """Browse the caller's public and granted private skills.""" + catalog: Final = await list_plugins(enabled_only=False, user_api_key_dict=caller) + needle: Final = query.casefold() + matched: Final = tuple( + plugin + for plugin in sorted(catalog.plugins, key=lambda item: item.name) + if needle in " ".join((plugin.name, plugin.description or "", *(plugin.keywords or ()))).casefold() + ) + page: Final = matched[offset : offset + limit] + counts: Final[Mapping[str, int]] = ( + await _require_store(uploads).download_counts(tuple(plugin.id for plugin in page)) + if is_proxy_admin(caller) + else {} + ) + return SkillHubList( + skills=tuple(_skill_response(request, plugin, counts.get(plugin.id)) for plugin in page), + count=len(matched), + next_offset=offset + limit if offset + limit < len(matched) else None, + ) + + +@router.get("/{name}", response_model=SkillHubSkill) +async def get_skill( + request: Request, + name: str, + uploads: PluginUploadRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> SkillHubSkill: + """Read one visible skill and its portable download URL.""" + store: Final = _require_store(uploads) + metadata: Final = await store.get_metadata(name) + if metadata is None or not (await resolve_skill_visibility(caller)).allows(metadata): + raise HTTPException(404, "Skill not found") + return skill_metadata_response(request, metadata) + + +def skill_metadata_response(request: Request, metadata: PluginArchiveMetadata) -> SkillHubSkill: + manifest: Final = MANIFEST_FIELDS.validate_json(metadata.manifest_json or "{}") + return _skill_response( + request, + PluginListItem.model_validate( + { + "id": metadata.id or metadata.name, + "name": metadata.name, + "description": metadata.description, + "version": metadata.version, + "source": dict(metadata.source), + "has_agent_skill_metadata": metadata.has_agent_skill_metadata, + "enabled": metadata.enabled, + "created_at": None, + "updated_at": metadata.updated_at.isoformat() if metadata.updated_at is not None else None, + "domain": manifest.get("domain"), + "namespace": manifest.get("namespace"), + "keywords": manifest.get("keywords"), + } + ), + ) + + +def _require_store(uploads: PluginUploadStore | None) -> PluginUploadStore: + if uploads is None: + raise HTTPException(503, "Skills Hub requires a connected database") + return uploads + + +def _require_admin(caller: UserAPIKeyAuth) -> None: + if not is_proxy_admin(caller): + raise HTTPException(403, "Only proxy admins may publish skills") + + +@router.put("/{name}", response_model=SkillHubMutationResponse) +async def replace_skill_metadata( + request: Request, + name: str, + body: SkillHubMetadata, + uploads: PluginUploadRepositoryDep, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> SkillHubMutationResponse: + """Replace registry metadata, preserving files. Members submit an immutable proposal for admin review.""" + existing: Final = await write_target(name, True, uploads, caller, body.expected_updated_at) + if isinstance(existing, SkillMutationError): + raise HTTPException(existing.status, existing.message) + if existing is None: + raise HTTPException(404, "Skill not found") + publication: Final = PluginPublication( + name=name, + version=body.version, + description=body.description, + manifest={ + **{key: value for key, value in existing.manifest.items() if key not in _EDITABLE_METADATA}, + **body.model_dump(exclude_none=True), + "name": name, + }, + created_by=caller.user_id, + ) + return await _mutate(request, publication, existing, _require_store(uploads), submissions, caller) + + +@router.post("/{name}/publish", response_model=SkillHubSkill) +async def publish_skill( + request: Request, + name: str, + uploads: PluginUploadRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + _body: SkillHubPublication | None = None, +) -> SkillHubSkill: + """Make a skill public using the existing Hub publication policy.""" + _require_admin(caller) + await enable_plugin(plugin_name=name, user_api_key_dict=caller) + return await get_skill(request, name, uploads, caller) + + +@router.post("/{name}/unpublish", response_model=SkillHubSkill) +async def unpublish_skill( + request: Request, + name: str, + uploads: PluginUploadRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + _body: SkillHubPublication | None = None, +) -> SkillHubSkill: + """Remove a skill from public discovery while retaining explicit grants.""" + _require_admin(caller) + await disable_plugin(plugin_name=name, user_api_key_dict=caller) + return await get_skill(request, name, uploads, caller) + + +@router.delete("/{name}", status_code=204, response_class=Response) +async def remove_skill(name: str, caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)]) -> Response: + """Delete a skill and its managed archive.""" + _require_admin(caller) + await delete_plugin(plugin_name=name, user_api_key_dict=caller) + return Response(status_code=204) + + +@router.post("/upload", response_model=SkillHubMutationResponse) +async def upload_skill( + request: Request, + uploads: PluginUploadRepositoryDep, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(archive_writer)], + file: Annotated[UploadFile, File(description="Agent Skills ZIP with SKILL.md and optional supporting files")], + version: Annotated[str | None, Form()] = None, +) -> SkillHubMutationResponse: + """Publish a portable skill while retaining Claude marketplace compatibility.""" + store: Final = _require_store(uploads) + content: Final = await file.read(MAX_IMPORT_BYTES + 1) + if len(content) > MAX_IMPORT_BYTES: + raise HTTPException(413, "Skill ZIP must be 10 MiB or smaller") + archive: Final = await asyncio.to_thread(build_agent_skill_archive, content) + if archive is None: + raise HTTPException(422, "Upload a valid Agent Skill with name and description in SKILL.md") + return await _publish(request, store, submissions, caller, archive, version, None) + + +def import_http_client() -> httpx.AsyncClient: + return get_async_httpx_client(llm_provider="skill_hub_import").client + + +@router.post("/import", response_model=SkillHubMutationResponse) +async def import_skill( + request: Request, + body: ImportSkillRequest, + uploads: PluginUploadRepositoryDep, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(archive_writer)], + client: Annotated[httpx.AsyncClient, Depends(import_http_client)], +) -> SkillHubMutationResponse: + """Import a public skills.sh or GitHub skill as a portable snapshot.""" + store: Final = _require_store(uploads) + source: Final = parse_skill_source(body.url, skill=body.skill, ref=body.ref, path=body.path) + if isinstance(source, SkillImportError): + raise HTTPException(source.status, source.message) + content: Final = await fetch_skill_source(source, client) + if isinstance(content, SkillImportError): + raise HTTPException(content.status, content.message) + archive: Final = await asyncio.to_thread(select_imported_skill, content, source) + if isinstance(archive, SkillImportError): + raise HTTPException(archive.status, archive.message) + return await _publish(request, store, submissions, caller, archive, body.version, body.url) + + +async def _publish( + request: Request, + store: PluginUploadStore, + submissions: SkillSubmissionStore | None, + caller: UserAPIKeyAuth, + archive: SkillArchive, + version: str | None, + origin: str | None, +) -> SkillHubMutationResponse: + existing: Final = await write_target(archive.declared_name or "", False, store, caller) + if isinstance(existing, SkillMutationError): + raise HTTPException(existing.status, existing.message) + publication: Final = await prepare_skill_publication(request, archive, version, origin, caller.user_id) + return await _mutate(request, publication, existing, store, submissions, caller) + + +async def _mutate( + request: Request, + publication: PluginPublication, + existing: PluginSnapshot | None, + store: PluginUploadStore, + submissions: SkillSubmissionStore | None, + caller: UserAPIKeyAuth, +) -> SkillHubMutationResponse: + result: Final = await write_skill(publication, existing, store, submissions, caller) + if isinstance(result, SkillMutationError): + raise HTTPException(result.status, result.message) + if isinstance(result, SkillSubmission): + return SkillHubPendingSkill(submission_id=result.id, proposal=skill_proposal_response(request, result)) + return skill_publication_response(request, publication, result) + + +def skill_proposal_response(request: Request, submission: SkillSubmission) -> SkillHubProposal: + proposal: Final = submission.proposal + return SkillHubProposal.model_validate( + { + **proposal.manifest, + "name": proposal.name, + "description": proposal.description, + "version": proposal.version, + **( + { + "source": { + "source": "archive", + "url": _url(request, f"skill-hub/submissions/{quote(submission.id, safe='')}/archive"), + } + } + if submission.has_archive + else {} + ), + } + ) + + +@router.put("/{name}/upload", response_model=SkillHubMutationResponse) +async def replace_skill_archive( + request: Request, + name: str, + uploads: PluginUploadRepositoryDep, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(archive_writer)], + file: Annotated[UploadFile, File(description="Replacement Agent Skills ZIP with the same declared skill name")], + version: Annotated[str | None, Form()] = None, + expected_updated_at: Annotated[datetime | None, Form()] = None, +) -> SkillHubMutationResponse: + """Replace skill files, retaining registry fields. Members propose an update against the original editor timestamp.""" + existing: Final = await write_target(name, True, uploads, caller, expected_updated_at) + if isinstance(existing, SkillMutationError): + raise HTTPException(existing.status, existing.message) + if existing is None: + raise HTTPException(404, "Skill not found") + content: Final = await file.read(MAX_IMPORT_BYTES + 1) + if len(content) > MAX_IMPORT_BYTES: + raise HTTPException(413, "Skill ZIP must be 10 MiB or smaller") + archive: Final = await asyncio.to_thread(build_agent_skill_archive, content) + if archive is None or archive.declared_name != name: + raise HTTPException(422, "Replacement SKILL.md must declare the existing skill name and a description") + prepared: Final = await prepare_skill_publication(request, archive, version, None, caller.user_id) + publication: Final = replace(prepared, manifest={**existing.manifest, **prepared.manifest}) + return await _mutate(request, publication, existing, _require_store(uploads), submissions, caller) + + +def skill_publication_response( + request: Request, publication: PluginPublication, published: PublishedPlugin +) -> SkillHubSkill: + return _skill_response( + request, + PluginListItem.model_validate( + { + "source": publication.manifest.get("source", {}), + "id": published.id, + "name": published.name, + "description": publication.description, + "version": published.version, + "enabled": published.enabled, + "has_agent_skill_metadata": has_agent_skill_metadata(publication.manifest), + "created_at": None, + "updated_at": published.updated_at.isoformat(), + "domain": publication.manifest.get("domain"), + "namespace": publication.manifest.get("namespace"), + "keywords": publication.manifest.get("keywords"), + } + ), + ) + + +async def prepare_skill_publication( + request: Request, + archive: SkillArchive, + version: str | None, + origin: str | None, + created_by: str | None, +) -> PluginPublication: + name: Final = archive.declared_name + if name is None: + raise HTTPException(422, "SKILL.md must declare a skill name") + plugin: Final = await asyncio.to_thread(build_uploaded_skill_archive, archive.content, name) + if plugin is None: + raise HTTPException(422, "Skill ZIP exceeds the supported archive limits") + portable: Final = await asyncio.to_thread(build_agent_skill_archive, plugin, plugin_name=name) + if portable is None or portable.digest != archive.digest: + raise HTTPException(422, "Skill ZIP exceeds the supported archive limits after compatibility packaging") + digest: Final = hashlib.sha256(plugin).hexdigest() + source: Final = { + "source": "archive", + "url": plugin_archive_url(request, name, digest), + "sha256": digest, + } + manifest: Final = { + "name": name, + "description": archive.declared_description, + "version": version, + "source": source, + "agent_skill": { + "name": name, + "description": archive.declared_description, + "digest": archive.digest, + "origin": origin, + }, + } + return PluginPublication( + name=name, + version=version, + description=archive.declared_description, + manifest=manifest, + content=plugin, + created_by=created_by, + ) + + +async def _version_metadata(name: str, store: PluginUploadStore, caller: UserAPIKeyAuth) -> PluginArchiveMetadata: + metadata: Final = await store.get_metadata(name) + if metadata is None or metadata.id is None or not (await resolve_skill_visibility(caller)).allows(metadata): + raise HTTPException(404, "Skill not found") + return metadata + + +def _version_response(request: Request, name: str, version: PluginVersion) -> SkillHubVersion: + return SkillHubVersion( + id=version.id, + version=version.version, + description=version.description, + manifest=version.manifest, + published_at=version.published_at, + published_by=version.published_by, + content_kind="stored_archive" if version.archive_available else "source_reference", + download_url=( + _url(request, f"skill-hub/skills/{quote(name, safe='')}/versions/{quote(version.id, safe='')}/archive") + if version.archive_available + else None + ), + ) + + +@router.get("/{name}/versions", response_model=SkillHubVersions) +async def list_skill_versions( + request: Request, + name: str, + uploads: PluginUploadRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + limit: Annotated[int, Query(ge=1, le=100)] = 20, + offset: Annotated[int, Query(ge=0)] = 0, +) -> SkillHubVersions: + """Read published metadata history. External references are not retained artifacts.""" + store: Final = _require_store(uploads) + metadata: Final = await _version_metadata(name, store, caller) + versions: Final = await store.versions(metadata, limit + 1, offset) + return SkillHubVersions( + versions=tuple(_version_response(request, name, version) for version in versions[:limit]), + next_offset=offset + limit if len(versions) > limit else None, + ) + + +@router.get("/{name}/versions/{version_id}", response_model=SkillHubVersion) +async def get_skill_version( + request: Request, + name: str, + version_id: str, + uploads: PluginUploadRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> SkillHubVersion: + """Read a retained publication under the current skill's visibility policy.""" + store: Final = _require_store(uploads) + metadata: Final = await _version_metadata(name, store, caller) + version: Final = await store.version(metadata, version_id) + if version is None: + raise HTTPException(404, "Skill version not found") + return _version_response(request, name, version) + + +@router.get( + "/{name}/versions/{version_id}/archive", + response_class=Response, + responses={200: {"content": {"application/zip": {"schema": {"type": "string", "format": "binary"}}}}}, +) +async def download_skill_version( + name: str, + version_id: str, + uploads: PluginUploadRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> Response: + """Download a retained publication as a portable ZIP with SKILL.md at its root.""" + store: Final = _require_store(uploads) + metadata: Final = await _version_metadata(name, store, caller) + version: Final = await store.version(metadata, version_id) + if version is None: + raise HTTPException(404, "Skill version not found") + content: Final = await store.version_content(metadata, version_id) + archive: Final = ( + await asyncio.to_thread(build_agent_skill_archive, content, plugin_name=name) if content is not None else None + ) + if archive is None: + raise HTTPException(404, "This version has no retained Agent Skills archive") + require_download_recorded( + await store.record_download( + replace(metadata, version=version.version), + archive.digest.removeprefix("sha256:"), + plugin_downloader(caller), + ) + ) + return Response( + archive.content, + media_type="application/zip", + headers={ + "Content-Disposition": f'attachment; filename="{name}.zip"', + "Cache-Control": "private, no-store", + "X-Content-Type-Options": "nosniff", + "ETag": f'"{archive.digest}"', + }, + ) + + +async def _visible_archive(request: Request, name: str, uploads: PluginUploadStore | None) -> VisibleSkillArchive: + anonymous_allowed: Final = litellm.public_skills_index is True and not marketplace_auth_required() + caller: Final = ( + await header_key_caller(request) if marketplace_key_header_present(request) or not anonymous_allowed else None + ) + store: Final = _require_store(uploads) + metadata: Final = await store.get_metadata(name) + if metadata is None or not (await resolve_skill_visibility(caller)).allows(metadata): + raise HTTPException(404, "Skill not found") + if not plugin_archive_source_matches( + request, + name, + metadata.source, + metadata.source.get("sha256", ""), + has_agent_skill_metadata=metadata.has_agent_skill_metadata, + ): + raise HTTPException(404, "This skill no longer refers to a managed archive") + content: Final = await store.get_content(metadata) + archive: Final = ( + await asyncio.to_thread(build_agent_skill_archive, content, plugin_name=name) if content is not None else None + ) + if archive is None: + raise HTTPException(404, "This skill has no stored Agent Skills archive") + return VisibleSkillArchive(archive, metadata, plugin_downloader(caller)) + + +@router.get( + "/{name}/archive", + response_class=Response, + responses={200: {"content": {"application/zip": {"schema": {"type": "string", "format": "binary"}}}}}, +) +async def download_skill( + request: Request, name: str, uploads: PluginUploadRepositoryDep, digest: str | None = None +) -> Response: + """Download the complete portable skill, with SKILL.md at the ZIP root.""" + visible: Final = await _visible_archive(request, name, uploads) + archive: Final = visible.archive + if digest is not None and digest != archive.digest: + raise HTTPException(404, "Skill archive version not found") + require_download_recorded( + await _require_store(uploads).record_download( + visible.metadata, archive.digest.removeprefix("sha256:"), visible.downloader + ) + ) + return Response( + archive.content, + media_type="application/zip", + headers={ + "Content-Disposition": f'attachment; filename="{name}.zip"', + "Cache-Control": "private, no-store", + "X-Content-Type-Options": "nosniff", + "ETag": f'"{archive.digest}"', + }, + ) + + +@router.get("/{name}/content", response_class=Response, responses={200: {"content": {"text/markdown": {}}}}) +async def skill_content(request: Request, name: str, uploads: PluginUploadRepositoryDep) -> Response: + """Read the original SKILL.md without downloading supporting files.""" + visible: Final = await _visible_archive(request, name, uploads) + return Response( + visible.archive.manifest_content, + media_type="text/markdown", + headers={"Cache-Control": "private, no-store", "X-Content-Type-Options": "nosniff"}, + ) + + +async def public_hub_entries(request: Request, store: PluginUploadStore | None) -> tuple[AgentSkillsIndexEntry, ...]: + if store is None or marketplace_auth_required(): + return () + entries: Final[list[AgentSkillsIndexEntry]] = [] # mutable-ok: bounded inventory avoids copying on each append + async for item in store.public_metadata(): + entry: AgentSkillsIndexEntry | None = _index_entry(request, item) + if entry is None: + continue + entries.append(entry) + if len(entries) == 1000: + break + return tuple(entries) + + +def _index_entry(request: Request, item: PluginArchiveMetadata) -> AgentSkillsIndexEntry | None: + manifest: Final = MANIFEST_FIELDS.validate_json(item.manifest_json or "{}") + raw: Final = manifest.get("agent_skill") + if ( + not isinstance(raw, dict) + or not item.enabled + or not plugin_archive_source_matches( + request, + item.name, + item.source, + item.source.get("sha256", ""), + has_agent_skill_metadata=item.has_agent_skill_metadata, + ) + ): + return None + skill: Final = MANIFEST_FIELDS.validate_python(raw) + digest: Final = skill.get("digest") + description: Final = item.description or skill.get("description") + if not isinstance(digest, str) or not isinstance(description, str): + return None + return AgentSkillsIndexEntry( + name=item.name, + type="archive", + description=description, + url=f"{_download_url(request, item.name)}?digest={quote(digest, safe=':')}", + digest=digest, + ) diff --git a/litellm/proxy/discovery_endpoints/skill_hub_mutations.py b/litellm/proxy/discovery_endpoints/skill_hub_mutations.py new file mode 100644 index 00000000000..63c519923c6 --- /dev/null +++ b/litellm/proxy/discovery_endpoints/skill_hub_mutations.py @@ -0,0 +1,172 @@ +import hashlib +from dataclasses import dataclass +from datetime import datetime +from typing import Annotated, Final + +from fastapi import Depends, HTTPException +from typing_extensions import assert_never + +from litellm.models.plugin_upload import ( + PluginPublication, + PluginPublicationBase, + PluginPublicationError, + PluginSnapshot, + PluginUploadStore, + PublishedPlugin, +) +from litellm.models.skill_submission import ( + SkillSubmission, + SkillSubmissionStore, + SubmissionConflict, + SubmissionLimitExceeded, + SubmissionMissing, + SubmissionReviewed, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import resolve_skill_visibility +from litellm.proxy.auth.auth_exception_handler import DB_UNAVAILABLE_FALLBACK_USER_ID +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.resource_ownership import is_proxy_admin +from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 +from litellm.repositories.plugin_publication import resolve_plugin_version +from litellm.types.proxy.auth.user_roles import LitellmUserRoles + + +@dataclass(frozen=True, slots=True) +class SkillMutationError: + status: int + message: str + + +def writer_error(caller: UserAPIKeyAuth) -> SkillMutationError | None: + if is_proxy_admin(caller): + return None + if ( + not caller.user_id + or caller.user_id == DB_UNAVAILABLE_FALLBACK_USER_ID + or caller.user_role in (None, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) + ): + return SkillMutationError(403, "Submitting a skill requires a writable user identity.") + return None + + +def submission_owner(caller: UserAPIKeyAuth) -> str | None | SkillMutationError: + if is_proxy_admin(caller) or caller.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY: + return None + if not caller.user_id or caller.user_id == DB_UNAVAILABLE_FALLBACK_USER_ID: + return SkillMutationError(403, "Reading submissions requires a user identity.") + return caller.user_id + + +async def archive_writer(caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)]) -> UserAPIKeyAuth: + from litellm.proxy.proxy_server import proxy_logging_obj + + error: Final = writer_error(caller) + if error is not None: + raise HTTPException(error.status, error.message) + limiter: Final = proxy_logging_obj.get_proxy_hook("parallel_request_limiter") + if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): + raise HTTPException(503, "Skills Hub publishing requires the default v3 rate limiter") + key_identity: Final = caller.api_key or caller.token + identity: Final = ( + f"user:{caller.user_id}" if caller.user_id else f"key:{key_identity}" if key_identity else "local-admin" + ) + result: Final = await limiter.atomic_check_and_increment_by_n( + descriptors=[ + { + "key": "skill_hub_publish", + "value": hashlib.sha256(identity.encode()).hexdigest(), + "rate_limit": {"requests_per_unit": 10, "window_size": 60}, + } + ], + increments=[{"requests": 1}], + parent_otel_span=caller.parent_otel_span, + ) + if result["overall_code"] == "OVER_LIMIT": + raise HTTPException( + 429, "Skills Hub allows 10 uploads or imports per minute per user", headers={"Retry-After": "60"} + ) + return caller + + +async def write_target( + name: str, + updating: bool, + uploads: PluginUploadStore | None, + caller: UserAPIKeyAuth, + expected_updated_at: datetime | None = None, +) -> PluginSnapshot | SkillMutationError | None: + error: Final = writer_error(caller) + if error is not None: + return error + if uploads is None: + return SkillMutationError(503, "Skills Hub requires a connected database") + existing: Final = await uploads.get_plugin(name) + if updating and (existing is None or not (await resolve_skill_visibility(caller)).allows(existing)): + return SkillMutationError(404, "Skill not found") + if existing is not None and expected_updated_at is not None and existing.updated_at != expected_updated_at: + return SkillMutationError(409, "The published skill changed. Reload it before submitting.") + if not updating and existing is not None: + return SkillMutationError(409, f"A skill named '{name}' already exists. Update the existing skill instead.") + return existing + + +def publication_error(error: PluginPublicationError) -> SkillMutationError: + return SkillMutationError({"invalid_version": 400, "conflict": 409, "not_found": 404}[error.kind], error.message) + + +async def write_skill( + publication: PluginPublication, + existing: PluginSnapshot | None, + uploads: PluginUploadStore, + submissions: SkillSubmissionStore | None, + caller: UserAPIKeyAuth, +) -> PublishedPlugin | SkillSubmission | SkillMutationError: + error: Final = writer_error(caller) + if error is not None: + return error + validation: Final = resolve_plugin_version( + publication.version, None if existing is None else existing.version, updating=existing is not None + ) + if isinstance(validation, PluginPublicationError): + return publication_error(validation) + expected: Final = None if existing is None else PluginPublicationBase(existing.id, existing.updated_at) + if is_proxy_admin(caller): + published: Final = await uploads.publish(publication, expected) + return publication_error(published) if isinstance(published, PluginPublicationError) else published + if submissions is None or not caller.user_id: + return SkillMutationError(503, "Skills Hub requires a connected database") + submission: Final = await submissions.submit(publication, caller.user_id, expected) + return ( + SkillMutationError(429, submission.message) if isinstance(submission, SubmissionLimitExceeded) else submission + ) + + +async def review_skill( + submission_id: str, + submissions: SkillSubmissionStore | None, + caller: UserAPIKeyAuth, + approve: bool, + reason: str | None, + publication_base_url: str | None = None, +) -> SubmissionReviewed | SkillMutationError: + if not is_proxy_admin(caller): + return SkillMutationError(403, "Only proxy admins may review skill submissions.") + if submissions is None: + return SkillMutationError(503, "Skills Hub requires a connected database") + result: Final = await submissions.review( + submission_id, caller.user_id or "proxy-admin", approve, reason, publication_base_url=publication_base_url + ) + match result: + case SubmissionMissing(): + return SkillMutationError(404, "Skill submission not found") + case SubmissionConflict(message=message): + return SkillMutationError(409, message) + case PluginPublicationError(kind="invalid_version") as error: + return publication_error(error) + case PluginPublicationError(message=message): + return SkillMutationError(409, message) + case SubmissionReviewed(): + return result + case _: + assert_never(result) diff --git a/litellm/proxy/discovery_endpoints/skill_hub_submissions.py b/litellm/proxy/discovery_endpoints/skill_hub_submissions.py new file mode 100644 index 00000000000..c1bb422b2a0 --- /dev/null +++ b/litellm/proxy/discovery_endpoints/skill_hub_submissions.py @@ -0,0 +1,192 @@ +import asyncio +from typing import Annotated, Final + +from fastapi import APIRouter, Depends, HTTPException, Query, Request, Response + +from litellm.models.plugin_upload import PluginPublication +from litellm.models.skill_submission import SkillSubmission, SkillSubmissionStore, SubmissionStatus +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_skill_access import resolve_skill_visibility +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.deps.repositories import PluginUploadRepositoryDep, SkillSubmissionRepositoryDep +from litellm.proxy.discovery_endpoints.agent_skills_archive import build_agent_skill_archive +from litellm.proxy.discovery_endpoints.skill_hub_endpoints import ( + skill_metadata_response, + skill_proposal_response, + skill_publication_response, +) +from litellm.proxy.discovery_endpoints.skill_hub_mutations import SkillMutationError, review_skill, submission_owner +from litellm.types.proxy.skill_hub import ( + SkillHubRejection, + SkillHubReviewResponse, + SkillHubSubmission, + SkillHubSubmissionDetail, + SkillHubSubmissionList, +) + +router: Final = APIRouter(prefix="/skill-hub/submissions", tags=["Skills Hub"]) + + +def _store(submissions: SkillSubmissionStore | None) -> SkillSubmissionStore: + if submissions is None: + raise HTTPException(503, "Skills Hub requires a connected database") + return submissions + + +def _owner(caller: UserAPIKeyAuth) -> str | None: + owner: Final = submission_owner(caller) + if isinstance(owner, SkillMutationError): + raise HTTPException(owner.status, owner.message) + return owner + + +def _response(request: Request, submission: SkillSubmission) -> SkillHubSubmission: + return SkillHubSubmission( + id=submission.id, + skill_name=submission.plugin_name, + operation=submission.operation, + status=submission.status, + proposal=skill_proposal_response(request, submission), + has_archive=submission.has_archive, + submitted_by=submission.submitted_by, + created_at=submission.created_at, + base_skill_id=submission.base_plugin_id, + base_updated_at=submission.base_updated_at, + reviewed_by=submission.reviewed_by, + reviewed_at=submission.reviewed_at, + rejection_reason=submission.rejection_reason, + ) + + +@router.get("", response_model=SkillHubSubmissionList) +async def list_submissions( + request: Request, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + status: SubmissionStatus | None = None, + limit: Annotated[int, Query(ge=1, le=100)] = 20, + offset: Annotated[int, Query(ge=0)] = 0, +) -> SkillHubSubmissionList: + """Read a page of owned proposals. Proxy admins and admin viewers can read every owner's proposals.""" + owner: Final = _owner(caller) + page: Final = await _store(submissions).list(owner, status, limit, offset) + return SkillHubSubmissionList( + submissions=tuple(_response(request, row) for row in page.submissions), + count=len(page.submissions), + has_more=page.has_more, + ) + + +@router.get("/{submission_id}", response_model=SkillHubSubmissionDetail) +async def get_submission( + request: Request, + submission_id: str, + submissions: SkillSubmissionRepositoryDep, + uploads: PluginUploadRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> SkillHubSubmissionDetail: + """Read an immutable proposal and its currently visible published skill.""" + owner: Final = _owner(caller) + submission: Final = await _store(submissions).get(submission_id, owner) + if submission is None: + raise HTTPException(404, "Skill submission not found") + if uploads is None: + raise HTTPException(503, "Skills Hub requires a connected database") + current: Final = await uploads.get_metadata(submission.plugin_name) + visible: Final = current is not None and (owner is None or (await resolve_skill_visibility(caller)).allows(current)) + return SkillHubSubmissionDetail( + submission=_response(request, submission), + current_skill=skill_metadata_response(request, current) if current is not None and visible else None, + ) + + +@router.get( + "/{submission_id}/archive", + response_class=Response, + responses={200: {"content": {"application/zip": {"schema": {"type": "string", "format": "binary"}}}}}, +) +async def download_submission( + submission_id: str, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> Response: + """Download an owned proposal as a private portable ZIP. Legacy nonportable archives return 404.""" + owner: Final = _owner(caller) + store: Final = _store(submissions) + submission: Final = await store.get(submission_id, owner) + if submission is None: + raise HTTPException(404, "Skill submission not found") + content: Final = await store.get_archive(submission_id, owner) + archive: Final = ( + await asyncio.to_thread(build_agent_skill_archive, content, plugin_name=submission.plugin_name) + if content is not None + else None + ) + if archive is None: + raise HTTPException(404, "This submission has no portable Agent Skills archive") + return Response( + archive.content, + media_type="application/zip", + headers={ + "Content-Disposition": f'attachment; filename="{submission.plugin_name}.zip"', + "Cache-Control": "private, no-store", + "X-Content-Type-Options": "nosniff", + "ETag": f'"{archive.digest}"', + }, + ) + + +async def _review( + request: Request, + submission_id: str, + submissions: SkillSubmissionStore | None, + caller: UserAPIKeyAuth, + approve: bool, + reason: str | None, +) -> SkillHubReviewResponse: + from litellm.proxy.utils import get_custom_url + + result: Final = await review_skill( + submission_id, submissions, caller, approve, reason, get_custom_url(str(request.base_url)) + ) + if isinstance(result, SkillMutationError): + raise HTTPException(result.status, result.message) + proposal: Final = result.submission.proposal + publication: Final = PluginPublication( + name=proposal.name, + version=proposal.version, + description=proposal.description, + manifest={ + **proposal.manifest, + **({"source": dict(result.published_source)} if result.published_source else {}), + }, + created_by=result.submission.submitted_by, + ) + return SkillHubReviewResponse( + action="approved" if approve else "rejected", + submission=_response(request, result.submission), + skill=None if result.plugin is None else skill_publication_response(request, publication, result.plugin), + ) + + +@router.put("/{submission_id}/approve", response_model=SkillHubReviewResponse) +async def approve_submission( + request: Request, + submission_id: str, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> SkillHubReviewResponse: + """Publish a pending proposal and record its approval atomically. Only proxy admins may review.""" + return await _review(request, submission_id, submissions, caller, True, None) + + +@router.put("/{submission_id}/reject", response_model=SkillHubReviewResponse) +async def reject_submission( + request: Request, + submission_id: str, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + body: SkillHubRejection | None = None, +) -> SkillHubReviewResponse: + """Reject a pending proposal without changing the published skill. Only proxy admins may review.""" + return await _review(request, submission_id, submissions, caller, False, None if body is None else body.reason) diff --git a/litellm/proxy/discovery_endpoints/skill_import.py b/litellm/proxy/discovery_endpoints/skill_import.py new file mode 100644 index 00000000000..a74e3abef10 --- /dev/null +++ b/litellm/proxy/discovery_endpoints/skill_import.py @@ -0,0 +1,126 @@ +import io +import re +import zipfile +from dataclasses import dataclass +from typing import Final +from urllib.parse import quote, unquote, urlsplit + +import httpx + +from litellm.proxy.discovery_endpoints.agent_skills_archive import ( + MAX_ARCHIVE_ENTRIES, + MAX_ARCHIVE_UNPACKED_BYTES, + MAX_SKILL_MANIFEST_BYTES, + SKILL_ARCHIVE_ERRORS, + SkillArchive, + agent_skill_frontmatter, + build_agent_skill_archive, + read_skill_archive_entry, + safe_skill_archive_entry, + skill_archive_entries, +) + +MAX_IMPORT_BYTES: Final = 10 * 1024 * 1024 + + +@dataclass(frozen=True, slots=True) +class SkillImportError: + status: int + message: str + + +@dataclass(frozen=True, slots=True) +class GitHubSkillSource: + owner: str + repository: str + ref: str + skill: str | None + path: str | None + + @property + def archive_url(self) -> str: + return f"https://codeload.github.com/{self.owner}/{self.repository}/zip/{quote(self.ref, safe='')}" + + +def parse_skill_source( + url: str, *, skill: str | None, ref: str, path: str | None +) -> GitHubSkillSource | SkillImportError: + try: + parsed: Final = urlsplit(url) + except ValueError: + return SkillImportError(400, "Invalid skill source URL") + if parsed.scheme != "https" or parsed.netloc not in ("skills.sh", "github.com") or parsed.query or parsed.fragment: + return SkillImportError(400, "Use an HTTPS skills.sh skill URL or a public github.com repository URL") + parts: Final = tuple(unquote(part) for part in parsed.path.strip("/").split("/")) + if len(parts) < 2 or any(re.fullmatch(r"[A-Za-z0-9_.-]+", part) is None for part in parts[:2]): + return SkillImportError(400, "The source URL must identify a GitHub owner and repository") + if parsed.netloc == "skills.sh": + if len(parts) != 3 or (skill is not None and skill != parts[2]): + return SkillImportError(400, "Use skills.sh/// with a matching skill selection") + return GitHubSkillSource(parts[0], parts[1], ref, parts[2], path) + if len(parts) == 2: + return GitHubSkillSource(parts[0], parts[1].removesuffix(".git"), ref, skill, path) + if len(parts) < 5 or parts[2] != "tree" or path is not None or ref != "HEAD": + return SkillImportError(400, "Use a repository URL with ref and path, or a GitHub tree URL") + return GitHubSkillSource(parts[0], parts[1], parts[3], skill, "/".join(parts[4:])) + + +async def fetch_skill_source(source: GitHubSkillSource, client: httpx.AsyncClient) -> bytes | SkillImportError: + buffer: Final = io.BytesIO() + request: Final = httpx.Request( + "GET", + source.archive_url, + headers={"Accept": "application/zip"}, + extensions={"timeout": httpx.Timeout(30).as_dict()}, + ) + try: + response: Final = await client.send(request, stream=True, follow_redirects=False, auth=None) + try: + if response.status_code != 200: + return SkillImportError(422, "Unable to download the public GitHub repository at the requested ref") + async for chunk in response.aiter_bytes(): + if buffer.tell() + len(chunk) > MAX_IMPORT_BYTES: + return SkillImportError(413, "Repository ZIP must be 10 MiB or smaller") + buffer.write(chunk) + finally: + await response.aclose() + except httpx.HTTPError: + return SkillImportError(502, "The skill source could not be reached") + return buffer.getvalue() + + +def select_imported_skill(content: bytes, source: GitHubSkillSource) -> SkillArchive | SkillImportError: + try: + with zipfile.ZipFile(io.BytesIO(content)) as archive: + infos: Final = archive.infolist() + if len(infos) > MAX_ARCHIVE_ENTRIES or sum(info.file_size for info in infos) > MAX_ARCHIVE_UNPACKED_BYTES: + return SkillImportError(413, "Repository exceeds the skill archive limits") + entries: Final = skill_archive_entries(archive, for_upload=True) + if entries is None: + return SkillImportError(422, "Repository contains unsafe archive entries") + candidates: Final = tuple(info for info, path in entries if path.endswith("/SKILL.md")) + if any(not safe_skill_archive_entry(info) for info in candidates): + return SkillImportError(422, "Repository contains unsafe skill manifests") + selected: Final = tuple(info for info in candidates if _matches_selection(info, archive, source)) + if len(selected) != 1: + return SkillImportError(422, "Select exactly one skill using its name or repository-relative path") + root: Final = selected[0].filename.rsplit("/", 1)[0] + skill: Final = build_agent_skill_archive(content, skill_path=root) + return ( + skill + if skill is not None + else SkillImportError(422, "The selected directory is not a valid Agent Skill") + ) + except SKILL_ARCHIVE_ERRORS: + return SkillImportError(422, "The repository did not return a readable skill ZIP") + + +def _matches_selection(info: zipfile.ZipInfo, archive: zipfile.ZipFile, source: GitHubSkillSource) -> bool: + if info.file_size > MAX_SKILL_MANIFEST_BYTES: + return False + relative: Final = info.filename.split("/", 1)[1].rpartition("/")[0] + if source.path is not None and relative != source.path.strip("/"): + return False + if source.skill is None: + return True + return agent_skill_frontmatter(read_skill_archive_entry(archive, info)).get("name") == source.skill diff --git a/litellm/proxy/discovery_endpoints/skill_scanning.py b/litellm/proxy/discovery_endpoints/skill_scanning.py new file mode 100644 index 00000000000..ee236482fbf --- /dev/null +++ b/litellm/proxy/discovery_endpoints/skill_scanning.py @@ -0,0 +1,186 @@ +import asyncio +import io +import json +from datetime import datetime, timedelta, timezone +from typing import Annotated, Final +from uuid import uuid4 + +import httpx +from fastapi import APIRouter, BackgroundTasks, Body, Depends, HTTPException +from pydantic import ValidationError + +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # shared HTTP factory has untyped optional params +) +from litellm.models.skill_submission import SkillScanInput, SkillSubmissionStore, SubmissionConflict, SubmissionMissing +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.resource_ownership import is_proxy_admin +from litellm.proxy.deps.repositories import SkillSubmissionRepositoryDep +from litellm.secret_managers.main import get_secret_str +from litellm.types.proxy.skill_scanning import ( + SkillScan, + SkillScannerConfig, + SkillScanReport, + SkillScanRequest, + SkillScanResponse, +) + +router: Final = APIRouter() +_MAX_REPORT_BYTES: Final = 1024 * 1024 + + +def scanner_config() -> SkillScannerConfig | str | None: + from litellm.proxy.proxy_server import general_settings_view + + value: Final = general_settings_view().get("skill_scanner") + if value is None: + return None + try: + return SkillScannerConfig.model_validate(value) + except ValidationError: + return "The skill scanner configuration is invalid" + + +async def scan_report( + client: httpx.AsyncClient, config: SkillScannerConfig, api_key: str, target: SkillScanInput +) -> SkillScanReport | str: + try: + async with client.stream( + "POST", + config.endpoint, + headers={"Authorization": f"Bearer {api_key}", "Accept-Encoding": "identity"}, + data={ + "source": json.dumps(dict(target.source)), + "use_llm": str(config.use_llm).lower(), + "expires_at": target.scan.expires_at.isoformat(), + }, + files=None if target.content is None else {"file": ("skill.zip", target.content, "application/zip")}, + timeout=httpx.Timeout(config.timeout_seconds, connect=10), + follow_redirects=False, + ) as response: + if response.status_code != 200: + return f"Scanner returned HTTP {response.status_code}. Check the scanner service and authentication." + output: Final = io.BytesIO() + async for chunk in response.aiter_bytes(chunk_size=65536): + if output.tell() + len(chunk) > _MAX_REPORT_BYTES: + return "Scanner report exceeded the 1 MiB limit." + output.write(chunk) + report: Final = SkillScanReport.model_validate_json(output.getvalue()) + if config.use_llm and report.scan_mode != "static+llm": + return report.model_copy( + update={ + "complete": False, + "limitations": (*report.limitations[:99], "Requested semantic analysis did not run."), + } + ) + return report + except (httpx.HTTPError, ValidationError): + return "The scanner was unreachable or returned an invalid report." + + +async def complete_scan( + submissions: SkillSubmissionStore, + submission_id: str, + target: SkillScanInput, + config: SkillScannerConfig, + api_key: str, +) -> None: + try: + client: Final = get_async_httpx_client( + "skill_scanner", params={"ssl_verify": True, "follow_redirects": False} + ).client + report: Final = await asyncio.wait_for( + scan_report(client, config, api_key, target), + timeout=(target.scan.expires_at - datetime.now(timezone.utc)).total_seconds(), + ) + completed: Final = target.scan.model_copy( + update={ + "status": "completed" if isinstance(report, SkillScanReport) else "failed", + "report": report if isinstance(report, SkillScanReport) else None, + "error": report if isinstance(report, str) else None, + "completed_at": datetime.now(timezone.utc), + } + ) + await submissions.finish_scan(submission_id, completed) + except (TimeoutError, asyncio.CancelledError) as error: + await submissions.finish_scan( + submission_id, + target.scan.model_copy( + update={ + "status": "failed", + "error": "Scan timed out or was interrupted. Try scanning again.", + "completed_at": datetime.now(timezone.utc), + } + ), + ) + if isinstance(error, asyncio.CancelledError): + raise + + +@router.get( + "/skill-hub/submissions/{submission_id}/scan", + tags=["Skills Hub"], + response_model=SkillScanResponse, +) +async def get_skill_scan( + submission_id: str, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], +) -> SkillScanResponse: + """Read the latest manual security scan. Only proxy admins may read reports.""" + if not is_proxy_admin(caller): + raise HTTPException(403, "Proxy admin access is required to review security scans") + if submissions is None: + raise HTTPException(503, "Database is not connected") + scan: Final = await submissions.get_scan(submission_id) + if isinstance(scan, SubmissionMissing): + raise HTTPException(404, "Skill submission not found") + return SkillScanResponse( + configured=isinstance(scanner_config(), SkillScannerConfig), + scan=None if scan is None else scan.at(datetime.now(timezone.utc)), + ) + + +@router.post( + "/skill-hub/submissions/{submission_id}/scan", + tags=["Skills Hub"], + response_model=SkillScanResponse, +) +async def start_skill_scan( + submission_id: str, + background_tasks: BackgroundTasks, + submissions: SkillSubmissionRepositoryDep, + caller: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + _request: Annotated[SkillScanRequest | None, Body()] = None, +) -> SkillScanResponse: + """Start an optional scan of a pending submission. Scans never change approval decisions.""" + if not is_proxy_admin(caller): + raise HTTPException(403, "Proxy admin access is required to scan skills") + if submissions is None: + raise HTTPException(503, "Database is not connected") + config: Final = scanner_config() + if config is None: + raise HTTPException(503, "Security scanning is not configured") + if isinstance(config, str): + raise HTTPException(503, config) + api_key: Final = get_secret_str(config.api_key_env) + if not api_key: + raise HTTPException(503, "The scanner authentication secret is not configured") + now: Final = datetime.now(timezone.utc) + target: Final = await submissions.start_scan( + submission_id, + SkillScan( + id=str(uuid4()), + status="running", + requested_by=caller.user_id or "proxy-admin", + started_at=now, + expires_at=now + timedelta(seconds=config.timeout_seconds), + ), + ) + if isinstance(target, SubmissionMissing): + raise HTTPException(404, "Skill submission not found") + if isinstance(target, SubmissionConflict): + raise HTTPException(409, target.message) + background_tasks.add_task(complete_scan, submissions, submission_id, target, config, api_key) + return SkillScanResponse(configured=True, scan=target.scan) diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index 6558c82f80c..4cbc2155af0 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -347,47 +347,45 @@ async def _cache_access_group_record(record: _AccessGroupRecord) -> None: # --------------------------------------------------------------------------- -async def _sync_add_access_group_to_teams(tx: _AccessGroupTx, team_ids: list[str], access_group_id: str) -> None: - """Add access_group_id to each team's access_group_ids (idempotent).""" - for team_id in team_ids: +async def _sync_access_group_on_teams( + tx: _AccessGroupTx, add: Sequence[str], remove: Sequence[str], access_group_id: str +) -> None: + """Rows are locked in sorted team_id order, the order every spend writer locks team rows in, so this + cannot deadlock with them.""" + to_add: Final = frozenset(add) + for team_id in sorted(to_add | frozenset(remove)): team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id}) - if team is not None and access_group_id not in (team.access_group_ids or []): + if team is None: + continue + current = list(team.access_group_ids or []) + if team_id in to_add and access_group_id not in current: + await tx.litellm_teamtable.update( + where={"team_id": team_id}, data={"access_group_ids": [*current, access_group_id]} + ) + elif team_id not in to_add and access_group_id in current: await tx.litellm_teamtable.update( where={"team_id": team_id}, - data={"access_group_ids": list(team.access_group_ids or []) + [access_group_id]}, + data={"access_group_ids": list(_without_access_group(current, access_group_id))}, ) -async def _sync_remove_access_group_from_teams(tx: _AccessGroupTx, team_ids: list[str], access_group_id: str) -> None: - """Remove access_group_id from each team's access_group_ids (idempotent).""" - for team_id in team_ids: - team = await tx.litellm_teamtable.find_unique(where={"team_id": team_id}) - if team is not None and access_group_id in (team.access_group_ids or []): - await tx.litellm_teamtable.update( - where={"team_id": team_id}, - data={"access_group_ids": [ag for ag in (team.access_group_ids or ()) if ag != access_group_id]}, - ) - - -async def _sync_add_access_group_to_keys(tx: _AccessGroupTx, key_tokens: list[str], access_group_id: str) -> None: - """Add access_group_id to each key's access_group_ids (idempotent).""" - for token in key_tokens: +async def _sync_access_group_on_keys( + tx: _AccessGroupTx, add: Sequence[str], remove: Sequence[str], access_group_id: str +) -> None: + """Rows are locked in sorted token order, the order every spend writer locks key rows in.""" + to_add: Final = frozenset(add) + for token in sorted(to_add | frozenset(remove)): key = await tx.litellm_verificationtoken.find_unique(where={"token": token}) - if key is not None and access_group_id not in (key.access_group_ids or []): + if key is None: + continue + current = list(key.access_group_ids or []) + if token in to_add and access_group_id not in current: await tx.litellm_verificationtoken.update( - where={"token": token}, - data={"access_group_ids": list(key.access_group_ids or []) + [access_group_id]}, + where={"token": token}, data={"access_group_ids": [*current, access_group_id]} ) - - -async def _sync_remove_access_group_from_keys(tx: _AccessGroupTx, key_tokens: list[str], access_group_id: str) -> None: - """Remove access_group_id from each key's access_group_ids (idempotent).""" - for token in key_tokens: - key = await tx.litellm_verificationtoken.find_unique(where={"token": token}) - if key is not None and access_group_id in (key.access_group_ids or []): + elif token not in to_add and access_group_id in current: await tx.litellm_verificationtoken.update( - where={"token": token}, - data={"access_group_ids": [ag for ag in (key.access_group_ids or ()) if ag != access_group_id]}, + where={"token": token}, data={"access_group_ids": list(_without_access_group(current, access_group_id))} ) @@ -582,8 +580,8 @@ async def create_access_group( ) # Sync team and key tables to reference the new access group - await _sync_add_access_group_to_teams(tx, data.assigned_team_ids or [], record.access_group_id) - await _sync_add_access_group_to_keys(tx, data.assigned_key_ids or [], record.access_group_id) + await _sync_access_group_on_teams(tx, data.assigned_team_ids or [], (), record.access_group_id) + await _sync_access_group_on_keys(tx, data.assigned_key_ids or [], (), record.access_group_id) except HTTPException: raise except Exception as e: @@ -718,10 +716,8 @@ async def update_access_group( data=update_data, ) - await _sync_add_access_group_to_teams(tx, teams_to_add, access_group_id) - await _sync_remove_access_group_from_teams(tx, teams_to_remove, access_group_id) - await _sync_add_access_group_to_keys(tx, keys_to_add, access_group_id) - await _sync_remove_access_group_from_keys(tx, keys_to_remove, access_group_id) + await _sync_access_group_on_teams(tx, teams_to_add, teams_to_remove, access_group_id) + await _sync_access_group_on_keys(tx, keys_to_add, keys_to_remove, access_group_id) except HTTPException: raise except Exception as e: @@ -790,25 +786,8 @@ async def delete_access_group( ) affected_key_tokens = list(all_affected_key_tokens) - # Update teams returned by find_many directly — we already have their data. - for team in teams_with_group: - await tx.litellm_teamtable.update( - where={"team_id": team.team_id}, - data={"access_group_ids": [ag for ag in (team.access_group_ids or ()) if ag != access_group_id]}, - ) - # Use _sync_remove only for out-of-sync teams not found by the hasSome query. - out_of_sync_team_ids: Final = set(existing.assigned_team_ids or []) - {t.team_id for t in teams_with_group} - await _sync_remove_access_group_from_teams(tx, list(out_of_sync_team_ids), access_group_id) - - # Update keys returned by find_many directly — we already have their data. - for key in keys_with_group: - await tx.litellm_verificationtoken.update( - where={"token": key.token}, - data={"access_group_ids": [ag for ag in (key.access_group_ids or ()) if ag != access_group_id]}, - ) - # Use _sync_remove only for out-of-sync keys not found by the hasSome query. - out_of_sync_key_tokens: Final = set(existing.assigned_key_ids or []) - {k.token for k in keys_with_group} - await _sync_remove_access_group_from_keys(tx, list(out_of_sync_key_tokens), access_group_id) + await _sync_access_group_on_teams(tx, (), affected_team_ids, access_group_id) + await _sync_access_group_on_keys(tx, (), affected_key_tokens, access_group_id) detached_agent_ids: Final = await _detach_access_group_from_agents(tx, access_group_id) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 6e4ea0bd0cb..bbd6a1c9021 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -287,12 +287,12 @@ def compute_tag_metadata_totals(records: Sequence[DailySpendRecord]) -> SpendMet def _entity_metadata( - entity_metadata_field: Mapping[str, dict[str, object]] | None, + entity_metadata_field: Mapping[str, Mapping[str, object]] | None, entity_id: str, ) -> dict[str, object]: """The metadata payload for one entity breakdown bucket, empty when the caller passed none.""" stored: Final = entity_metadata_field.get(entity_id) if entity_metadata_field else None - return stored if stored is not None else {} + return dict(stored) if stored is not None else {} def provider_bucket_metadata(provider: str, base: Mapping[str, object]) -> Mapping[str, object]: @@ -1039,7 +1039,7 @@ def _fold_entity_rollups_sync( results: Sequence[DailySpendData], entity_rows: Sequence[EntityRollupRow], api_key_metadata: Mapping[str, KeyMetadataRow], - entity_metadata_field: Mapping[str, dict[str, object]] | None, # mutable-ok: shared field shape + entity_metadata_field: Mapping[str, Mapping[str, object]] | None, ) -> None: """Write breakdown.entities onto the already-built per-day results.""" by_date: Final = {day.date.strftime("%Y-%m-%d"): day for day in results} @@ -1113,6 +1113,19 @@ async def _guardrail_spend_or_empty( return () +async def _user_entity_metadata( + repository: DailyActivityRepository, entity_rows: Sequence[EntityRollupRow] +) -> Mapping[str, Mapping[str, object]]: + user_ids: Final = frozenset(row.entity_id for row in entity_rows if row.entity_id) + users: Final = await repository.user_metadata(user_ids) + return MappingProxyType( + { + user_id: MappingProxyType({"user_email": user.user_email, "user_alias": user.user_alias}) + for user_id, user in users.items() + } + ) + + async def get_daily_activity_aggregated( repository: DailyActivityRepository, scope: DailyActivityScope, @@ -1157,12 +1170,17 @@ async def get_daily_activity_aggregated( if entity_api_keys else MappingProxyType({}) ) + user_entity_metadata: Final[Mapping[str, Mapping[str, object]]] = ( + await _user_entity_metadata(repository, entity_records) + if scope.table is DailyActivityTable.USER + else MappingProxyType({}) + ) await asyncio.to_thread( _fold_entity_rollups_sync, results=aggregated["results"], entity_rows=entity_records, api_key_metadata=entity_key_metadata, - entity_metadata_field=entity_metadata_field, + entity_metadata_field=MappingProxyType({**user_entity_metadata, **(entity_metadata_field or {})}), ) results: Final = ( list(await asyncio.to_thread(_with_guardrail_breakdown, aggregated["results"], guardrail_rows)) diff --git a/litellm/proxy/management_endpoints/daily_activity_scopes.py b/litellm/proxy/management_endpoints/daily_activity_scopes.py index 7ff73ac6e2e..1eb6ab3210b 100644 --- a/litellm/proxy/management_endpoints/daily_activity_scopes.py +++ b/litellm/proxy/management_endpoints/daily_activity_scopes.py @@ -363,7 +363,7 @@ USER_RESOLVER = EntityScopeResolver( tags=("Budget & Spend Tracking", "Internal User management"), query=_user_query, resolve=_resolve_user, - include_entity_breakdown=False, + include_entity_breakdown=True, include_guardrail_breakdown=True, operation_names=MappingProxyType( { diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index d2aac9cae38..93a429a2fcc 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -83,6 +83,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: handle_update_object_permission_common, set_object_permission, ) +from litellm.proxy.management_helpers.team_roster_sync import user_tx_db from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.utils import handle_exception_on_proxy, hash_password from litellm.repositories.organization_repository import OrganizationRepository @@ -125,6 +126,11 @@ _USER_MODEL_BUDGET_ADAPTER: Final = TypeAdapter(dict[str, float | BudgetConfig]) _USER_CACHE_INVALIDATION_BATCH_SIZE: Final = 50 +_LOCK_ALL_USERS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_UserTable" WHERE TRUE ORDER BY user_id COLLATE "C" FOR UPDATE' +) + + @with_config(ConfigDict(extra="allow", strict=True)) class _AdminTeamAndKeyRows(TypedDict): keys: ReadOnly[Sequence[object] | None] @@ -2011,14 +2017,16 @@ async def bulk_user_update( try: # Perform bulk database update - await UserRepository(prisma_client).table.update_many( - where={}, - data=( - {**non_default_values, "model_max_budget": json.dumps(non_default_values["model_max_budget"])} - if "model_max_budget" in non_default_values - else non_default_values - ), - ) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ALL_USERS_SQL) + _ = await user_tx_db(transaction).update_many( + where={}, + data=( + {**non_default_values, "model_max_budget": json.dumps(non_default_values["model_max_budget"])} + if "model_max_budget" in non_default_values + else non_default_values + ), + ) for start in range(0, len(all_users_in_db), _USER_CACHE_INVALIDATION_BATCH_SIZE): await asyncio.gather( diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 8a1b9c112dc..4e511aee24a 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -1047,12 +1047,12 @@ async def delete_organization( deleted_orgs: Final = [] for organization_id in data.organization_ids: - # delete all teams in the organization - await _table(TeamRepository(prisma_client)).delete_many(where={"organization_id": organization_id}) - # delete all members in the organization - await _table(OrganizationMembershipRepository(prisma_client)).delete_many( - where={"organization_id": organization_id} - ) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ORG_TEAMS_SQL, organization_id) + _ = await transaction.litellm_teamtable.delete_many(where={"organization_id": organization_id}) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ORG_MEMBERS_SQL, organization_id) + _ = await transaction.litellm_organizationmembership.delete_many(where={"organization_id": organization_id}) await _delete_organization_keys( organization_id=organization_id, prisma_client=prisma_client, @@ -1074,6 +1074,19 @@ async def delete_organization( return deleted_orgs +_LOCK_ORG_TEAMS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TeamTable" WHERE organization_id = $1 ORDER BY team_id COLLATE "C" FOR UPDATE' +) +_LOCK_ORG_MEMBERS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_OrganizationMembership" WHERE organization_id = $1 ' + 'ORDER BY organization_id COLLATE "C", user_id COLLATE "C" FOR UPDATE' +) +_LOCK_ORG_KEYS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE organization_id = $1 ' + 'ORDER BY token COLLATE "C" FOR UPDATE' +) + + async def _delete_organization_keys( organization_id: str, prisma_client: PrismaClient, @@ -1087,7 +1100,9 @@ async def _delete_organization_keys( hashed_tokens=hashed_tokens_to_delete, prisma_client=prisma_client, ) - await _table(VerificationTokenRepository(prisma_client)).delete_many(where=key_filter) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ORG_KEYS_SQL, organization_id) + _ = await transaction.litellm_verificationtoken.delete_many(where={"organization_id": organization_id}) await delete_cache_key_objects( hashed_tokens=hashed_tokens_to_delete, user_api_key_cache=user_api_key_cache, diff --git a/litellm/proxy/management_helpers/access_group_model_sync.py b/litellm/proxy/management_helpers/access_group_model_sync.py index 890715bf9c2..20ea2ed0338 100644 --- a/litellm/proxy/management_helpers/access_group_model_sync.py +++ b/litellm/proxy/management_helpers/access_group_model_sync.py @@ -6,6 +6,7 @@ the arrays alone strands every group on a name nothing serves any more. """ from collections.abc import Sequence +from contextlib import AbstractAsyncContextManager from typing import Final, Protocol from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper @@ -53,6 +54,15 @@ _REMOVE_MODEL_NAME_SQL: Final = ( ) +class RawTransactions(Protocol): + def tx(self) -> AbstractAsyncContextManager[RawExecutor]: ... + + +def raw_transactions(prisma_client: object) -> RawTransactions: + db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client + return writer_wrapper(db) # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin + + def raw_executor(prisma_client: object) -> RawExecutor: db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client return writer_wrapper(db) # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin diff --git a/litellm/proxy/management_helpers/model_allowlist_rename_sync.py b/litellm/proxy/management_helpers/model_allowlist_rename_sync.py index f1a9faf6b2e..0e65e98c557 100644 --- a/litellm/proxy/management_helpers/model_allowlist_rename_sync.py +++ b/litellm/proxy/management_helpers/model_allowlist_rename_sync.py @@ -6,18 +6,27 @@ Those allowlists store public model names, not ids, so a deployment rename that alone denies the new name while the old entry grants a name nothing serves any more. """ -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final +from pydantic import ConfigDict + from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache -from litellm.proxy.management_helpers.access_group_model_sync import raw_executor, still_backed +from litellm.proxy.management_helpers.access_group_model_sync import raw_executor, raw_transactions, still_backed +from litellm.proxy.management_helpers.access_group_team_sync import TEAM_ADVISORY_LOCK_SQL from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel +class _ListingTeam(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + team_id: str + + class _TouchedRow(LiteLLMBaseModel): kind: str object_id: str @@ -80,11 +89,56 @@ def _rewrite_sql(set_clause: str, where_clause: str) -> str: return f"WITH {ctes} {rows}" +_TEAMS_LISTING_SQL: Final = ( + 'SELECT team_id FROM "LiteLLM_TeamTable" WHERE $1 = ANY("models") ORDER BY team_id COLLATE "C"' +) + +_LOCK_ORGS_LISTING_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_OrganizationTable" WHERE $1 = ANY("models") ' + 'ORDER BY organization_id COLLATE "C" FOR UPDATE' +) +_LOCK_TEAMS_LISTING_SQL: Final = ( + 'SELECT team_id FROM "LiteLLM_TeamTable" WHERE $1 = ANY("models") ORDER BY team_id COLLATE "C" FOR UPDATE' +) +_LOCK_USERS_LISTING_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_UserTable" WHERE $1 = ANY("models") ORDER BY user_id COLLATE "C" FOR UPDATE' +) +_LOCK_KEYS_LISTING_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE $1 = ANY("models") ORDER BY token COLLATE "C" FOR UPDATE' +) +_LOCK_PROJECTS_LISTING_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_ProjectTable" WHERE $1 = ANY("models") ORDER BY project_id COLLATE "C" FOR UPDATE' +) + _REPLACE_SQL: Final = _rewrite_sql('array_replace(array_remove("models", $2), $1, $2)', '$1 = ANY("models")') _APPEND_SQL: Final = _rewrite_sql('array_append("models", $2)', '$1 = ANY("models") AND NOT ($2 = ANY("models"))') +def _team_ids(rows: Sequence[object]) -> frozenset[str]: + return frozenset(_ListingTeam.model_validate(row).team_id for row in rows) + + +async def _rewrite_allowlists( + prisma_client: object, old_name: str, new_name: str, *, append: bool, attempts: int +) -> Sequence[object]: + """Team advisory locks come from a snapshot, so a team that starts listing the name before its row lock would + be locked without one. That attempt commits having written nothing, and the next one locks the fresh set.""" + async with raw_transactions(prisma_client).tx() as transaction: + listed: Final = _team_ids(await transaction.query_raw(_TEAMS_LISTING_SQL, old_name)) + for team_id in sorted(listed): + _ = await transaction.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id) + _ = await transaction.query_raw(_LOCK_ORGS_LISTING_SQL, old_name) + locked: Final = _team_ids(await transaction.query_raw(_LOCK_TEAMS_LISTING_SQL, old_name)) + retry: Final = attempts > 1 and not locked <= listed + if not retry: + _ = await transaction.query_raw(_LOCK_USERS_LISTING_SQL, old_name) + _ = await transaction.query_raw(_LOCK_KEYS_LISTING_SQL, old_name) + _ = await transaction.query_raw(_LOCK_PROJECTS_LISTING_SQL, old_name) + return await transaction.query_raw(_APPEND_SQL if append else _REPLACE_SQL, old_name, new_name) + return await _rewrite_allowlists(prisma_client, old_name, new_name, append=append, attempts=attempts - 1) + + async def sync_model_allowlists_for_renamed_model( prisma_client: object, *, @@ -98,8 +152,8 @@ async def sync_model_allowlists_for_renamed_model( return executor: Final = raw_executor(prisma_client) old_name_still_backed: Final = await still_backed(executor, llm_router, old_name, model_id) - touched_rows: Final = await executor.query_raw( - _APPEND_SQL if old_name_still_backed else _REPLACE_SQL, old_name, new_name + touched_rows: Final = await _rewrite_allowlists( + prisma_client, old_name, new_name, append=old_name_still_backed, attempts=3 ) touched: Final = tuple(_TouchedRow.model_validate(row) for row in touched_rows) await evict_and_broadcast( diff --git a/litellm/proxy/middleware/request_size_limit_middleware.py b/litellm/proxy/middleware/request_size_limit_middleware.py index f2883168d65..329a3c38d91 100644 --- a/litellm/proxy/middleware/request_size_limit_middleware.py +++ b/litellm/proxy/middleware/request_size_limit_middleware.py @@ -5,6 +5,7 @@ from typing import Final from fastapi import HTTPException from starlette._utils import get_route_path +from starlette.routing import compile_path from starlette.types import ASGIApp, Message, Receive, Scope, Send from litellm.proxy.common_utils.anthropic_error_payload import route_error_content @@ -25,7 +26,20 @@ def effective_max_request_size_mb( ) -> float | None: global_limit: Final = configured if configured is not None and configured > 0 else None method: Final = scope.get("method") - route_limit: Final = route_limits.get((method, get_route_path(scope))) if isinstance(method, str) else None + path: Final = get_route_path(scope) + exact_limit: Final = route_limits.get((method, path)) if isinstance(method, str) else None + route_limit: Final = ( + exact_limit + if exact_limit is not None + else next( + ( + limit + for (route_method, route), limit in route_limits.items() + if route_method == method and "{" in route and compile_path(route)[0].fullmatch(path) is not None + ), + None, + ) + ) if route_limit is None: return global_limit return min(route_limit, global_limit) if global_limit is not None else route_limit diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9a80761b477..c1af5d6f980 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2934,7 +2934,11 @@ def configured_max_request_size_mb() -> int | float | None: _ROUTE_MAX_REQUEST_SIZE_MB: Final[Mapping[tuple[str, str], float]] = MappingProxyType( - {("POST", "/claude-code/plugins/upload"): 11} + { + ("POST", "/claude-code/plugins/upload"): 11, + ("POST", "/skill-hub/skills/upload"): 11, + ("PUT", "/skill-hub/skills/{name}/upload"): 11, + } ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 78f16d9f1fd..818f2dc4858 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1616,10 +1616,69 @@ model LiteLLM_ClaudeCodePluginTable { created_by String? uploaded_archive LiteLLM_ClaudeCodePluginUpload? + versions LiteLLM_ClaudeCodePluginVersion[] + downloads LiteLLM_ClaudeCodePluginDownload[] + anonymous_download_count BigInt @default(0) @@map("LiteLLM_ClaudeCodePluginTable") } +model LiteLLM_ClaudeCodePluginVersion { + id String @id @default(uuid()) + plugin_id String + version String? + description String? + manifest_json String? + published_at DateTime? + published_by String? + archive_available Boolean @default(false) + content Bytes? + plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_id], references: [id], onDelete: Cascade, onUpdate: Cascade) + + @@unique([plugin_id, version]) + @@index([plugin_id, published_at]) +} + + +model LiteLLM_SkillSubmission { + id String @id @default(uuid()) + plugin_name String + operation String + status String @default("pending_review") + proposal_json String + archive_content Bytes? + submitted_by String + created_at DateTime @default(now()) + base_plugin_id String? + base_updated_at DateTime? + reviewed_by String? + reviewed_at DateTime? + rejection_reason String? + scan_json String? + + @@index([submitted_by, created_at(sort: Desc), id(sort: Desc)]) + @@index([submitted_by, status, created_at(sort: Desc), id(sort: Desc)]) + @@index([status, created_at(sort: Desc), id(sort: Desc)]) + @@index([created_at(sort: Desc), id(sort: Desc)]) +} + +model LiteLLM_ClaudeCodePluginDownload { + id BigInt @id @default(autoincrement()) + plugin_id String + downloaded_at DateTime @default(now()) + version String? + digest String + authenticated Boolean + user_id String? + user_email String? + key_alias String? + team_id String? + team_alias String? + plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_id], references: [id], onDelete: Cascade, onUpdate: Cascade) + + @@index([plugin_id, id]) +} + model LiteLLM_ClaudeCodePluginUpload { plugin_name String @id content Bytes diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index fcb9eb4a280..0348d8e1f82 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -77,9 +77,6 @@ from litellm.proxy.utils import handle_exception_on_proxy, query_raw_within_stat from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import SpendLogsRepository from litellm.repositories.team_repository import TeamRepository -from litellm.repositories.verification_token_repository import ( - VerificationTokenRepository, -) from litellm.types.proxy.spend_capture_rate import CaptureRateReport, SpendCaptureProvider if TYPE_CHECKING: @@ -292,24 +289,22 @@ class _TeamTable(Protocol): async def update_many(self, *, data: Mapping[str, float], where: Mapping[str, object]) -> int: ... -class _VerificationTokenTable(Protocol): - """The subset of the Prisma verification token table API this module uses.""" - - async def update_many(self, *, data: Mapping[str, float], where: Mapping[str, object]) -> int: ... - - def _spend_logs_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_SpendLogs"]: return SpendLogsRepository(prisma_client).table +_LOCK_ALL_TEAMS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TeamTable" WHERE TRUE ORDER BY team_id COLLATE "C" FOR UPDATE' +) +_LOCK_ALL_KEYS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE TRUE ORDER BY token COLLATE "C" FOR UPDATE' +) + + def _team_table(prisma_client: PrismaClient) -> _TeamTable: return TeamRepository(prisma_client).table -def _verification_token_table(prisma_client: PrismaClient) -> _VerificationTokenTable: - return VerificationTokenRepository(prisma_client).table - - def _spend_logs_daily_summary_sql( *, start_date_iso: str, @@ -3626,8 +3621,11 @@ async def global_spend_reset(): code=status.HTTP_401_UNAUTHORIZED, ) - await _verification_token_table(prisma_client).update_many(data={"spend": 0.0}, where={}) - await _team_table(prisma_client).update_many(data={"spend": 0.0}, where={}) + async with prisma_client.tx() as transaction: + _ = await transaction.query_raw(_LOCK_ALL_TEAMS_SQL) + _ = await transaction.litellm_teamtable.update_many(data={"spend": 0.0}, where={}) + _ = await transaction.query_raw(_LOCK_ALL_KEYS_SQL) + _ = await transaction.litellm_verificationtoken.update_many(data={"spend": 0.0}, where={}) return { "message": "Spend for all API Keys and Teams reset successfully", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 71ecae74058..00187790901 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -28,6 +28,7 @@ from dataclasses import dataclass, field from datetime import date, datetime, timedelta, timezone from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText +from email.utils import formatdate from functools import partial, reduce from itertools import takewhile from operator import attrgetter @@ -1952,6 +1953,7 @@ class ProxyLogging: self.daily_report_started: bool = False self.hanging_requests_check_started: bool = False self.deprecation_check_started: bool = False + self.deprecation_email_check_started: bool = False def startup_event( self, @@ -1995,6 +1997,20 @@ class ProxyLogging: self.hanging_requests_check_started = True self._ensure_deprecation_check_scheduled() + self._ensure_deprecation_email_check_scheduled() + + def _ensure_deprecation_email_check_scheduled(self) -> None: + """The email event opt-in controls reminders independently of Slack integration settings.""" + if self.deprecation_email_check_started: + return + try: + asyncio.get_running_loop() + except RuntimeError: + return + from litellm.integrations.model_deprecation_email import run_scheduled_model_deprecation_emails + + asyncio.create_task(run_scheduled_model_deprecation_emails()) + self.deprecation_email_check_started = True def _ensure_deprecation_check_scheduled(self) -> None: """Alerting can be configured at startup or by a later config reload, so schedule from either path""" @@ -2041,6 +2057,7 @@ class ProxyLogging: if updated_slack_alerting is True: self._ensure_deprecation_check_scheduled() + self._ensure_deprecation_email_check_scheduled() self.slack_alerting_instance.update_values( alerting=self.alerting, alerting_threshold=self.alerting_threshold, @@ -5528,6 +5545,20 @@ class _StaleReadEngine: return self.wrapper is current and self.generation == current.engine_generation +_LOCK_KEYS_BY_TOKEN_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE token = ANY($1::text[]) ' + 'ORDER BY token COLLATE "C" FOR UPDATE' +) +_LOCK_KEYS_BY_TEAM_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE team_id = ANY($1::text[]) ' + 'ORDER BY token COLLATE "C" FOR UPDATE' +) +_LOCK_TEAMS_BY_ID_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TeamTable" WHERE team_id = ANY($1::text[]) ' + 'ORDER BY team_id COLLATE "C" FOR UPDATE' +) + + class PrismaClient: spend_log_transactions: list = [] _spend_log_transactions_lock = asyncio.Lock() @@ -6828,24 +6859,26 @@ class PrismaClient: else: hashed_token = token hashed_tokens.append(hashed_token) - filter_query: dict[str, object] = {} - if user_id is not None: - filter_query = {"AND": [{"token": {"in": hashed_tokens}}, {"user_id": user_id}]} - else: - filter_query = {"token": {"in": hashed_tokens}} - - deleted_tokens: Final[int] = await VerificationTokenRepository(self).table.delete_many( - where=filter_query - ) + async with self.tx() as transaction: + _ = await transaction.query_raw(_LOCK_KEYS_BY_TOKEN_SQL, hashed_tokens) + deleted_tokens: Final[int] = await transaction.litellm_verificationtoken.delete_many( + where={"token": {"in": hashed_tokens}, "user_id": user_id} + if user_id is not None + else {"token": {"in": hashed_tokens}} + ) verbose_proxy_logger.debug("deleted_tokens: %s", deleted_tokens) return {"deleted_keys": deleted_tokens} elif table_name == "team" and team_id_list is not None and isinstance(team_id_list, list): # admin only endpoint -> `/team/delete` - await TeamRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}}) + async with self.tx() as transaction: + _ = await transaction.query_raw(_LOCK_TEAMS_BY_ID_SQL, team_id_list) + _ = await transaction.litellm_teamtable.delete_many(where={"team_id": {"in": team_id_list}}) return {"deleted_teams": team_id_list} elif table_name == "key" and team_id_list is not None and isinstance(team_id_list, list): # admin only endpoint -> `/team/delete` - await VerificationTokenRepository(self).table.delete_many(where={"team_id": {"in": team_id_list}}) + async with self.tx() as transaction: + _ = await transaction.query_raw(_LOCK_KEYS_BY_TEAM_SQL, team_id_list) + _ = await transaction.litellm_verificationtoken.delete_many(where={"team_id": {"in": team_id_list}}) except Exception as e: import traceback @@ -8263,7 +8296,9 @@ async def send_email( receiver_email: str | None = None, subject: str | None = None, html: str | None = None, -): + text: str | None = None, + message_id: str | None = None, +) -> bool: """ smtp_host, smtp_port, @@ -8289,16 +8324,21 @@ async def send_email( raise ValueError(f"No HTML body provided for SMTP email. {html}") ## EMAIL SETUP ## - email_message: Final = MIMEMultipart() + email_message: Final = MIMEMultipart("alternative" if text is not None else "mixed") email_message["From"] = sender_email email_message["To"] = receiver_email email_message["Subject"] = subject + if message_id is not None: + email_message["Message-ID"] = message_id + email_message["Date"] = formatdate(usegmt=True) verbose_proxy_logger.debug("sending email from %s to %s", sender_email, receiver_email) if smtp_host is None: raise ValueError("Trying to use SMTP, but SMTP_HOST is not set") # Attach the body to the email + if text is not None: + email_message.attach(MIMEText(text, "plain")) email_message.attach(MIMEText(html, "html")) try: @@ -8317,6 +8357,8 @@ async def send_email( except Exception as e: verbose_proxy_logger.exception("An error occurred while sending the email:" + str(e)) + return False + return True def hash_token(token: str): diff --git a/litellm/repositories/accounting_delivery_repository.py b/litellm/repositories/accounting_delivery_repository.py index 7dd92aecb2d..d5bbe891089 100644 --- a/litellm/repositories/accounting_delivery_repository.py +++ b/litellm/repositories/accounting_delivery_repository.py @@ -5,7 +5,7 @@ from itertools import chain from types import TracebackType from typing import TYPE_CHECKING, Final, Literal, Protocol -from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter from litellm.repositories.prisma_protocols import TableActions @@ -41,7 +41,10 @@ class _Delivery(BaseModel): delivery_id: str attempts: int - accounting_payload: JsonValue + + +class _ClaimedDelivery(_Delivery): + accounting_payload: str @dataclass(frozen=True, slots=True) @@ -68,9 +71,11 @@ class DeliveryMetrics(BaseModel): _DELIVERIES: Final = TypeAdapter(tuple[_Delivery, ...]) +_CLAIMED: Final = TypeAdapter(tuple[_ClaimedDelivery, ...]) _PAYLOAD: Final = TypeAdapter(dict[str, object]) _TIMEOUT: Final = timedelta(seconds=60) -_CLAIM: Final = """SELECT delivery_id, attempts, accounting_payload FROM "LiteLLM_SpendLogDelivery" +_CLAIM: Final = """SELECT delivery_id, attempts, accounting_payload::text AS accounting_payload +FROM "LiteLLM_SpendLogDelivery" WHERE state = 'pending' AND next_attempt_at <= CURRENT_TIMESTAMP AND (attempts > 0) = $2::boolean ORDER BY next_attempt_at, delivery_id LIMIT $1::integer FOR UPDATE SKIP LOCKED""" @@ -79,10 +84,20 @@ SET state = 'applied', applied_at = CURRENT_TIMESTAMP, last_error = NULL WHERE delivery_id = ANY($1::text[]) AND state = 'pending' """ _FAILED: Final = """UPDATE "LiteLLM_SpendLogDelivery" SET attempts = attempts + 1, last_error = $3, - next_attempt_at = CURRENT_TIMESTAMP + make_interval(secs => LEAST(300, power(2, LEAST(attempts, 8)))::integer), + next_attempt_at = CURRENT_TIMESTAMP + make_interval(secs => LEAST(300, power(2, LEAST(attempts, 8)) * (0.5 + random()))), state = $4::"AccountingDeliveryState" WHERE delivery_id = $1 AND state = 'pending' AND attempts = $2::integer -RETURNING delivery_id, attempts, accounting_payload""" +RETURNING delivery_id, attempts""" +_UNDECODABLE: Final = """UPDATE "LiteLLM_SpendLogDelivery" +SET state = 'parked', attempts = attempts + 1, last_error = 'UndecodablePayload' +WHERE delivery_id = ANY($1::text[]) AND state = 'pending'""" + + +def _decoded(delivery: _ClaimedDelivery) -> DeliveryProjection | None: + try: + return DeliveryProjection(_PAYLOAD.validate_json(delivery.accounting_payload)) + except ValueError: + return None def _is_permanent_failure(error: Exception) -> bool: @@ -138,26 +153,25 @@ class AccountingDeliveryRepository: try: async with self._database.tx(timeout=_TIMEOUT) as transaction: retries: Final = ( - _DELIVERIES.validate_python(await transaction.query_raw(_CLAIM, 1, True)) + _CLAIMED.validate_python(await transaction.query_raw(_CLAIM, 1, True)) if retry_only is not False else () ) - claimed = retries or ( # rebind-ok: retain claimed identities after rollback - _DELIVERIES.validate_python(await transaction.query_raw(_CLAIM, batch_size, False)) + rows: Final = retries or ( + _CLAIMED.validate_python(await transaction.query_raw(_CLAIM, batch_size, False)) if retry_only is not True else () ) - if not claimed: - return DeliveryBatchResult() - await apply( - transaction, - tuple( - DeliveryProjection(_PAYLOAD.validate_python(delivery.accounting_payload)) - for delivery in claimed - ), - ) - await transaction.execute_raw(_APPLIED, tuple(delivery.delivery_id for delivery in claimed)) - return DeliveryBatchResult(applied=len(claimed)) + decoded: Final = tuple((row, _decoded(row)) for row in rows) + undecodable: Final = tuple(row.delivery_id for row, projection in decoded if projection is None) + applicable: Final = tuple((row, projection) for row, projection in decoded if projection is not None) + if undecodable: + await transaction.execute_raw(_UNDECODABLE, undecodable) + claimed = tuple(row for row, _ in applicable) # rebind-ok: retain claimed identities after rollback + if claimed: + await apply(transaction, tuple(projection for _, projection in applicable)) + await transaction.execute_raw(_APPLIED, tuple(row.delivery_id for row in claimed)) + return DeliveryBatchResult(applied=len(claimed), parked=undecodable) except Exception as error: # noqa: BLE001 # roll back all destinations before recording a retry if not claimed: raise diff --git a/litellm/repositories/config_repository.py b/litellm/repositories/config_repository.py index 8768d79b0dd..36c6ab7f711 100644 --- a/litellm/repositories/config_repository.py +++ b/litellm/repositories/config_repository.py @@ -4,6 +4,7 @@ from __future__ import annotations import json from collections.abc import Mapping, Sequence +from datetime import datetime from typing import TYPE_CHECKING, Final, Protocol, cast if TYPE_CHECKING: @@ -35,6 +36,8 @@ class _ConfigTable(Protocol): async def delete(self, *, where: Mapping[str, str]) -> _ConfigRow | None: ... + async def delete_many(self, *, where: Mapping[str, object]) -> int: ... + class _ConfigDatabase(Protocol): async def query_raw(self, query: str, *args: object) -> object: ... @@ -130,3 +133,9 @@ class ConfigRepository: param_value = _decoded_json(param_value) result[record.param_name] = param_value return result + + async def delete_params_older_than(self, prefix: str, before: datetime) -> int: + """Expire operational config records using the timestamp written by revision claims.""" + return await self._config_table.delete_many( + where={"param_name": {"startswith": prefix}, "last_run_at": {"lt": before}} + ) diff --git a/litellm/repositories/plugin_publication.py b/litellm/repositories/plugin_publication.py new file mode 100644 index 00000000000..f5fab458c6b --- /dev/null +++ b/litellm/repositories/plugin_publication.py @@ -0,0 +1,175 @@ +import hashlib +import json +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Final, Protocol + +from pydantic import TypeAdapter +from semver import Version + +from litellm.models.plugin_upload import ( + PluginPublication, + PluginPublicationBase, + PluginPublicationError, + PublishedPlugin, +) +from litellm.repositories.prisma_protocols import TableActions + +if TYPE_CHECKING: + from prisma.models import ( + LiteLLM_ClaudeCodePluginTable, + LiteLLM_ClaudeCodePluginUpload, + LiteLLM_ClaudeCodePluginVersion, + ) + +_MANIFEST: Final = TypeAdapter(dict[str, object]) +_SOURCE: Final = TypeAdapter(dict[str, str]) + + +class PluginPublicationTransaction(Protocol): + @property + def litellm_claudecodeplugintable(self) -> TableActions["LiteLLM_ClaudeCodePluginTable"]: ... + + @property + def litellm_claudecodepluginupload(self) -> TableActions["LiteLLM_ClaudeCodePluginUpload"]: ... + + @property + def litellm_claudecodepluginversion(self) -> TableActions["LiteLLM_ClaudeCodePluginVersion"]: ... + + async def query_raw(self, query: str, *args: object) -> object: ... + + +def resolve_plugin_version( + requested: str | None, current: str | None = None, *, updating: bool = False +) -> str | PluginPublicationError: + try: + previous: Final = Version.parse(current) if current is not None else None + except ValueError: + return _legacy_version(requested) + if updating and previous is None: + return _legacy_version(requested) + try: + candidate: Final = Version.parse(requested) if requested is not None else None + except ValueError: + return PluginPublicationError("invalid_version", "Version must be a valid semantic version, for example 1.0.1") + if candidate is None: + return str(previous.bump_patch()) if previous is not None else "1.0.0" + if previous is not None and candidate.compare(previous) <= 0: + return PluginPublicationError("invalid_version", f"Version must be strictly higher than {current}") + return str(candidate) + + +def _legacy_version(requested: str | None) -> str | PluginPublicationError: + if requested is None: + return PluginPublicationError( + "conflict", + "The published version is missing or is not SemVer. Supply an explicit semantic version to update it", + ) + return resolve_plugin_version(requested) + + +def verified_archive(content: bytes | None, manifest_json: str | None) -> bytes | None: + if content is None: + return None + manifest: Final = _MANIFEST.validate_json(manifest_json or "{}") + source: Final = _SOURCE.validate_python(manifest.get("source", {})) + if source.get("source") != "archive" or source.get("sha256", "").lower() != hashlib.sha256(content).hexdigest(): + return None + return content + + +async def _snapshot( + tx: PluginPublicationTransaction, + plugin: "LiteLLM_ClaudeCodePluginTable", + content: bytes | None, + published_by: str | None, +) -> None: + from prisma import Base64 + + archive: Final = verified_archive(content, plugin.manifest_json) + await tx.litellm_claudecodepluginversion.create( + data={ + "plugin_id": plugin.id, + "version": plugin.version, + "description": plugin.description, + "manifest_json": plugin.manifest_json, + "published_at": plugin.updated_at, + "published_by": published_by, + "content": Base64.encode(archive) if archive is not None else None, + "archive_available": archive is not None, + } + ) + + +async def publish_plugin( + tx: PluginPublicationTransaction, + publication: PluginPublication, + expected: PluginPublicationBase | None = None, +) -> PublishedPlugin | PluginPublicationError: + from prisma import Base64 + + await tx.query_raw('SELECT id FROM "LiteLLM_ClaudeCodePluginTable" WHERE name = $1 FOR UPDATE', publication.name) + current: Final = await tx.litellm_claudecodeplugintable.find_unique(where={"name": publication.name}) + if current is not None and expected is None: + return PluginPublicationError("conflict", f"A skill named '{publication.name}' already exists") + if expected is not None and current is None: + return PluginPublicationError("not_found", f"Plugin '{publication.name}' not found") + if current is not None and ( + expected is not None and (current.id != expected.id or current.updated_at != expected.updated_at) + ): + return PluginPublicationError("conflict", "The published skill changed. Reload it before publishing an update") + version: Final = resolve_plugin_version( + publication.version, current.version if current is not None else None, updating=expected is not None + ) + if isinstance(version, PluginPublicationError): + return version + manifest_json: Final = json.dumps({**publication.manifest, "version": version}) + if publication.content is not None and verified_archive(publication.content, manifest_json) is None: + return PluginPublicationError("conflict", "Uploaded archive does not match its published source digest") + timestamp: Final = datetime.now(timezone.utc) + data: Final = { + "version": version, + "description": publication.description, + "manifest_json": manifest_json, + "files_json": "{}", + "updated_at": timestamp, + } + if current is None: + created: Final = await tx.litellm_claudecodeplugintable.create( + data={**data, "name": publication.name, "created_by": publication.created_by} + ) + if publication.content is not None: + await tx.litellm_claudecodepluginupload.create( + data={"plugin_name": publication.name, "content": Base64.encode(publication.content)} + ) + await _snapshot(tx, created, publication.content, publication.created_by) + return PublishedPlugin(created.id, created.name, version, created.enabled, created.updated_at or timestamp) + previous_archive: Final = await tx.litellm_claudecodepluginupload.find_unique(where={"plugin_name": current.name}) + previous_content: Final = ( + verified_archive(previous_archive.content.decode(), current.manifest_json) + if previous_archive is not None + else None + ) + history_count: Final = await tx.litellm_claudecodepluginversion.count(where={"plugin_id": current.id}, take=1) + if history_count == 0: + await _snapshot(tx, current, previous_content, None) + unchanged_content: Final = ( + previous_content + if _MANIFEST.validate_json(current.manifest_json or "{}").get("source") == publication.manifest.get("source") + else None + ) + content: Final = publication.content if publication.content is not None else unchanged_content + updated: Final = await tx.litellm_claudecodeplugintable.update(where={"id": current.id}, data=data) + if updated is None: + raise RuntimeError("Locked plugin disappeared during publication") + if content is None: + await tx.litellm_claudecodepluginupload.delete_many(where={"plugin_name": current.name}) + else: + await tx.litellm_claudecodepluginupload.upsert( + where={"plugin_name": current.name}, + data={ + "create": {"plugin_name": current.name, "content": Base64.encode(content)}, + "update": {"content": Base64.encode(content)}, + }, + ) + await _snapshot(tx, updated, content, publication.created_by) + return PublishedPlugin(updated.id, updated.name, version, updated.enabled, updated.updated_at or timestamp) diff --git a/litellm/repositories/plugin_upload_repository.py b/litellm/repositories/plugin_upload_repository.py index 99f1f76abb5..99e226e442b 100644 --- a/litellm/repositories/plugin_upload_repository.py +++ b/litellm/repositories/plugin_upload_repository.py @@ -1,60 +1,197 @@ -import json +from collections.abc import AsyncIterator, Mapping, Sequence +from contextlib import AbstractAsyncContextManager +from dataclasses import dataclass from types import MappingProxyType -from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # Prisma dynamically exposes its typed table actions +from typing import ( + TYPE_CHECKING, + Final, + Protocol, + cast, # noqa: TID251 # Prisma dynamically exposes its typed table actions +) -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError -from litellm.models.plugin_upload import PluginArchiveMetadata, PluginUpload +from litellm.models.plugin_upload import ( + PluginArchiveMetadata, + PluginDownloader, + PluginDownloadEvent, + PluginDownloadResult, + PluginDownloadsResponse, + PluginPublication, + PluginPublicationBase, + PluginPublicationError, + PluginSnapshot, + PluginUpload, + PluginVersion, + PublishedPlugin, + has_agent_skill_metadata, +) +from litellm.repositories.plugin_publication import PluginPublicationTransaction, publish_plugin, verified_archive from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ClaudeCodePluginRepository if TYPE_CHECKING: from prisma import models as prisma_models + from litellm.proxy.db.create_views import SupportsExecuteRaw from litellm.proxy.utils import PrismaClient _MANIFEST: Final = TypeAdapter(dict[str, object]) _SOURCE: Final = TypeAdapter(dict[str, str]) +PAGE_SIZE: Final = 100 +_INCREMENT_ANONYMOUS_DOWNLOADS: Final = """ +UPDATE "LiteLLM_ClaudeCodePluginTable" +SET "anonymous_download_count" = "anonymous_download_count" + 1 +WHERE "id" = $1 +""" +_DOWNLOAD_COUNTS_QUERY: Final = """ +SELECT p.id, p.anonymous_download_count + COUNT(d.id) AS download_count +FROM "LiteLLM_ClaudeCodePluginTable" p +LEFT JOIN "LiteLLM_ClaudeCodePluginDownload" d ON d.plugin_id = p.id +WHERE p.id = ANY($1::text[]) +GROUP BY p.id +""" + + +@dataclass(frozen=True, slots=True) +class _DownloadCount: + id: str + download_count: int + + +_DOWNLOAD_COUNTS: Final = TypeAdapter(tuple[_DownloadCount, ...]) +_VERSIONS: Final = TypeAdapter(tuple[PluginVersion, ...]) +_VERSION_QUERY: Final = """ + SELECT v.id, v.version, v.description, COALESCE(v.manifest_json, '{}')::jsonb AS manifest, + v.published_at, v.published_by, v.archive_available + FROM "LiteLLM_ClaudeCodePluginVersion" v JOIN "LiteLLM_ClaudeCodePluginTable" p ON p.id = v.plugin_id + WHERE p.id = $1 AND p.enabled = $2 +""" + + +class _PublicationDatabase(PluginPublicationTransaction, Protocol): + def tx(self) -> AbstractAsyncContextManager[PluginPublicationTransaction]: ... class PluginUploadRepository: def __init__(self, prisma_client: "PrismaClient") -> None: + self._db = cast( # cast-ok: the generated Prisma transaction delegate is untyped at this boundary + _PublicationDatabase, prisma_client.writer_db + ) + self._writer = cast( # cast-ok: PrismaWrapper dynamically delegates raw statements to its generated client + "SupportsExecuteRaw", prisma_client.writer_db + ) self._plugins = ClaudeCodePluginRepository(prisma_client, use_writer=True).table self._archives = cast( # cast-ok: Prisma exposes generated table actions through an untyped delegate TableActions["prisma_models.LiteLLM_ClaudeCodePluginUpload"], prisma_client.writer_db.litellm_claudecodepluginupload, ) + self._downloads = cast( # cast-ok: Prisma exposes generated table actions through an untyped delegate + TableActions["prisma_models.LiteLLM_ClaudeCodePluginDownload"], + prisma_client.writer_db.litellm_claudecodeplugindownload, + ) async def create(self, upload: PluginUpload) -> str | None: - from prisma import Base64 + result: Final = await self.publish( + PluginPublication( + name=upload.name, + version=upload.version, + description=upload.description, + manifest=upload.manifest, + content=upload.content, + created_by=upload.created_by, + ) + ) + return result.id if isinstance(result, PublishedPlugin) else None + + async def publish( + self, publication: PluginPublication, expected: PluginPublicationBase | None = None + ) -> PublishedPlugin | PluginPublicationError: from prisma.errors import UniqueViolationError try: - plugin: Final = await self._plugins.create( - data={ - "name": upload.name, - "version": upload.version, - "description": upload.description, - "manifest_json": json.dumps(dict(upload.manifest)), - "created_by": upload.created_by, - "uploaded_archive": {"create": {"content": Base64.encode(upload.content)}}, - } - ) + async with self._db.tx() as tx: + return await publish_plugin(tx, publication, expected) except UniqueViolationError: + return PluginPublicationError("conflict", "A skill with this name or version already exists") + + async def versions(self, metadata: PluginArchiveMetadata, limit: int, offset: int) -> tuple[PluginVersion, ...]: + rows: Final = await self._db.query_raw( + _VERSION_QUERY + " ORDER BY v.published_at DESC NULLS LAST, v.id DESC LIMIT $3 OFFSET $4", + metadata.id, + metadata.enabled, + limit, + offset, + ) + return _VERSIONS.validate_python(rows) + + async def version(self, metadata: PluginArchiveMetadata, version_id: str) -> PluginVersion | None: + rows: Final = await self._db.query_raw( + _VERSION_QUERY + " AND v.id = $3", + metadata.id, + metadata.enabled, + version_id, + ) + versions: Final = _VERSIONS.validate_python(rows) + return versions[0] if versions else None + + async def version_content(self, metadata: PluginArchiveMetadata, version_id: str) -> bytes | None: + version: Final = await self._db.litellm_claudecodepluginversion.find_first( + where={ + "id": version_id, + "plugin_id": metadata.id, + "plugin": { + "is": { + "enabled": metadata.enabled, + } + }, + } + ) + if version is None or version.content is None: return None - return plugin.id + return verified_archive(version.content.decode(), version.manifest_json) + + async def get_plugin(self, name: str) -> PluginSnapshot | None: + row: Final = await self._plugins.find_unique(where={"name": name}) + if row is None: + return None + return PluginSnapshot( + id=row.id, + name=row.name, + version=row.version, + description=row.description, + manifest=MappingProxyType(_MANIFEST.validate_json(row.manifest_json or "{}")), + enabled=row.enabled, + created_at=row.created_at, + updated_at=row.updated_at, + created_by=row.created_by, + ) async def get_metadata(self, name: str) -> PluginArchiveMetadata | None: plugin: Final = await self._plugins.find_unique(where={"name": name}) if plugin is None: return None - manifest: Final = _MANIFEST.validate_json(plugin.manifest_json or "{}") - return PluginArchiveMetadata( - name=plugin.name, - enabled=plugin.enabled, - source=MappingProxyType(_SOURCE.validate_python(manifest.get("source", {}))), - manifest_json=plugin.manifest_json, - ) + return _archive_metadata(plugin) + + async def public_metadata(self) -> AsyncIterator[PluginArchiveMetadata]: + after: str | None = None # rebind-ok: bounded keyset traversal advances after each raw page + while True: + plugins: Sequence[prisma_models.LiteLLM_ClaudeCodePluginTable] = await self._plugins.find_many( + where={"enabled": True} if after is None else {"enabled": True, "name": {"gt": after}}, + take=PAGE_SIZE, + order={"name": "asc"}, + ) + if not plugins: + return + for plugin in plugins: + try: + metadata: PluginArchiveMetadata = _archive_metadata(plugin) + except ValidationError: + continue + yield metadata + if len(plugins) < PAGE_SIZE: + return + after = plugins[-1].name async def get_content(self, metadata: PluginArchiveMetadata) -> bytes | None: archive: Final = await self._archives.find_first( @@ -62,6 +199,7 @@ class PluginUploadRepository: "plugin_name": metadata.name, "plugin": { "is": { + "id": metadata.id, "manifest_json": metadata.manifest_json, "enabled": metadata.enabled, } @@ -69,3 +207,84 @@ class PluginUploadRepository: } ) return None if archive is None else archive.content.decode() + + async def record_download( + self, metadata: PluginArchiveMetadata, digest: str, downloader: PluginDownloader + ) -> PluginDownloadResult: + from prisma.errors import ForeignKeyViolationError + + if metadata.id is None: + return PluginDownloadResult.SKILL_NOT_FOUND + if not downloader.authenticated: + # A counter bounds anonymous storage; raw SQL preserves Prisma's @updatedAt metadata. + affected: Final = await self._writer.execute_raw(_INCREMENT_ANONYMOUS_DOWNLOADS, metadata.id) + return PluginDownloadResult.RECORDED if affected == 1 else PluginDownloadResult.SKILL_NOT_FOUND + try: + await self._downloads.create( + data={ + "plugin_id": metadata.id, + "version": metadata.version, + "digest": digest, + "authenticated": downloader.authenticated, + "user_id": downloader.user_id, + "user_email": downloader.user_email, + "key_alias": downloader.key_alias, + "team_id": downloader.team_id, + "team_alias": downloader.team_alias, + } + ) + except ForeignKeyViolationError: + return PluginDownloadResult.SKILL_NOT_FOUND + return PluginDownloadResult.RECORDED + + async def downloads(self, name: str, limit: int, cursor: int | None) -> PluginDownloadsResponse | None: + plugin: Final = await self._plugins.find_unique(where={"name": name}) + if plugin is None: + return None + total: Final = await self._downloads.count(where={"plugin_id": plugin.id}) + rows: Final = await self._downloads.find_many( + where={"plugin_id": plugin.id, **({"id": {"lt": cursor}} if cursor is not None else {})}, + order={"id": "desc"}, + take=limit + 1, + ) + return PluginDownloadsResponse( + total_downloads=total + plugin.anonymous_download_count, + anonymous_downloads=plugin.anonymous_download_count, + downloads=tuple( + PluginDownloadEvent( + id=str(row.id), + downloaded_at=row.downloaded_at, + version=row.version, + digest=row.digest, + authenticated=row.authenticated, + user_id=row.user_id, + user_email=row.user_email, + key_alias=row.key_alias, + team_id=row.team_id, + team_alias=row.team_alias, + ) + for row in rows[:limit] + ), + next_cursor=str(rows[limit - 1].id) if len(rows) > limit else None, + ) + + async def download_counts(self, plugin_ids: tuple[str, ...]) -> Mapping[str, int]: + if not plugin_ids: + return MappingProxyType({}) + rows: Final = _DOWNLOAD_COUNTS.validate_python(await self._db.query_raw(_DOWNLOAD_COUNTS_QUERY, plugin_ids)) + return MappingProxyType({row.id: row.download_count for row in rows}) + + +def _archive_metadata(plugin: "prisma_models.LiteLLM_ClaudeCodePluginTable") -> PluginArchiveMetadata: + manifest: Final = _MANIFEST.validate_json(plugin.manifest_json or "{}") + return PluginArchiveMetadata( + name=plugin.name, + enabled=plugin.enabled, + source=MappingProxyType(_SOURCE.validate_python(manifest.get("source", {}))), + manifest_json=plugin.manifest_json, + id=plugin.id, + version=plugin.version, + description=plugin.description, + has_agent_skill_metadata=has_agent_skill_metadata(manifest), + updated_at=plugin.updated_at, + ) diff --git a/litellm/repositories/skill_submission_repository.py b/litellm/repositories/skill_submission_repository.py new file mode 100644 index 00000000000..829aa0aadbf --- /dev/null +++ b/litellm/repositories/skill_submission_repository.py @@ -0,0 +1,368 @@ +import hashlib +import json +from collections.abc import Mapping +from contextlib import AbstractAsyncContextManager +from dataclasses import replace +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Final, Protocol, cast # noqa: TID251 # Prisma's generated delegates are dynamic +from urllib.parse import urlsplit + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from litellm.models.plugin_upload import ( + PluginPublication, + PluginPublicationBase, + PluginPublicationError, + PublishedPlugin, + uploaded_plugin_source, +) +from litellm.models.skill_submission import ( + SkillProposal, + SkillScanInput, + SkillSubmission, + SubmissionConflict, + SubmissionLimitExceeded, + SubmissionMissing, + SubmissionOperation, + SubmissionPage, + SubmissionReviewed, + SubmissionReviewResult, + SubmissionStatus, +) +from litellm.repositories.plugin_publication import PluginPublicationTransaction, publish_plugin, verified_archive +from litellm.repositories.prisma_protocols import TableActions +from litellm.types.proxy.skill_scanning import SkillScan + +if TYPE_CHECKING: + from prisma import models as prisma_models + + from litellm.proxy.utils import PrismaClient + + +class _StoredProposal(BaseModel): + model_config = ConfigDict(frozen=True) + + manifest: Mapping[str, object] + version: str | None + + def proposal(self, name: str) -> SkillProposal: + return SkillProposal( + name, self.version, _DESCRIPTION.validate_python(self.manifest.get("description")), self.manifest + ) + + +class _PendingUsage(BaseModel): + model_config = ConfigDict(frozen=True) + + count: int + storage_bytes: int + + +class _SubmissionRow(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str + plugin_name: str + operation: SubmissionOperation + status: SubmissionStatus + proposal_json: str + has_archive: bool + submitted_by: str + created_at: datetime + base_plugin_id: str | None + base_updated_at: datetime | None + reviewed_by: str | None + reviewed_at: datetime | None + rejection_reason: str | None + + def submission(self) -> SkillSubmission: + stored: Final = _StoredProposal.model_validate_json(self.proposal_json) + return SkillSubmission( + id=self.id, + plugin_name=self.plugin_name, + operation=self.operation, + status=self.status, + proposal=stored.proposal(self.plugin_name), + has_archive=self.has_archive, + submitted_by=self.submitted_by, + created_at=self.created_at, + base_plugin_id=self.base_plugin_id, + base_updated_at=self.base_updated_at, + reviewed_by=self.reviewed_by, + reviewed_at=self.reviewed_at, + rejection_reason=self.rejection_reason, + ) + + +_SCAN_ROWS: Final = TypeAdapter(tuple[Mapping[str, str | None], ...]) +_ROWS: Final = TypeAdapter(tuple[_SubmissionRow, ...]) +_DESCRIPTION: Final = TypeAdapter[str | None](str | None) +_SOURCE: Final = TypeAdapter(dict[str, str]) +_USAGE: Final = TypeAdapter(tuple[_PendingUsage, ...]) +_MAX_PENDING_SUBMISSIONS: Final = 20 +_MAX_PENDING_BYTES: Final = 50 * 1024 * 1024 +_SELECT: Final = """ +SELECT id, plugin_name, operation, status, proposal_json, + archive_content IS NOT NULL AS has_archive, submitted_by, created_at, + base_plugin_id, base_updated_at, reviewed_by, reviewed_at, rejection_reason +FROM "LiteLLM_SkillSubmission" +""" + + +class SubmissionTransaction(PluginPublicationTransaction, Protocol): + @property + def litellm_skillsubmission(self) -> TableActions["prisma_models.LiteLLM_SkillSubmission"]: ... + + +class SubmissionDatabase(SubmissionTransaction, Protocol): + def tx(self) -> AbstractAsyncContextManager[SubmissionTransaction]: ... + + +def _submission(row: "prisma_models.LiteLLM_SkillSubmission") -> SkillSubmission: + return _SubmissionRow.model_validate( + {**row.model_dump(), "has_archive": row.archive_content is not None} + ).submission() + + +class SkillSubmissionRepository: + def __init__(self, prisma_client: "PrismaClient") -> None: + self._database = cast( # cast-ok: Prisma dynamically exposes generated delegates and transactions + SubmissionDatabase, prisma_client.writer_db + ) + + async def get_scan(self, submission_id: str) -> SkillScan | SubmissionMissing | None: + rows: Final = _SCAN_ROWS.validate_python( + await self._database.query_raw( + 'SELECT scan_json FROM "LiteLLM_SkillSubmission" WHERE id = $1', submission_id + ) + ) + if not rows: + return SubmissionMissing() + value: Final = rows[0]["scan_json"] + return None if value is None else SkillScan.model_validate_json(value) + + async def start_scan( + self, submission_id: str, scan: SkillScan + ) -> SkillScanInput | SubmissionMissing | SubmissionConflict: + async with self._database.tx() as transaction: + await transaction.query_raw( + 'SELECT id FROM "LiteLLM_SkillSubmission" WHERE id = $1 FOR UPDATE', submission_id + ) + row: Final = await transaction.litellm_skillsubmission.find_unique(where={"id": submission_id}) + if row is None: + return SubmissionMissing() + if row.status != "pending_review": + return SubmissionConflict("Only pending submissions can be scanned.") + previous: Final = None if row.scan_json is None else SkillScan.model_validate_json(row.scan_json) + if previous is not None and previous.at(scan.started_at).status == "running": + return SubmissionConflict("A scan is already running for this submission.") + stored: Final = _StoredProposal.model_validate_json(row.proposal_json) + source: Final = _SOURCE.validate_python(stored.manifest.get("source", {})) + inherited: Final = ( + await transaction.litellm_claudecodepluginupload.find_unique(where={"plugin_name": row.plugin_name}) + if row.archive_content is None and source.get("source") == "archive" + else None + ) + content: Final = ( + row.archive_content.decode() + if row.archive_content is not None + else verified_archive( + None if inherited is None else inherited.content.decode(), json.dumps(stored.manifest) + ) + ) + managed_archive: Final = urlsplit(source.get("url", "")).path.endswith( + f"/claude-code/plugins/{row.plugin_name}/archive/{source.get('sha256', '')}.zip" + ) + if source.get("source") == "archive" and managed_archive and content is None: + return SubmissionConflict("The proposal's archive is no longer available. Submit a new ZIP to scan it.") + started: Final = scan.model_copy( + update={"archive_sha256": None if content is None else hashlib.sha256(content).hexdigest()} + ) + await transaction.litellm_skillsubmission.update( + where={"id": submission_id}, data={"scan_json": started.model_dump_json()} + ) + return SkillScanInput(started, source, content) + + async def finish_scan(self, submission_id: str, scan: SkillScan) -> None: + async with self._database.tx() as transaction: + rows: Final = _SCAN_ROWS.validate_python( + await transaction.query_raw( + 'SELECT scan_json FROM "LiteLLM_SkillSubmission" WHERE id = $1 FOR UPDATE', submission_id + ) + ) + value: Final = rows[0]["scan_json"] if rows else None + if value is None: + return + current: Final = SkillScan.model_validate_json(value) + if current.id != scan.id or current.status != "running": + return + result: Final = ( + current.at(scan.completed_at) + if scan.completed_at is not None and scan.completed_at >= current.expires_at + else scan + ) + await transaction.litellm_skillsubmission.update( + where={"id": submission_id}, data={"scan_json": result.model_dump_json()} + ) + + async def submit( + self, upload: PluginPublication, submitted_by: str, base: PluginPublicationBase | None + ) -> SkillSubmission | SubmissionLimitExceeded: + from prisma import Base64 + + proposal_json: Final = _StoredProposal(manifest=dict(upload.manifest), version=upload.version).model_dump_json() + incoming_bytes: Final = len(proposal_json.encode("utf-8")) + len(upload.content or b"") + async with self._database.tx() as transaction: + await transaction.query_raw( + "SELECT 1 AS locked FROM pg_advisory_xact_lock(hashtext('litellm.skill_submissions'), hashtext($1))", + submitted_by, + ) + usage: Final = _USAGE.validate_python( + await transaction.query_raw( + """SELECT COUNT(*) AS count, + COALESCE(SUM(octet_length(proposal_json) + + COALESCE(octet_length(archive_content), 0)), 0) AS storage_bytes + FROM "LiteLLM_SkillSubmission" WHERE submitted_by = $1 AND status = 'pending_review'""", + submitted_by, + ) + )[0] + if usage.count >= _MAX_PENDING_SUBMISSIONS: + return SubmissionLimitExceeded( + "You already have 20 pending skill submissions. " + "Wait for an admin to review one before submitting another." + ) + if usage.storage_bytes + incoming_bytes > _MAX_PENDING_BYTES: + return SubmissionLimitExceeded( + "Pending skill submissions cannot exceed 50 MiB of stored metadata and archives per user. " + "Wait for an admin review or submit less data." + ) + row: Final = await transaction.litellm_skillsubmission.create( + data={ + "plugin_name": upload.name, + "operation": "create" if base is None else "update", + "proposal_json": proposal_json, + "archive_content": None if upload.content is None else Base64.encode(upload.content), + "submitted_by": submitted_by, + "base_plugin_id": None if base is None else base.id, + "base_updated_at": None if base is None else base.updated_at, + } + ) + return _submission(row) + + async def list( + self, submitted_by: str | None, status: SubmissionStatus | None, limit: int, offset: int + ) -> SubmissionPage: + rows: Final = _ROWS.validate_python( + await self._database.query_raw( + _SELECT + + """WHERE ($1::text IS NULL OR submitted_by = $1) + AND ($2::text IS NULL OR status = $2) + ORDER BY created_at DESC, id DESC LIMIT $3 OFFSET $4""", + submitted_by, + status, + limit + 1, + offset, + ) + ) + return SubmissionPage(tuple(row.submission() for row in rows[:limit]), has_more=len(rows) > limit) + + async def get(self, submission_id: str, submitted_by: str | None) -> SkillSubmission | None: + rows: Final = _ROWS.validate_python( + await self._database.query_raw( + _SELECT + "WHERE id = $1 AND ($2::text IS NULL OR submitted_by = $2)", + submission_id, + submitted_by, + ) + ) + return None if not rows else rows[0].submission() + + async def get_archive(self, submission_id: str, submitted_by: str | None) -> bytes | None: + row: Final = await self._database.litellm_skillsubmission.find_first( + where={"id": submission_id, **({"submitted_by": submitted_by} if submitted_by is not None else {})} + ) + return None if row is None or row.archive_content is None else row.archive_content.decode() + + async def review( + self, + submission_id: str, + reviewer: str, + approve: bool, + reason: str | None, + *, + publication_base_url: str | None = None, + ) -> SubmissionReviewResult: + from prisma.errors import UniqueViolationError + + try: + async with self._database.tx() as transaction: + await transaction.query_raw( + 'SELECT id FROM "LiteLLM_SkillSubmission" WHERE id = $1 FOR UPDATE', submission_id + ) + row: Final = await transaction.litellm_skillsubmission.find_unique(where={"id": submission_id}) + if row is None: + return SubmissionMissing() + submission: Final = _submission(row) + if submission.status != "pending_review": + return SubmissionConflict("This submission has already been reviewed.") + content: Final = None if row.archive_content is None else row.archive_content.decode() + if approve and content is not None and not publication_base_url: + return SubmissionConflict("Uploaded skill approval requires a trusted publication URL.") + source: Final = ( + uploaded_plugin_source(publication_base_url, row.plugin_name, hashlib.sha256(content).hexdigest()) + if approve and content is not None and publication_base_url + else _SOURCE.validate_python(submission.proposal.manifest.get("source")) + ) + publication: Final = await _publish_submission(transaction, row, source, content) if approve else None + if isinstance(publication, PluginPublicationError): + return publication + reviewed_at: Final = datetime.now(timezone.utc) + status: Final[SubmissionStatus] = "approved" if approve else "rejected" + await transaction.litellm_skillsubmission.update( + where={"id": submission_id}, + data={ + "status": status, + "reviewed_by": reviewer, + "reviewed_at": reviewed_at, + "rejection_reason": None if approve else reason, + }, + ) + return SubmissionReviewed( + submission=replace( + submission, + status=status, + reviewed_by=reviewer, + reviewed_at=reviewed_at, + rejection_reason=None if approve else reason, + ), + plugin=publication, + published_source=source if approve else None, + ) + except UniqueViolationError: + return SubmissionConflict("A skill with this name is already published. Submit a new update proposal.") + + +async def _publish_submission( + transaction: PluginPublicationTransaction, + row: "prisma_models.LiteLLM_SkillSubmission", + source: Mapping[str, str], + content: bytes | None, +) -> PublishedPlugin | PluginPublicationError: + stored: Final = _StoredProposal.model_validate_json(row.proposal_json) + proposal: Final = stored.proposal(row.plugin_name) + base: Final = ( + None + if row.operation == "create" + else PluginPublicationBase(id=row.base_plugin_id or "", updated_at=row.base_updated_at) + ) + return await publish_plugin( + transaction, + PluginPublication( + name=row.plugin_name, + version=stored.version, + description=proposal.description, + manifest={**stored.manifest, "source": dict(source)}, + content=content, + created_by=row.submitted_by, + ), + base, + ) diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 10db81eca4b..d010455faa3 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -387,3 +387,7 @@ class BackgroundInteractionSettlementRepository( class RetiredAgentRepository(PrismaTableRepository["prisma_models.LiteLLM_RetiredAgent"]): table_name = "litellm_retiredagent" + + +class DailyUserSpendRepository(PrismaTableRepository["prisma_models.LiteLLM_DailyUserSpend"]): + table_name = "litellm_dailyuserspend" diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index f6aad883479..b78f7e5a9b7 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -109,6 +109,46 @@ class TeamMembershipSpendResetWrites: self.batch.execute_raw(_TEAM_MEMBERSHIP_TOTAL_SPEND_REPAIR_SQL, list(budget_ids)) +_LOCK_BUDGET_ORGS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_OrganizationTable" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY organization_id COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_MEMBERSHIPS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TeamMembership" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY team_id COLLATE "C", user_id COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_KEYS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_VerificationToken" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY token COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_PROJECTS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_ProjectTable" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY project_id COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_TAGS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_TagTable" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY tag_name COLLATE "C" FOR UPDATE' +) +_LOCK_BUDGET_ACCESS_GROUPS_SQL: Final = ( + 'SELECT 1 AS locked FROM "LiteLLM_ModelAccessGroupBudgetTable" WHERE budget_id = ANY($1::text[]) ' + 'ORDER BY access_group_name COLLATE "C" FOR UPDATE' +) + + +@dataclass(frozen=True, slots=True) +class BudgetLinkedRowLocks: + batch: PrismaBatch + + def queue_row_locks(self, budget_ids: Sequence[str]) -> None: + """Locks every row a cascade may reset, in the shared row lock order, before the filtered resets run.""" + self.batch.execute_raw(_LOCK_BUDGET_ORGS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_MEMBERSHIPS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_KEYS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_PROJECTS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_TAGS_SQL, list(budget_ids)) + self.batch.execute_raw(_LOCK_BUDGET_ACCESS_GROUPS_SQL, list(budget_ids)) + + @dataclass(frozen=True, slots=True) class BudgetWindowWrites: table: BatchTable @@ -137,6 +177,7 @@ class BudgetCascadeUnitOfWork: ``enduser_spend_reset_unit_of_work`` because one tier can gate millions. """ + row_locks: BudgetLinkedRowLocks team_memberships: TeamMembershipSpendResetWrites keys: LinkedSpendResetWrites organizations: LinkedSpendResetWrites @@ -163,6 +204,7 @@ async def budget_cascade_unit_of_work( ) -> AsyncGenerator[BudgetCascadeUnitOfWork, None]: batch = new_batch() yield BudgetCascadeUnitOfWork( + row_locks=BudgetLinkedRowLocks(batch=batch), team_memberships=TeamMembershipSpendResetWrites( resets=LinkedSpendResetWrites(table=batch.litellm_teammembership), batch=batch ), diff --git a/litellm/types/integrations/model_deprecation_email.py b/litellm/types/integrations/model_deprecation_email.py new file mode 100644 index 00000000000..f196bb47544 --- /dev/null +++ b/litellm/types/integrations/model_deprecation_email.py @@ -0,0 +1,35 @@ +"""Model retirement email milestones and recipient-scoped content.""" + +from dataclasses import dataclass +from datetime import date +from enum import IntEnum +from typing import Final + +from litellm.types.proxy.model_deprecation import ModelDeprecationInfo + +MODEL_DEPRECATION_EMAIL_EVENT: Final = "Model Deprecations" + + +class DeprecationEmailMilestone(IntEnum): + THIRTY_DAYS = 30 + ONE_DAY = 1 + + +@dataclass(frozen=True, slots=True) +class ModelDeprecationEmail: + model: ModelDeprecationInfo + usage_model: str + recipient_email: str + milestone: DeprecationEmailMilestone + key_count: int + requests: int + team_names: tuple[str, ...] + start_date: date + end_date: date + + +@dataclass(frozen=True, slots=True) +class RenderedDeprecationEmail: + subject: str + html: str + text: str diff --git a/litellm/types/proxy/claude_code_endpoints.py b/litellm/types/proxy/claude_code_endpoints.py index 49dd8b9ef91..5d8479b8ce1 100644 --- a/litellm/types/proxy/claude_code_endpoints.py +++ b/litellm/types/proxy/claude_code_endpoints.py @@ -2,6 +2,7 @@ Claude Code Marketplace endpoint types for LiteLLM Proxy """ +from datetime import datetime from typing import Final, Literal, TypeAlias from pydantic import AliasChoices, ConfigDict, Field, TypeAdapter, model_validator @@ -32,7 +33,7 @@ class PluginMetadata(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) - version: str | None = Field("1.0.0", description="Semantic version") + version: str | None = Field("1.0.0", description="Semantic version, defaults to 1.0.0 on creation") description: str | None = Field(None, description="Plugin description") author: PluginAuthor | None = Field(None, description="Plugin author") homepage: str | None = Field(None, description="Plugin homepage URL") @@ -105,11 +106,14 @@ class UpdatePluginRequest(PluginSpec): The plugin name is the resource identity and is supplied as the path parameter, so it cannot be changed here. This is a full replace: omitted - fields reset to their defaults, so version is cleared rather than - defaulting to the create-time "1.0.0". + fields reset to their defaults, except version which automatically increments + the published patch version when omitted or null. """ - version: str | None = Field(None, description="Semantic version; cleared if omitted") + version: str | None = Field(None, description="Explicit higher SemVer, or omit/null for an automatic patch bump") + expected_updated_at: datetime | None = Field( + None, description="Published updated_at read by the editor. A stale value returns 409", exclude=True + ) class PluginResponse(LiteLLMBaseModel): @@ -150,6 +154,7 @@ class PluginListItem(LiteLLMBaseModel): enabled: bool created_at: str | None updated_at: str | None + has_agent_skill_metadata: bool = Field(default=False, exclude=True) class ListPluginsResponse(LiteLLMBaseModel): diff --git a/litellm/types/proxy/skill_hub.py b/litellm/types/proxy/skill_hub.py new file mode 100644 index 00000000000..d575defcda7 --- /dev/null +++ b/litellm/types/proxy/skill_hub.py @@ -0,0 +1,165 @@ +from collections.abc import Mapping +from datetime import datetime +from typing import Literal, TypeAlias + +from pydantic import ConfigDict, Field + +from litellm.models.skill_submission import SubmissionOperation, SubmissionStatus +from litellm.types.llms.base import LiteLLMBaseModel + + +class SkillHubSkill(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + id: str + name: str + description: str | None = None + version: str | None = None + updated_at: datetime | None = None + visibility: Literal["public", "private"] + source: Mapping[str, str] + download_url: str | None = None + content_url: str | None = None + domain: str | None = None + namespace: str | None = None + keywords: tuple[str, ...] | None = None + download_count: int | None = Field( + None, description="Gateway ZIP downloads, included in proxy-admin catalog listings" + ) + + +class SkillHubList(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + skills: tuple[SkillHubSkill, ...] + count: int + next_offset: int | None + + +class ImportSkillRequest(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + url: str = Field(description="Public skills.sh skill URL or GitHub repository URL") + skill: str | None = Field(None, description="Skill name to select when a repository contains several skills") + ref: str = Field("HEAD", description="Git commit, tag or branch to import") + path: str | None = Field(None, description="Repository-relative skill directory") + version: str | None = Field(None, description="Semantic version, or omit/null to start at 1.0.0") + + +class SkillHubMetadata(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + description: str | None = Field(None, max_length=1024) + version: str | None = Field(None, description="Explicit higher SemVer, or omit/null for an automatic patch bump") + expected_updated_at: datetime | None = Field( + None, description="Published updated_at read by the editor. A stale value returns 409", exclude=True + ) + domain: str | None = None + namespace: str | None = None + keywords: tuple[str, ...] | None = None + + +class SkillHubPublication(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + +class SkillHubVersion(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + id: str + version: str | None + description: str | None + manifest: Mapping[str, object] + published_at: datetime | None + published_by: str | None + content_kind: Literal["stored_archive", "source_reference"] + download_url: str | None + + +class SkillHubVersions(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + versions: tuple[SkillHubVersion, ...] + next_offset: int | None + + +class SkillHubAgentSkill(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + name: str + description: str + digest: str + origin: str | None + + +class SkillHubProposal(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + name: str + description: str | None + version: str | None + domain: str | None = None + namespace: str | None = None + keywords: tuple[str, ...] | None = None + source: Mapping[str, str] + agent_skill: SkillHubAgentSkill | None = None + + +class SkillHubPendingSkill(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + status: Literal["pending_review"] = "pending_review" + action: Literal["submitted"] = "submitted" + submission_id: str + proposal: SkillHubProposal + + +SkillHubMutationResponse: TypeAlias = SkillHubSkill | SkillHubPendingSkill + + +class SkillHubSubmission(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + id: str + skill_name: str + operation: SubmissionOperation + status: SubmissionStatus + proposal: SkillHubProposal + has_archive: bool + submitted_by: str + created_at: datetime + base_skill_id: str | None + base_updated_at: datetime | None + reviewed_by: str | None + reviewed_at: datetime | None + rejection_reason: str | None + + +class SkillHubSubmissionList(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + submissions: tuple[SkillHubSubmission, ...] + count: int + has_more: bool + + +class SkillHubSubmissionDetail(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + submission: SkillHubSubmission + current_skill: SkillHubSkill | None + + +class SkillHubReviewResponse(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + status: Literal["success"] = "success" + action: Literal["approved", "rejected"] + submission: SkillHubSubmission + skill: SkillHubSkill | None + + +class SkillHubRejection(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + reason: str | None = Field(None, max_length=2000) diff --git a/litellm/types/proxy/skill_scanning.py b/litellm/types/proxy/skill_scanning.py new file mode 100644 index 00000000000..19577c8a9ac --- /dev/null +++ b/litellm/types/proxy/skill_scanning.py @@ -0,0 +1,91 @@ +from datetime import datetime +from typing import Annotated, Final, Literal +from urllib.parse import urlsplit + +from pydantic import BaseModel, ConfigDict, Field, field_validator + + +class SkillScannerConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + endpoint: str = Field(max_length=2048, description="External scanner's authenticated POST endpoint") + api_key_env: str = Field( + pattern=r"^[A-Za-z_][A-Za-z0-9_]*$", description="Environment secret name for scanner auth" + ) + use_llm: bool = False + timeout_seconds: int = Field(default=180, ge=10, le=600) + + @field_validator("endpoint") + @classmethod + def validate_endpoint(cls, value: str) -> str: + parsed: Final = urlsplit(value) + if ( + not parsed.hostname + or parsed.username is not None + or parsed.password is not None + or parsed.query + or parsed.fragment + or not ( + parsed.scheme == "https" + or (parsed.scheme == "http" and parsed.hostname in {"127.0.0.1", "::1", "localhost"}) + ) + ): + raise ValueError( + "Scanner endpoint must use HTTPS without credentials, query or fragment (HTTP loopback allowed)" + ) + return value + + +class SkillScanFinding(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + rule_id: str = Field(max_length=200) + severity: Literal["LOW", "MEDIUM", "HIGH", "CRITICAL"] + message: str = Field(max_length=8000) + file: str = Field(max_length=2048) + start_line: int | None = Field(default=None, ge=1) + explanation: str | None = Field(default=None, max_length=8000) + remediation: str | None = Field(default=None, max_length=8000) + + +class SkillScanReport(BaseModel): + model_config = ConfigDict(frozen=True, extra="ignore") + + scanner_version: str = Field(min_length=1, max_length=100) + complete: bool + scan_mode: Literal["static-only", "static+llm"] + risk_score: int = Field(ge=0, le=100) + findings: tuple[SkillScanFinding, ...] = Field(max_length=500) + limitations: tuple[Annotated[str, Field(max_length=2000)], ...] = Field(default=(), max_length=100) + + +class SkillScan(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str + status: Literal["running", "completed", "failed"] + requested_by: str + started_at: datetime + expires_at: datetime + completed_at: datetime | None = None + archive_sha256: str | None = None + report: SkillScanReport | None = None + error: str | None = None + + def at(self, now: datetime) -> "SkillScan": + if self.status == "running" and self.expires_at <= now: + return self.model_copy( + update={"status": "failed", "error": "Scan timed out or was interrupted. Try scanning again."} + ) + return self + + +class SkillScanRequest(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + +class SkillScanResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + configured: bool + scan: SkillScan | None diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 76fd85d001b..aba1957f03c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -32320,7 +32320,9 @@ "deprecation_source": "https://ai.google.dev/gemini-api/docs/deprecations", "successor_model": "gemini/gemini-3.8-flash", "input_cost_per_token_above_200k_tokens_batches": 1.25e-06, - "output_cost_per_token_above_200k_tokens_batches": 7.5e-06 + "input_cost_per_token_above_200k_tokens_flex": 1.25e-06, + "output_cost_per_token_above_200k_tokens_batches": 7.5e-06, + "output_cost_per_token_above_200k_tokens_flex": 7.5e-06 }, "gemini/gemini-2.5-computer-use-preview-10-2025": { "deprecation_date": "2026-07-28", @@ -32899,7 +32901,9 @@ "output_cost_per_token_flex": 6e-06, "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, "input_cost_per_token_above_200k_tokens_batches": 2e-06, - "output_cost_per_token_above_200k_tokens_batches": 9e-06 + "input_cost_per_token_above_200k_tokens_flex": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, + "output_cost_per_token_above_200k_tokens_flex": 9e-06 }, "gemini/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -32965,7 +32969,9 @@ "output_cost_per_token_flex": 6e-06, "cache_read_input_token_cost_above_200k_tokens_batches": 4e-07, "input_cost_per_token_above_200k_tokens_batches": 2e-06, - "output_cost_per_token_above_200k_tokens_batches": 9e-06 + "input_cost_per_token_above_200k_tokens_flex": 2e-06, + "output_cost_per_token_above_200k_tokens_batches": 9e-06, + "output_cost_per_token_above_200k_tokens_flex": 9e-06 }, "gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, @@ -60260,6 +60266,7 @@ "rpm": 10, "gemini_audio_only_live": true, "input_cost_per_second": 8.33333333333e-05, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "supports_response_schema": false, "supports_reasoning": true }, @@ -72140,6 +72147,7 @@ "supports_audio_input": true, "tpm": 250000, "rpm": 10, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -72161,6 +72169,7 @@ "supports_audio_input": true, "tpm": 250000, "rpm": 10, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "supports_function_calling": true, "supports_response_schema": false, "supports_vision": true, @@ -83130,6 +83139,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_prompt_caching": true, "supports_function_calling": true, "supports_response_schema": true, @@ -83146,6 +83160,11 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_prompt_caching": true, "supports_function_calling": true, "supports_response_schema": true, @@ -83155,6 +83174,7 @@ "litellm_provider": "empiriolabs", "mode": "chat", "max_input_tokens": 202745, + "max_output_tokens": 128000, "input_cost_per_token": 8.25e-07, "output_cost_per_token": 3.301e-06, "input_cost_per_token_above_32k_tokens": 1.1e-06, @@ -83178,6 +83198,11 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_function_calling": true, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/glm-5-2" @@ -83230,6 +83255,11 @@ "supports_system_messages": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_function_calling": true, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/glm-5-3" @@ -83245,6 +83275,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.033, + "search_context_size_medium": 0.033, + "search_context_size_high": 0.033 + }, "supports_function_calling": true, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/glm-5-3-flash" @@ -83721,6 +83756,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-122b-a10b" }, @@ -83737,6 +83777,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-27b" }, @@ -83753,6 +83798,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-35b-a3b" }, @@ -83769,6 +83819,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-5-397b-a17b" }, @@ -83830,6 +83885,10 @@ "max_output_tokens": 65536, "input_cost_per_token": 2.9e-08, "output_cost_per_token": 2.87e-07, + "input_cost_per_token_above_128k_tokens": 1.15e-07, + "output_cost_per_token_above_128k_tokens": 1.147e-06, + "input_cost_per_token_above_256k_tokens": 1.72e-07, + "output_cost_per_token_above_256k_tokens": 1.72e-06, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, @@ -83852,6 +83911,11 @@ "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, "supports_audio_input": true, "source": "https://empiriolabs.ai/models/qwen3-5-omni-flash" }, @@ -83865,6 +83929,11 @@ "supports_system_messages": true, "supports_vision": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.015, + "search_context_size_medium": 0.015, + "search_context_size_high": 0.015 + }, "supports_response_schema": true, "supports_audio_input": true, "source": "https://empiriolabs.ai/models/qwen3-5-omni-plus" @@ -83924,6 +83993,11 @@ "supports_vision": true, "supports_reasoning": true, "supports_web_search": true, + "search_context_cost_per_query": { + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01, + "search_context_size_high": 0.01 + }, "supports_response_schema": true, "source": "https://empiriolabs.ai/models/qwen3-6-27b" }, @@ -84421,7 +84495,7 @@ "max_output_tokens": 131072, "input_cost_per_token": 1e-06, "output_cost_per_token": 2.7e-06, - "cache_read_input_token_cost": 5e-08, + "cache_read_input_token_cost": 1e-07, "supports_system_messages": true, "supports_vision": true, "supports_reasoning": true, @@ -85310,6 +85384,24 @@ "output_cost_per_token": 6.09e-07, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/XiaomiMiMo/MiMo-V2.6-Flash": { + "cache_read_input_token_cost": 2.24e-09, + "input_cost_per_token": 1.12e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 2.24e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/XiaomiMiMo/MiMo-V2.6-Pro": { + "cache_read_input_token_cost": 2.88e-09, + "input_cost_per_token": 3.48e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1050000, + "mode": "chat", + "output_cost_per_token": 6.96e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/deepseek-ai/DeepSeek-Prover-V2-671B": { "input_cost_per_token": 5e-07, "litellm_provider": "gmi", @@ -85434,6 +85526,15 @@ "output_cost_per_token": 3.168e-06, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/deepseek-ai/DeepSeek-V4.1-Flash": { + "cache_read_input_token_cost": 3.6e-09, + "input_cost_per_token": 1.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/deepseek-ai/deepseek-v4-flash-vision-exp": { "cache_read_input_token_cost": 1.4e-08, "input_cost_per_token": 4.4e-07, @@ -85443,6 +85544,24 @@ "output_cost_per_token": 1.32e-06, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/deepseek-ai/deepseek-v4.1-flash-1017": { + "cache_read_input_token_cost": 3.6e-09, + "input_cost_per_token": 1.8e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1048575, + "mode": "chat", + "output_cost_per_token": 7.2e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, + "gmi/gmicloud/Ripple-65B-Preview": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 6e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/google/gemini-3.1-flash-lite-preview": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_token": 2.5e-07, @@ -85759,6 +85878,14 @@ "output_cost_per_token": 2.5e-07, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/openai/gpt-oss-20b": { + "input_cost_per_token": 4e-08, + "litellm_provider": "gmi", + "max_input_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-07, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/stepfun-ai/Step-3.7-Flash": { "cache_read_input_token_cost": 4e-08, "input_cost_per_token": 2e-07, @@ -85803,6 +85930,15 @@ "output_cost_per_token": 6e-07, "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" }, + "gmi/tencent/hy4": { + "cache_read_input_token_cost": 2.1e-08, + "input_cost_per_token": 4.17e-07, + "litellm_provider": "gmi", + "max_input_tokens": 1000000, + "mode": "chat", + "output_cost_per_token": 1.2505e-06, + "source": "https://ie-api.gmicloud.ai/v1/artifact/get_public_artifacts" + }, "gmi/tencent/hy4-preview": { "cache_read_input_token_cost": 4.2e-08, "input_cost_per_token": 8.34e-07, @@ -86176,5 +86312,41 @@ "output_cost_per_token": 6e-07, "source": "https://api.together.ai/v1/models", "supports_prompt_caching": true + }, + "vertex_ai/gemini-omni-1.2-flash": { + "input_cost_per_token": 1.5e-06, + "litellm_provider": "vertex_ai", + "max_input_tokens": 131072, + "max_output_tokens": 260640, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.2e-05, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", + "supported_endpoints": [ + "/v1beta/interactions" + ], + "supported_modalities": [ + "text", + "image", + "video" + ], + "supported_output_modalities": [ + "text", + "video" + ], + "supports_audio_input": false, + "supports_audio_output": false, + "supports_computer_use": false, + "supports_function_calling": false, + "supports_prompt_caching": false, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_system_messages": false, + "supports_tool_choice": false, + "supports_url_context": false, + "supports_video_input": true, + "supports_vision": true, + "supports_web_search": false } } diff --git a/pyproject.toml b/pyproject.toml index 2cc2d0f70a3..ee6cb6c1dcc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,6 +95,7 @@ proxy = [ "pyroscope-io>=0.8.16,<1.0; sys_platform != 'win32'", "expression>=5.6.0,<6.0", "backports-zstd>=1.0.0,<2.0; python_version < '3.14'", + "semver>=3.0.4,<4", ] # Thin client install for the `lite` CLI on developer laptops. The CLI's heavy # imports are all guarded, so it runs on the base SDK plus these packages, and @@ -269,6 +270,7 @@ proxy-dev = [ "a2a-sdk==1.1.0", "backports-zstd==1.8.0; python_version < '3.14'", "httptools>=0.7.1,<0.8", + "semver>=3.0.4,<4", ] ci = [ "psutil==7.2.2", diff --git a/schema.prisma b/schema.prisma index 78f16d9f1fd..818f2dc4858 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1616,10 +1616,69 @@ model LiteLLM_ClaudeCodePluginTable { created_by String? uploaded_archive LiteLLM_ClaudeCodePluginUpload? + versions LiteLLM_ClaudeCodePluginVersion[] + downloads LiteLLM_ClaudeCodePluginDownload[] + anonymous_download_count BigInt @default(0) @@map("LiteLLM_ClaudeCodePluginTable") } +model LiteLLM_ClaudeCodePluginVersion { + id String @id @default(uuid()) + plugin_id String + version String? + description String? + manifest_json String? + published_at DateTime? + published_by String? + archive_available Boolean @default(false) + content Bytes? + plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_id], references: [id], onDelete: Cascade, onUpdate: Cascade) + + @@unique([plugin_id, version]) + @@index([plugin_id, published_at]) +} + + +model LiteLLM_SkillSubmission { + id String @id @default(uuid()) + plugin_name String + operation String + status String @default("pending_review") + proposal_json String + archive_content Bytes? + submitted_by String + created_at DateTime @default(now()) + base_plugin_id String? + base_updated_at DateTime? + reviewed_by String? + reviewed_at DateTime? + rejection_reason String? + scan_json String? + + @@index([submitted_by, created_at(sort: Desc), id(sort: Desc)]) + @@index([submitted_by, status, created_at(sort: Desc), id(sort: Desc)]) + @@index([status, created_at(sort: Desc), id(sort: Desc)]) + @@index([created_at(sort: Desc), id(sort: Desc)]) +} + +model LiteLLM_ClaudeCodePluginDownload { + id BigInt @id @default(autoincrement()) + plugin_id String + downloaded_at DateTime @default(now()) + version String? + digest String + authenticated Boolean + user_id String? + user_email String? + key_alias String? + team_id String? + team_alias String? + plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_id], references: [id], onDelete: Cascade, onUpdate: Cascade) + + @@index([plugin_id, id]) +} + model LiteLLM_ClaudeCodePluginUpload { plugin_name String @id content Bytes diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index f0275c34c06..f4ca7fa2965 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -90,6 +90,7 @@ IGNORE_FUNCTIONS = [ "_json", # router_rust index views: max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by raising ValueError at the cap. "to_wire", # router_rust model info: bounded by the nesting depth of a stored deployment or cost-map entry (finite JSON trees: pydantic refuses a cycle when it dumps a deployment). "from_wire", # router_rust model info: walks only what router-core merged from to_wire's output, so it never goes deeper than to_wire did. + "_valid_frontmatter_node", # max depth 32, manifest size capped at 1 MiB, and YAML aliases rejected before traversal. ] diff --git a/tests/integration/_support/wire.py b/tests/integration/_support/wire.py index e278a796146..d4cbccfc43e 100644 --- a/tests/integration/_support/wire.py +++ b/tests/integration/_support/wire.py @@ -1,16 +1,25 @@ from __future__ import annotations +import io +import selectors +import socket import ssl import threading import time from collections.abc import Callable, Generator, Iterator, Mapping -from contextlib import contextmanager +from contextlib import contextmanager, suppress from dataclasses import dataclass from functools import partial from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from queue import SimpleQueue from types import MappingProxyType -from typing import BinaryIO, Final +from typing import TYPE_CHECKING, BinaryIO, Final + +if TYPE_CHECKING: + from collections.abc import Buffer + +_POLL_SECONDS: Final = 0.05 +_WRITE_PIECE_BYTES: Final = 64 * 1024 @dataclass(frozen=True, slots=True) @@ -97,12 +106,108 @@ def _await_release( return gate.is_set() +class _RequestReader(io.BufferedReader): + cut = False + + def readline(self, size: int | None = -1, /) -> bytes: + line: Final = super().readline(size) + self.cut = self.cut or not line.endswith(b"\n") + return line + + +@dataclass(frozen=True, slots=True) +class _Blocked: + since: float + writing: bool + + +class _Connection(io.BufferedIOBase): + blocked: _Blocked | None = None + + def __init__(self, connection: socket.socket, closing: threading.Event) -> None: + super().__init__() + self.connection = connection + self.closing = closing + self.selector = selectors.DefaultSelector() + self.selector.register(connection, selectors.EVENT_READ) + + def readable(self) -> bool: + return True + + def writable(self) -> bool: + return True + + def readinto(self, buffer: Buffer, /) -> int: + if not self._wait_for_bytes(): + return 0 + with self._blocking(writing=False): + return self.connection.recv_into(buffer) + + def write(self, buffer: Buffer, /) -> int: + view: Final = memoryview(buffer).cast("B") + for start in range(0, view.nbytes, _WRITE_PIECE_BYTES): + with self._blocking(writing=True): + self.connection.sendall(view[start : start + _WRITE_PIECE_BYTES]) + return view.nbytes + + def close(self) -> None: + self.selector.close() + super().close() + + def stalled(self, grace: float, teardown_started: float) -> _Blocked | None: + blocked: Final = self.blocked + if blocked is None or time.monotonic() - max(blocked.since, teardown_started) < grace: + return None + return blocked + + def shut(self) -> None: + with suppress(OSError): + self.connection.shutdown(socket.SHUT_RDWR) + + def _wait_for_bytes(self) -> bool: + while not self.closing.is_set(): + if self._has_bytes(_POLL_SECONDS): + return True + return self._has_bytes(0) + + def _has_bytes(self, timeout: float) -> bool: + decrypted: Final = isinstance(self.connection, ssl.SSLSocket) and self.connection.pending() > 0 + return decrypted or bool(self.selector.select(timeout)) + + @contextmanager + def _blocking(self, writing: bool) -> Iterator[None]: + self.blocked = _Blocked(time.monotonic(), writing) + try: + yield + finally: + self.blocked = None + + +def _close_unless_stalled( + close: Callable[[], None], connections: SimpleQueue[_Connection], grace: float, errors: SimpleQueue[Exception] +) -> None: + teardown_started: Final = time.monotonic() + closer: Final = threading.Thread(target=close) + closer.start() + while closer.is_alive(): + for connection in tuple(connections.get_nowait() for _ in range(connections.qsize())): + stalled: Final = connection.stalled(grace, teardown_started) + if stalled is None: + connections.put(connection) + continue + if stalled.writing: + errors.put(TimeoutError(f"a reply write made no progress for {grace} s")) + connection.shut() + closer.join(timeout=_POLL_SECONDS) + + @contextmanager def wire_server( respond: Callable[[Request], Reply], tls: ssl.SSLContext | None = None, port: int = 0, keep_alive: bool = False, + teardown_grace: float = 10, ) -> Generator[Wire, None, None]: """Owned TCP peer; requests traverse the real HTTP client and serialization. With `keep_alive` the peer honours HTTP/1.1 persistent connections so `connections()` counts the client's TCP sessions.""" @@ -111,14 +216,20 @@ def wire_server( disconnected: Final[SimpleQueue[str]] = SimpleQueue() connected: Final[SimpleQueue[str]] = SimpleQueue() closed: Final[SimpleQueue[str]] = SimpleQueue() + streams: Final[SimpleQueue[_Connection]] = SimpleQueue() closing: Final = threading.Event() class Handler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" - timeout = 5 + request: socket.socket + reader: _RequestReader def setup(self) -> None: - super().setup() + stream: Final = _Connection(self.request, closing) + self.reader = _RequestReader(stream) + self.rfile = self.reader + self.wfile = stream + streams.put(stream) connected.put(f"{self.client_address[0]}:{self.client_address[1]}") def finish(self) -> None: @@ -130,6 +241,9 @@ def wire_server( self.wfile.flush() def respond(self) -> None: + if self.reader.cut: + self.close_connection = True + return headers: Final = {name.lower(): value for name, value in self.headers.items()} try: body: Final = read_body(headers, self.rfile) @@ -221,7 +335,7 @@ def wire_server( self.socket = tls.wrap_socket(self.socket, server_side=True) with OwnedHTTPServer(("127.0.0.1", port), Handler) as server: - thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}) + thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": _POLL_SECONDS}) thread.start() try: yield Wire( @@ -236,6 +350,6 @@ def wire_server( server.shutdown() thread.join(timeout=6) assert not thread.is_alive(), "Owned HTTP server survived cleanup" - server.server_close() + _close_unless_stalled(server.server_close, streams, teardown_grace, errors) failure: Final = None if errors.empty() else errors.get_nowait() assert failure is None, f"Owned HTTP peer failed: {failure!r}" diff --git a/tests/integration/database/test_skill_publication.py b/tests/integration/database/test_skill_publication.py new file mode 100644 index 00000000000..476c6bd4399 --- /dev/null +++ b/tests/integration/database/test_skill_publication.py @@ -0,0 +1,201 @@ +import io +import json +import os +import stat +import uuid +import zipfile +from concurrent.futures import ThreadPoolExecutor +from contextlib import AsyncExitStack, ExitStack +from typing import Final + +import psycopg +import pytest +from psycopg import sql + +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.plugin_upload_repository import PluginUploadRepository +from tests.integration._support.client import Gateway, list_value, object_value, string_value +from tests.integration._support.database import read_rows + + +def _archive(name: str) -> bytes: + buffer: Final = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + archive.writestr( + "SKILL.md", f"---\nname: {name}\ndescription: Published instructions\n---\nRetain these exact bytes.\n", + ) + return buffer.getvalue() + + +def _delete(gateway: Gateway, name: str) -> None: + response: Final = gateway.request("DELETE", f"/claude-code/plugins/{name}") + assert response.status_code == 200, response.text + + +@pytest.mark.asyncio +async def test_snapshot_failure_rolls_back_catalog_archive_and_history_then_source_replacement_preserves_bytes( + gateway: Gateway, +) -> None: + name: Final = f"integration-version-{uuid.uuid4().hex}" + path: Final = f"/claude-code/plugins/{name}" + history_path: Final = f"/skill-hub/skills/{name}/versions" + content: Final = _archive(name) + async with AsyncExitStack() as cleanup: + created: Final = gateway.request_multipart( + "/claude-code/plugins/upload", + {"metadata": json.dumps({"name": name})}, + {"file": ("skill.zip", content, "application/zip")}, + ) + assert created.status_code == 200, created.text + cleanup.callback(_delete, gateway, name) + client: Final = PrismaClient(os.environ["DATABASE_URL"], ProxyLogging(UserApiKeyCache())) + await client.connect() + cleanup.push_async_callback(client.disconnect) + repository: Final = PluginUploadRepository(client) + metadata: Final = await repository.get_metadata(name) + assert metadata is not None + snapshot: Final = (await repository.versions(metadata, 100, 0))[0] + snapshot_bytes: Final = await repository.version_content(metadata, snapshot.id) + assert snapshot_bytes is not None + before: Final = gateway.get(path) + history: Final = gateway.get(history_path) + original: Final = object_value(list_value(history["versions"])[0]) + archive_path: Final = string_value(original["download_url"]) + original_response: Final = gateway.request("GET", archive_path) + assert original_response.status_code == 200, original_response.text + published_content: Final = original_response.content + with zipfile.ZipFile(io.BytesIO(published_content)) as published, zipfile.ZipFile(io.BytesIO(content)) as uploaded: + assert published.read("SKILL.md") == uploaded.read("SKILL.md") + current_path: Final = string_value(object_value(before["source"])["url"]) + original_current: Final = gateway.request("GET", current_path) + assert original_current.status_code == 200, original_current.text + witness: Final = f"publication_{uuid.uuid4().hex}" + trigger: Final = f"{witness}_trigger" + function: Final = f"{witness}_function" + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection, ExitStack() as failure: + connection.execute(sql.SQL("CREATE SEQUENCE {}").format(sql.Identifier(witness))) + failure.callback(connection.execute, sql.SQL("DROP SEQUENCE {}").format(sql.Identifier(witness))) + connection.execute(sql.SQL("GRANT USAGE ON SEQUENCE {} TO PUBLIC").format(sql.Identifier(witness))) + connection.execute( + sql.SQL( + "CREATE FUNCTION {}() RETURNS trigger LANGUAGE plpgsql AS $$ BEGIN " + "IF EXISTS (SELECT 1 FROM \"LiteLLM_ClaudeCodePluginTable\" " + "WHERE id = NEW.plugin_id AND name = {}) THEN PERFORM nextval({}); " + "RAISE EXCEPTION 'publication snapshot failure'; END IF; RETURN NEW; END $$" + ).format(sql.Identifier(function), sql.Literal(name), sql.Literal(witness)) + ) + failure.callback(connection.execute, sql.SQL("DROP FUNCTION {}()").format(sql.Identifier(function))) + connection.execute( + sql.SQL('CREATE TRIGGER {} BEFORE INSERT ON "LiteLLM_ClaudeCodePluginVersion" ' + "FOR EACH ROW EXECUTE FUNCTION {}()").format(sql.Identifier(trigger), sql.Identifier(function)) + ) + failure.callback( + connection.execute, + sql.SQL('DROP TRIGGER {} ON "LiteLLM_ClaudeCodePluginVersion"').format(sql.Identifier(trigger)), + ) + rejected: Final = gateway.request( + "PUT", path, {"source": {"source": "github", "repo": "example/changed"}, "description": "must roll back"}, + ) + assert rejected.status_code == 500, rejected.text + assert connection.execute( + sql.SQL("SELECT is_called FROM {}").format(sql.Identifier(witness)) + ).fetchone() == (True,) + assert gateway.get(path) == before + assert gateway.get(history_path) == history + current_archive: Final = gateway.request("GET", current_path) + assert current_archive.status_code == 200 and current_archive.content == original_current.content + updated: Final = gateway.request( + "PUT", path, {"source": {"source": "github", "repo": "example/changed"}}, + ) + assert updated.status_code == 200, updated.text + after: Final = gateway.get(history_path) + versions: Final = tuple(object_value(value) for value in list_value(after["versions"])) + assert [version["version"] for version in versions] == ["1.0.1", "1.0.0"] + assert versions[0]["content_kind"] == "source_reference" + assert versions[1]["content_kind"] == "stored_archive" + retained: Final = gateway.request("GET", archive_path) + assert retained.status_code == 200 and retained.content == published_content + assert read_rows( + 'SELECT plugin_name FROM "LiteLLM_ClaudeCodePluginUpload" WHERE plugin_name=%s', (name,) + ) == [] + assert snapshot in await repository.versions(metadata, 100, 0) + assert await repository.version(metadata, snapshot.id) == snapshot + assert await repository.version_content(metadata, snapshot.id) == snapshot_bytes + hidden: Final = gateway.request("POST", f"{path}/disable") + assert hidden.status_code == 200, hidden.text + assert await repository.versions(metadata, 100, 0) == () + assert await repository.version(metadata, snapshot.id) is None + assert await repository.version_content(metadata, snapshot.id) is None + + +def test_neutral_publication_preserves_portable_files_and_atomic_metadata_history(gateway: Gateway) -> None: + name: Final = f"integration-neutral-{uuid.uuid4().hex}" + path: Final = f"/skill-hub/skills/{name}" + files: Final = { + "SKILL.md": f"---\nname: {name}\ndescription: Portable integration fixture.\n---\n\nRead references/notes.md.\n", + "references/notes.md": "Preserve the original supporting file.\n", + "scripts/check.sh": "#!/bin/sh\nprintf 'portable\\n'\n", + } + buffer: Final = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + for filename, content in files.items(): + entry: Final = zipfile.ZipInfo(filename) + entry.create_system = 3 + entry.external_attr = (stat.S_IFREG | (0o755 if filename.endswith(".sh") else 0o644)) << 16 + archive.writestr(entry, content) + with ExitStack() as cleanup: + created: Final = gateway.request_multipart( + "/skill-hub/skills/upload", {}, {"file": ("skill.zip", buffer.getvalue(), "application/zip")}, + ) + assert created.status_code == 200, created.text + cleanup.callback(_delete, gateway, name) + assert created.json()["version"] == "1.0.0" + before: Final = gateway.get(path) + assert string_value(before["updated_at"]) == created.json()["updated_at"] + history_before: Final = gateway.get(f"{path}/versions") + first: Final = object_value(list_value(history_before["versions"])[0]) + version_id: Final = string_value(first["id"]) + archive_path: Final = f"{path}/versions/{version_id}/archive" + downloaded: Final = gateway.request("GET", archive_path) + assert downloaded.status_code == 200, downloaded.text + with zipfile.ZipFile(io.BytesIO(downloaded.content)) as portable: + assert {member: portable.read(member).decode() for member in portable.namelist()} == files + assert stat.S_IMODE(portable.getinfo("scripts/check.sh").external_attr >> 16) == 0o755 + assert gateway.get(f"{path}/versions/{version_id}") == first + + def publish(index: int) -> int: + response: Final = gateway.request( + "PUT", path, + {"description": f"Registry edit {index}", "expected_updated_at": before["updated_at"]}, + ) + if response.status_code == 200: + assert string_value(response.json()["updated_at"]) != before["updated_at"] + return response.status_code + + with ThreadPoolExecutor(max_workers=2) as pool: + statuses: Final = tuple(pool.map(publish, range(2))) + assert sorted(statuses) == [200, 409] + assert publish(2) == 409 + current: Final = gateway.get(path) + history_after: Final = gateway.get(f"{path}/versions") + versions: Final = tuple(object_value(value) for value in list_value(history_after["versions"])) + assert [version["version"] for version in versions] == ["1.0.1", "1.0.0"] + latest: Final = versions[0] + manifest: Final = object_value(latest["manifest"]) + assert current["version"] == manifest["version"] == latest["version"] + assert current["description"] == manifest["description"] == latest["description"] + assert manifest["agent_skill"] == object_value(first["manifest"])["agent_skill"] + stored: Final = read_rows( + 'SELECT version, description, manifest_json FROM "LiteLLM_ClaudeCodePluginTable" WHERE name=%s', (name,) + ) + assert len(stored) == 1 + assert stored[0]["version"] == current["version"] + assert json.loads(string_value(stored[0]["manifest_json"])) == manifest + retained: Final = gateway.request("GET", archive_path) + latest_archive: Final = gateway.request("GET", f"{path}/versions/{latest['id']}/archive") + assert retained.content == latest_archive.content == downloaded.content + rejected: Final = gateway.request("PUT", path, {"version": "1.0.0", "description": "rejected"}) + assert rejected.status_code == 400, rejected.text + assert gateway.get(path) == current + assert gateway.get(f"{path}/versions") == history_after diff --git a/tests/integration/database/test_skill_submissions.py b/tests/integration/database/test_skill_submissions.py new file mode 100644 index 00000000000..cc00bf6e359 --- /dev/null +++ b/tests/integration/database/test_skill_submissions.py @@ -0,0 +1,454 @@ +import asyncio +import hashlib +import io +import os +import uuid +import zipfile +from collections.abc import AsyncIterator, Mapping, Sequence +from dataclasses import dataclass, replace +from typing import Final, Literal + +import psycopg +import pytest +import pytest_asyncio +from pydantic import JsonValue, TypeAdapter + +from litellm.models.plugin_upload import ( + PluginPublication, + PluginPublicationBase, + PluginPublicationError, + PluginUploadStore, + PublishedPlugin, +) +from litellm.models.skill_submission import ( + SkillSubmission, + SkillSubmissionStore, + SubmissionConflict, + SubmissionLimitExceeded, + SubmissionReviewed, +) +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.plugin_upload_repository import PluginUploadRepository +from litellm.repositories.skill_submission_repository import SkillSubmissionRepository +from tests.integration._support.client import eventually +from tests.integration._support.database import read_rows, write_rows + + +@dataclass(frozen=True, slots=True) +class PendingUpdate: + publications: PluginUploadStore + submissions: SkillSubmissionStore + original: PluginPublication + proposal: PluginPublication + base: PublishedPlugin + submission: SkillSubmission + + +def _publication(name: str, description: str, base_url: str = "https://skills.example.test") -> PluginPublication: + buffer: Final = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + archive.writestr("SKILL.md", f"---\nname: {name}\ndescription: {description}\n---\n{description}\n") + content: Final = buffer.getvalue() + digest: Final = hashlib.sha256(content).hexdigest() + return PluginPublication( + name=name, + version=None, + description=description, + manifest={ + "name": name, + "description": description, + "extension": {"retained": True}, + "source": { + "source": "archive", + "url": f"{base_url}/claude-code/plugins/{name}/archive/{digest}.zip", + "sha256": digest, + }, + }, + content=content, + created_by=f"{name}-owner", + ) + + +def _published_rows(name: str) -> Sequence[Mapping[str, JsonValue]]: + return read_rows( + """SELECT 'catalog' AS storage, to_jsonb(p) AS value + FROM "LiteLLM_ClaudeCodePluginTable" p WHERE name=%s + UNION ALL + SELECT 'archive', to_jsonb(a) FROM "LiteLLM_ClaudeCodePluginUpload" a WHERE plugin_name=%s + UNION ALL + SELECT 'history', to_jsonb(v) FROM "LiteLLM_ClaudeCodePluginVersion" v + JOIN "LiteLLM_ClaudeCodePluginTable" p ON p.id=v.plugin_id WHERE p.name=%s + ORDER BY storage, value""", + (name, name, name), + ) + + +@pytest_asyncio.fixture(loop_scope="function") +async def pending_update() -> AsyncIterator[PendingUpdate]: + name: Final = f"integration-submission-{uuid.uuid4().hex}" + client: Final = PrismaClient(os.environ["DATABASE_URL"], ProxyLogging(UserApiKeyCache())) + await client.connect() + try: + publications: Final = PluginUploadRepository(client) + submissions: Final = SkillSubmissionRepository(client) + original: Final = _publication(name, "Published instructions") + base: Final = await publications.publish(original) + assert isinstance(base, PublishedPlugin) + proposal: Final = _publication(name, "Proposed instructions", "https://member.example.test") + submission: Final = await submissions.submit( + proposal, f"{name}-owner", PluginPublicationBase(base.id, base.updated_at) + ) + assert isinstance(submission, SkillSubmission) + yield PendingUpdate(publications, submissions, original, proposal, base, submission) + finally: + try: + write_rows('DELETE FROM "LiteLLM_SkillSubmission" WHERE plugin_name=%s', (name,)) + write_rows('DELETE FROM "LiteLLM_ClaudeCodePluginTable" WHERE name=%s', (name,)) + finally: + await client.disconnect() + + +async def test_approval_commits_version_archive_history_and_review(pending_update: PendingUpdate) -> None: + case: Final = pending_update + before: Final = await case.publications.get_metadata(case.original.name) + assert before is not None + assert await case.publications.get_content(before) == case.original.content + assert case.submission.proposal.version is None + reviewed: Final = await case.submissions.review( + case.submission.id, "integration-reviewer", True, None, publication_base_url="https://skills.example.test" + ) + assert isinstance(reviewed, SubmissionReviewed) + assert reviewed.plugin is not None and reviewed.plugin.version == "1.0.1" + stored: Final = await case.submissions.get(case.submission.id, None) + assert stored is not None + assert (stored.status, stored.reviewed_by, stored.rejection_reason) == ( + "approved", + "integration-reviewer", + None, + ) + assert stored.reviewed_at is not None + assert stored.proposal == case.submission.proposal + assert await case.submissions.get_archive(stored.id, None) == case.proposal.content + current: Final = await case.publications.get_plugin(case.original.name) + assert current is not None + published_manifest: Final = { + **case.proposal.manifest, + "source": { + **TypeAdapter(dict[str, str]).validate_python(case.submission.proposal.manifest["source"]), + "url": f"https://skills.example.test/claude-code/plugins/{case.original.name}/archive/" + f"{hashlib.sha256(case.proposal.content or b'').hexdigest()}.zip", + }, + } + assert (current.id, current.version, current.description, current.manifest) == ( + case.base.id, + reviewed.plugin.version, + case.proposal.description, + {**published_manifest, "version": reviewed.plugin.version}, + ) + metadata: Final = await case.publications.get_metadata(case.original.name) + assert metadata is not None + assert await case.publications.get_content(metadata) == case.proposal.content + history: Final = await case.publications.versions(metadata, 10, 0) + assert tuple((version.version, version.description, version.archive_available) for version in history) == ( + ("1.0.1", case.proposal.description, True), + ("1.0.0", case.original.description, True), + ) + for version, publication, manifest in zip( + history, (case.proposal, case.original), (published_manifest, case.original.manifest), strict=True + ): + assert version.manifest == {**manifest, "version": version.version} + assert version.published_by == publication.created_by + assert await case.publications.version_content(metadata, version.id) == publication.content + + +async def test_simultaneous_approve_reject_has_one_winner(pending_update: PendingUpdate) -> None: + case: Final = pending_update + before: Final = _published_rows(case.original.name) + with psycopg.connect(os.environ["DATABASE_URL"]) as blocker: + blocker.execute('SELECT id FROM "LiteLLM_SkillSubmission" WHERE id=%s FOR UPDATE', (case.submission.id,)) + blocker_pid: Final = str(blocker.info.backend_pid) + reviews: Final = ( + asyncio.create_task( + case.submissions.review( + case.submission.id, + "integration-approver", + True, + None, + publication_base_url="https://skills.example.test", + ) + ), + asyncio.create_task( + case.submissions.review(case.submission.id, "integration-rejector", False, "Needs revision") + ), + ) + try: + await asyncio.to_thread( + eventually, + lambda: read_rows( + """WITH RECURSIVE blocked AS ( + SELECT %s::integer AS pid + UNION + SELECT a.pid FROM pg_stat_activity a + JOIN blocked b ON b.pid = ANY(pg_blocking_pids(a.pid)) + ) SELECT pid FROM blocked WHERE pid<>%s::integer""", + (blocker_pid, blocker_pid), + ), + lambda rows: len(rows) == 2, + seconds=3, + ) + finally: + blocker.rollback() + results: Final = await asyncio.gather(*reviews) + winners: Final = tuple(result for result in results if isinstance(result, SubmissionReviewed)) + assert len(winners) == 1 + assert sum(isinstance(result, SubmissionConflict) for result in results) == 1 + stored: Final = await case.submissions.get(case.submission.id, None) + assert stored is not None and stored.reviewed_at is not None + winner: Final = winners[0] + assert (stored.status, stored.reviewed_by, stored.rejection_reason) == ( + winner.submission.status, + winner.submission.reviewed_by, + winner.submission.rejection_reason, + ) + if isinstance(results[0], SubmissionConflict): + assert (stored.status, stored.reviewed_by, stored.rejection_reason) == ( + "rejected", + "integration-rejector", + "Needs revision", + ) + assert _published_rows(case.original.name) == before + return + assert (stored.status, stored.reviewed_by, stored.rejection_reason) == ( + "approved", + "integration-approver", + None, + ) + metadata: Final = await case.publications.get_metadata(case.original.name) + assert metadata is not None + assert await case.publications.get_content(metadata) == case.proposal.content + assert tuple(version.version for version in await case.publications.versions(metadata, 10, 0)) == ( + "1.0.1", + "1.0.0", + ) + + +async def test_publication_failure_rolls_back_and_leaves_submission_pending(pending_update: PendingUpdate) -> None: + case: Final = pending_update + collision_id: Final = str(uuid.uuid4()) + write_rows( + """INSERT INTO "LiteLLM_ClaudeCodePluginVersion" (id, plugin_id, version) + VALUES (%s, %s, '1.0.1')""", + (collision_id, case.base.id), + ) + before: Final = _published_rows(case.original.name) + failed: Final = await case.submissions.review( + case.submission.id, "integration-reviewer", True, None, publication_base_url="https://skills.example.test" + ) + assert isinstance(failed, SubmissionConflict) + assert _published_rows(case.original.name) == before + assert await case.submissions.get(case.submission.id, None) == case.submission + assert await case.submissions.get_archive(case.submission.id, None) == case.proposal.content + write_rows('DELETE FROM "LiteLLM_ClaudeCodePluginVersion" WHERE id=%s', (collision_id,)) + retried: Final = await case.submissions.review( + case.submission.id, "integration-reviewer", True, None, publication_base_url="https://skills.example.test" + ) + assert isinstance(retried, SubmissionReviewed) + assert retried.plugin is not None and retried.plugin.version == "1.0.1" + stored: Final = await case.submissions.get(case.submission.id, None) + assert stored is not None and stored.status == "approved" + + +@pytest.mark.parametrize("change", ("update", "recreate")) +async def test_stale_target_cannot_publish( + pending_update: PendingUpdate, change: Literal["update", "recreate"] +) -> None: + case: Final = pending_update + intervening: Final = _publication(case.original.name, "Intervening publication") + if change == "recreate": + write_rows('DELETE FROM "LiteLLM_ClaudeCodePluginTable" WHERE name=%s', (case.original.name,)) + replacement: Final = await case.publications.publish( + intervening, + PluginPublicationBase(case.base.id, case.base.updated_at) if change == "update" else None, + ) + assert isinstance(replacement, PublishedPlugin) + before: Final = _published_rows(case.original.name) + reviewed: Final = await case.submissions.review( + case.submission.id, "integration-reviewer", True, None, publication_base_url="https://skills.example.test" + ) + assert isinstance(reviewed, PluginPublicationError) and reviewed.kind == "conflict" + assert _published_rows(case.original.name) == before + assert await case.submissions.get(case.submission.id, None) == case.submission + + +async def test_concurrent_submissions_share_user_limit_and_review_releases_capacity( + pending_update: PendingUpdate, +) -> None: + case: Final = pending_update + owner: Final = case.submission.submitted_by + base: Final = PluginPublicationBase(case.base.id, case.base.updated_at) + for _ in range(18): + assert isinstance(await case.submissions.submit(case.proposal, owner, base), SkillSubmission) + peer: Final = PrismaClient(os.environ["DATABASE_URL"], ProxyLogging(UserApiKeyCache())) + await peer.connect() + try: + with psycopg.connect(os.environ["DATABASE_URL"]) as blocker: + blocker.execute( + "SELECT pg_advisory_xact_lock(hashtext('litellm.skill_submissions'), hashtext(%s))", (owner,) + ) + blocker_pid: Final = str(blocker.info.backend_pid) + requests: Final = ( + asyncio.create_task(case.submissions.submit(case.proposal, owner, base)), + asyncio.create_task(SkillSubmissionRepository(peer).submit(case.proposal, owner, base)), + ) + try: + await asyncio.to_thread( + eventually, + lambda: read_rows( + "SELECT pid FROM pg_stat_activity WHERE %s::integer = ANY(pg_blocking_pids(pid))", + (blocker_pid,), + ), + lambda rows: len(rows) == 2, + seconds=3, + ) + finally: + blocker.rollback() + results: Final = await asyncio.gather(*requests) + finally: + await peer.disconnect() + assert sum(isinstance(result, SkillSubmission) for result in results) == 1 + assert sum(isinstance(result, SubmissionLimitExceeded) for result in results) == 1 + full: Final = await case.submissions.list(owner, "pending_review", 20, 0) + assert (len(full.submissions), full.has_more) == (20, False) + other: Final = await case.submissions.submit(case.proposal, f"{owner}-other", base) + assert isinstance(other, SkillSubmission) + rejected: Final = await case.submissions.review( + case.submission.id, "integration-reviewer", False, "Review complete" + ) + assert isinstance(rejected, SubmissionReviewed) + assert await case.submissions.get_archive(case.submission.id, owner) == case.proposal.content + replacement: Final = await case.submissions.submit(case.proposal, owner, base) + assert isinstance(replacement, SkillSubmission) + after: Final = await case.submissions.list(owner, "pending_review", 20, 0) + assert len(after.submissions) == 20 + assert all(row.id != case.submission.id and row.submitted_by == owner for row in after.submissions) + + +async def test_submission_pages_filter_owner_and_status_before_limit(pending_update: PendingUpdate) -> None: + case: Final = pending_update + owner: Final = case.submission.submitted_by + other: Final = f"{owner}-other" + generated: Final = await asyncio.gather( + *( + case.submissions.submit(case.proposal, identity, PluginPublicationBase(case.base.id, case.base.updated_at)) + for identity in (owner, other, owner, other) + ) + ) + submitted: Final = tuple(result for result in generated if isinstance(result, SkillSubmission)) + assert len(submitted) == 4 + rejected: Final = await case.submissions.review(submitted[0].id, "integration-reviewer", False, None) + assert isinstance(rejected, SubmissionReviewed) + write_rows( + """UPDATE "LiteLLM_SkillSubmission" SET created_at=TIMESTAMP '9999-01-01' + WHERE plugin_name=%s""", + (case.original.name,), + ) + own_pending: Final = tuple(sorted((case.submission.id, submitted[2].id), reverse=True)) + first: Final = await case.submissions.list(owner, "pending_review", 1, 0) + second: Final = await case.submissions.list(owner, "pending_review", 1, 1) + empty: Final = await case.submissions.list(owner, "pending_review", 1, 2) + assert (tuple(row.id for row in first.submissions), first.has_more) == (own_pending[:1], True) + assert (tuple(row.id for row in second.submissions), second.has_more) == (own_pending[1:], False) + assert (empty.submissions, empty.has_more) == ((), False) + reviewed: Final = await case.submissions.list(owner, "rejected", 20, 0) + assert (tuple(row.id for row in reviewed.submissions), reviewed.has_more) == ((submitted[0].id,), False) + own_all: Final = await case.submissions.list(owner, None, 20, 0) + assert tuple(row.id for row in own_all.submissions) == tuple(sorted((*own_pending, submitted[0].id), reverse=True)) + other_page: Final = await case.submissions.list(other, None, 20, 0) + assert tuple(row.id for row in other_page.submissions) == tuple( + sorted((submitted[1].id, submitted[3].id), reverse=True) + ) + global_pending: Final = await case.submissions.list(None, "pending_review", 3, 0) + all_pending: Final = tuple(sorted((*own_pending, submitted[1].id, submitted[3].id), reverse=True)) + assert (tuple(row.id for row in global_pending.submissions), global_pending.has_more) == (all_pending[:3], True) + + +@pytest.mark.parametrize("completion_seconds", (13, 21), ids=("before-deadline", "at-deadline")) +async def test_scan_claim_retry_and_late_result_preserve_proposal( + pending_update: PendingUpdate, completion_seconds: int, +) -> None: + from datetime import datetime, timedelta, timezone + + from litellm.models.skill_submission import SkillScanInput + from litellm.types.proxy.skill_scanning import SkillScan, SkillScanReport + + case: Final = pending_update + now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + first: Final = SkillScan( + id="first", status="running", requested_by="reviewer", started_at=now, + expires_at=now + timedelta(seconds=10), + ) + started: Final = await case.submissions.start_scan(case.submission.id, first) + assert isinstance(started, SkillScanInput) + assert started.content == case.proposal.content + assert started.scan.archive_sha256 == hashlib.sha256(case.proposal.content or b"").hexdigest() + duplicate: Final = await case.submissions.start_scan(case.submission.id, first.model_copy(update={"id": "duplicate"})) + assert isinstance(duplicate, SubmissionConflict) + second: Final = first.model_copy(update={ + "id": "second", "started_at": now + timedelta(seconds=11), "expires_at": now + timedelta(seconds=21), + }) + retry: Final = await case.submissions.start_scan(case.submission.id, second) + assert isinstance(retry, SkillScanInput) + report: Final = SkillScanReport( + scanner_version="test-version", complete=False, scan_mode="static-only", risk_score=90, findings=(), + ) + await case.submissions.finish_scan(case.submission.id, started.scan.model_copy(update={ + "status": "completed", "report": report, "completed_at": now + timedelta(seconds=12), + })) + assert await case.submissions.get_scan(case.submission.id) == retry.scan + completed: Final = retry.scan.model_copy(update={ + "status": "completed", "report": report, "completed_at": now + timedelta(seconds=completion_seconds), + }) + expected: Final = completed if completion_seconds < 21 else retry.scan.at(now + timedelta(seconds=21)) + await case.submissions.finish_scan(case.submission.id, completed) + assert await case.submissions.get_scan(case.submission.id) == expected + assert await case.submissions.get(case.submission.id, None) == case.submission + approved: Final = await case.submissions.review(case.submission.id, "reviewer", True, None, publication_base_url="https://gateway.test") + assert isinstance(approved, SubmissionReviewed) + assert await case.submissions.get_scan(case.submission.id) == expected + denied: Final = await case.submissions.start_scan(case.submission.id, second) + assert isinstance(denied, SubmissionConflict) + + +@pytest.mark.parametrize("available, remote", ((True, False), (False, False), (False, True))) +async def test_metadata_scan_requires_the_matching_managed_archive( + pending_update: PendingUpdate, available: bool, remote: bool +) -> None: + from datetime import datetime, timedelta, timezone + + from litellm.models.skill_submission import SkillScanInput + from litellm.types.proxy.skill_scanning import SkillScan + + case: Final = pending_update + submission: Final = await case.submissions.submit( + replace(case.original, content=None, manifest={ + **case.original.manifest, + **({"source": {"source": "archive", "url": "https://example.org/skill.zip"}} if remote else {}), + }), "member", PluginPublicationBase(case.base.id, case.base.updated_at) + ) + if not available: + write_rows('DELETE FROM "LiteLLM_ClaudeCodePluginUpload" WHERE plugin_name=%s', (case.original.name,)) + now: Final = datetime.now(timezone.utc) + result: Final = await case.submissions.start_scan(submission.id, SkillScan( + id="metadata-scan", status="running", requested_by="admin", started_at=now, + expires_at=now + timedelta(seconds=10), + )) + if available or remote: + assert isinstance(result, SkillScanInput) + assert result.content == (None if remote else case.original.content) + if remote: + assert result.source == {"source": "archive", "url": "https://example.org/skill.zip"} + assert result.scan.archive_sha256 is None + else: + assert isinstance(result, SubmissionConflict) + assert await case.submissions.get_scan(submission.id) is None diff --git a/tests/integration/database/test_spend_lock_order.py b/tests/integration/database/test_spend_lock_order.py new file mode 100644 index 00000000000..a7e00b7fa87 --- /dev/null +++ b/tests/integration/database/test_spend_lock_order.py @@ -0,0 +1,197 @@ +import asyncio +import subprocess +import sys +from collections.abc import AsyncIterator, Callable, Iterator, Sequence +from pathlib import Path +from types import SimpleNamespace +from typing import Final + +import psycopg +import pytest +import pytest_asyncio +from prisma import Prisma + +from litellm.proxy._types import DBSpendUpdateTransactions +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.management_helpers.model_allowlist_rename_sync import sync_model_allowlists_for_renamed_model +from litellm.repositories import unit_of_work +from tests.integration._support.database import scratch_database + +_ROOT: Final = Path(__file__).resolve().parents[3] +_KEY: Final = "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683" +_CHARGES: Final = DBSpendUpdateTransactions( + user_list_transactions={"user": 1.0}, + end_user_list_transactions={}, + key_list_transactions={_KEY: 1.0}, + team_list_transactions={"team": 1.0}, + team_member_list_transactions={"team_id::team::user_id::user": 1.0}, + org_list_transactions={"org": 1.0}, + org_member_list_transactions={"org_id::org::user_id::user": 1.0}, + project_list_transactions={"project": 1.0}, + tag_list_transactions={"tag": 1.0}, + agent_list_transactions={}, + model_access_group_list_transactions={"group": 1.0}, +) +_SEED: Final = ( + "INSERT INTO \"LiteLLM_BudgetTable\" (budget_id, created_by, updated_by) VALUES ('budget', 't', 't')", + 'INSERT INTO "LiteLLM_OrganizationTable" (organization_id, organization_alias, budget_id, created_by, updated_by) ' + "VALUES ('org', 'org', 'budget', 't', 't')", + "INSERT INTO \"LiteLLM_UserTable\" (user_id) VALUES ('user')", + 'INSERT INTO "LiteLLM_TeamTable" (team_id, members_with_roles) ' + """VALUES ('team', '[{"user_id": "user", "role": "user"}]')""", + "INSERT INTO \"LiteLLM_TeamMembership\" (team_id, user_id, budget_id) VALUES ('team', 'user', 'budget')", + "INSERT INTO \"LiteLLM_OrganizationMembership\" (organization_id, user_id) VALUES ('org', 'user')", + f"INSERT INTO \"LiteLLM_VerificationToken\" (token, budget_id) VALUES ('{_KEY}', 'budget')", + "INSERT INTO \"LiteLLM_ProjectTable\" (project_id, budget_id, created_by, updated_by) " + "VALUES ('project', 'budget', 't', 't')", + "INSERT INTO \"LiteLLM_TagTable\" (tag_name, budget_id) VALUES ('tag', 'budget')", + "INSERT INTO \"LiteLLM_ModelAccessGroupBudgetTable\" (access_group_name, budget_id) VALUES ('group', 'budget')", +) + + +@pytest.fixture(scope="module") +def schema_sql() -> str: + return subprocess.run( + [sys.executable, "-I", "-m", "prisma", "migrate", "diff", "--from-empty", "--to-schema-datamodel", + str(_ROOT / "schema.prisma"), "--script"], + check=True, + capture_output=True, + text=True, + cwd=_ROOT, + ).stdout + + +@pytest.fixture +def database(schema_sql: str) -> Iterator[str]: + with scratch_database("spend_lock_order") as url, psycopg.connect(url, autocommit=True) as connection: + connection.execute(schema_sql) + for statement in _SEED: + connection.execute(statement) + yield url + + +@pytest_asyncio.fixture(loop_scope="function") +async def ledger(database: str) -> AsyncIterator[Prisma]: + client: Final = Prisma(datasource={"url": database}) + await client.connect() + try: + yield client + finally: + await client.disconnect() + + +async def _apply_charges(client: Prisma) -> None: + async with client.tx() as transaction: + await DBSpendUpdateWriter().apply_accounting_entities(client, None, transaction, _CHARGES) + + +async def _ledger_waits_for_a_lock(url: str) -> None: + with psycopg.connect(url, autocommit=True) as observer: + for _ in range(200): + if observer.execute("SELECT 1 FROM pg_locks WHERE NOT granted").fetchone(): + return + await asyncio.sleep(0.05) + raise AssertionError("the ledger never waited on the other writer") + + +def _spend(url: str) -> Sequence[object]: + with psycopg.connect(url) as connection: + return connection.execute( + 'SELECT (SELECT spend FROM "LiteLLM_UserTable"), (SELECT spend FROM "LiteLLM_TeamTable"), ' + '(SELECT spend FROM "LiteLLM_TeamMembership"), (SELECT spend FROM "LiteLLM_VerificationToken"), ' + '(SELECT spend FROM "LiteLLM_OrganizationTable"), (SELECT spend FROM "LiteLLM_ProjectTable")' + ).fetchone() or () + + +_Writer = Callable[[psycopg.Connection[tuple[object, ...]]], None] + + +def _roster_first(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute("SELECT pg_advisory_xact_lock(hashtext('team'))") + + +def _roster_rest(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute("SELECT 1 FROM \"LiteLLM_UserTable\" WHERE user_id = 'user' FOR UPDATE") + connection.execute("UPDATE \"LiteLLM_TeamTable\" SET members_with_roles = members_with_roles WHERE team_id = 'team'") + + +def _access_group_first(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute("UPDATE \"LiteLLM_TeamTable\" SET access_group_ids = ARRAY['ag'] WHERE team_id = 'team'") + + +def _access_group_rest(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute(f"UPDATE \"LiteLLM_VerificationToken\" SET access_group_ids = ARRAY['ag'] WHERE token = '{_KEY}'") + + +def _cascade_first(connection: psycopg.Connection[tuple[object, ...]]) -> None: + connection.execute(unit_of_work._LOCK_BUDGET_ORGS_SQL.replace("$1", "%s"), (["budget"],)) + + +def _cascade_rest(connection: psycopg.Connection[tuple[object, ...]]) -> None: + for statement in ( + unit_of_work._LOCK_BUDGET_MEMBERSHIPS_SQL, + unit_of_work._LOCK_BUDGET_KEYS_SQL, + unit_of_work._LOCK_BUDGET_PROJECTS_SQL, + unit_of_work._LOCK_BUDGET_TAGS_SQL, + unit_of_work._LOCK_BUDGET_ACCESS_GROUPS_SQL, + ): + connection.execute(statement.replace("$1", "%s"), (["budget"],)) + connection.execute("UPDATE \"LiteLLM_TagTable\" SET spend = 0 WHERE budget_id = 'budget'") + + +@pytest.mark.parametrize( + ("first", "rest"), + [(_roster_first, _roster_rest), (_access_group_first, _access_group_rest), (_cascade_first, _cascade_rest)], + ids=["team_roster", "access_group", "budget_cascade"], +) +@pytest.mark.asyncio(loop_scope="function") +async def test_charges_wait_for_a_concurrent_writer_instead_of_deadlocking( + database: str, ledger: Prisma, first: _Writer, rest: _Writer +) -> None: + """The other writer takes its first lock, the ledger starts and must block holding nothing the other writer + needs next, so the other writer finishes and the charge lands once after it.""" + with psycopg.connect(database) as writer: + writer.execute("SET lock_timeout = '10s'") + first(writer) + charges: Final = asyncio.create_task(_apply_charges(ledger)) + await _ledger_waits_for_a_lock(database) + await asyncio.to_thread(rest, writer) + writer.commit() + await asyncio.wait_for(charges, timeout=20) + + assert _spend(database) == (1.0, 1.0, 1.0, 1.0, 1.0, 1.0) + + +async def _rename(client: Prisma) -> None: + await sync_model_allowlists_for_renamed_model( + SimpleNamespace(db=client), + model_id="deployment", + old_name="old-model", + new_name="new-model", + llm_router=None, + user_api_key_cache=UserApiKeyCache(), + ) + + +@pytest.mark.asyncio(loop_scope="function") +async def test_a_model_rename_waits_for_a_roster_change_instead_of_deadlocking(database: str, ledger: Prisma) -> None: + """The roster change holds the team lock and the user row, so the rename must wait before it locks the team + row the roster change updates next.""" + with psycopg.connect(database, autocommit=True) as seed: + seed.execute("UPDATE \"LiteLLM_TeamTable\" SET models = ARRAY['old-model']") + seed.execute("UPDATE \"LiteLLM_UserTable\" SET models = ARRAY['old-model']") + with psycopg.connect(database) as writer: + writer.execute("SET lock_timeout = '10s'") + _roster_first(writer) + writer.execute("SELECT 1 FROM \"LiteLLM_UserTable\" WHERE user_id = 'user' FOR UPDATE") + rename: Final = asyncio.create_task(_rename(ledger)) + await _ledger_waits_for_a_lock(database) + await asyncio.to_thread(_roster_rest, writer) + writer.commit() + await asyncio.wait_for(rename, timeout=20) + + with psycopg.connect(database) as connection: + assert connection.execute( + 'SELECT (SELECT models FROM "LiteLLM_TeamTable"), (SELECT models FROM "LiteLLM_UserTable")' + ).fetchone() == (["new-model"], ["new-model"]) diff --git a/tests/integration/database/test_spendlogs_delivery.py b/tests/integration/database/test_spendlogs_delivery.py index 19f01a6b902..d1227e85565 100644 --- a/tests/integration/database/test_spendlogs_delivery.py +++ b/tests/integration/database/test_spendlogs_delivery.py @@ -131,9 +131,9 @@ async def test_duplicate_rows_in_one_flush_keep_first_accounting_payload(deliver "accounting_log_disabled": mode == "disabled", } assert await repository.insert_with_deliveries((first, later), write_spend_logs=mode != "url") == 1 - assert delivery_database.execute( - 'SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"' - ).fetchone() == (first["accounting_payload"],) + assert delivery_database.execute('SELECT accounting_payload FROM "LiteLLM_SpendLogDelivery"').fetchone() == ( + first["accounting_payload"], + ) assert delivery_database.execute('SELECT spend FROM "LiteLLM_SpendLogs"').fetchall() == ( [(first["spend"],)] if mode == "local" else [] ) @@ -415,11 +415,28 @@ def test_delivery_migration_preserves_unused_old_receipts(state): assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,) +@pytest.mark.parametrize( + ("payload", "huge_state"), + [ + ( + "jsonb_set(accounting_payload, '{metadata}', ('{\"tokens_saved\": ' || repeat('9', 401) || '}')::jsonb)", + "applied", + ), + ( + "jsonb_set(accounting_payload, '{metadata}', ('{\"tokens_saved\": ' || repeat('9', 5000) || '}')::jsonb)", + "parked", + ), + ("'[]'::jsonb", "parked"), + ], + ids=["beyond_float", "beyond_python_int_parse", "not_an_object"], +) @pytest.mark.asyncio -async def test_malformed_payload_is_parked_without_losing_other_deliveries(delivery_database, delivery_client): - _insert(delivery_database, (_row("bad"), _row("good"))) +async def test_a_payload_the_engine_cannot_decode_never_blocks_other_deliveries( + delivery_database, delivery_client, payload, huge_state +): + _insert(delivery_database, (_row("huge"), _row("good"))) delivery_database.execute( - 'UPDATE "LiteLLM_SpendLogDelivery" SET accounting_payload=$$ [] $$::jsonb WHERE request_id=$$bad$$' + f'UPDATE "LiteLLM_SpendLogDelivery" SET accounting_payload = {payload} WHERE request_id = $$huge$$' ) async def apply(transaction, deliveries): @@ -428,12 +445,16 @@ async def test_malformed_payload_is_parked_without_losing_other_deliveries(deliv ) repository = AccountingDeliveryRepository(delivery_client) - assert (await repository.apply_pending(apply, max_attempts=2)).failed == 2 - assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.0,) - delivery_database.execute('UPDATE "LiteLLM_SpendLogDelivery" SET next_attempt_at=CURRENT_TIMESTAMP') - assert (await repository.apply_pending(apply, max_attempts=2)).parked == ("bad",) - assert (await repository.apply_pending(apply, max_attempts=2)).applied == 1 - assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5,) + results = [await repository.apply_pending(apply, max_attempts=1) for _ in range(3)] + _insert(delivery_database, (_row("later"),)) + results.append(await repository.apply_pending(apply, max_attempts=1)) + + applied: Final = 3 if huge_state == "applied" else 2 + assert sum(result.applied for result in results) == applied + assert delivery_database.execute("SELECT spend FROM totals").fetchone() == (0.5 * applied,) + assert delivery_database.execute( + 'SELECT request_id, state::text FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id' + ).fetchall() == [("good", "applied"), ("huge", huge_state), ("later", "applied")] @pytest.mark.asyncio @@ -463,19 +484,29 @@ async def test_ledger_batch_retrieves_charge_once_in_every_logging_mode( from litellm.proxy.utils import jsonify_object completed = { - **_row("batch_completed"), "call_type": "aretrieve_batch", "status": "success", "spend": 0.5, + **_row("batch_completed"), + "call_type": "aretrieve_batch", + "status": "success", + "spend": 0.5, "metadata": '{"status":"success","private":"request detail"}', "accounting_payload": { - "request_id": "batch_completed", "call_type": "aretrieve_batch", "status": "success", "spend": 0.5, + "request_id": "batch_completed", + "call_type": "aretrieve_batch", + "status": "success", + "spend": 0.5, "metadata": {"accounting_savings": {"compression": 0.2}}, }, } monkeypatch.setattr(proxy_server, "general_settings", {"spend_logs_metadata_fields": {"include": ["status"]}}) client = SimpleNamespace(db=delivery_client, writer_db=delivery_client, jsonify_object=jsonify_object) writer = DBSpendUpdateWriter() - first = await writer._claim_settle_once_cost_row(payload=completed, prisma_client=client, disable_spend_logs=disabled) + first = await writer._claim_settle_once_cost_row( + payload=completed, prisma_client=client, disable_spend_logs=disabled + ) again = await writer._claim_settle_once_cost_row( - payload={**completed, "startTime": "2026-11-05T12:00:00+00:00"}, prisma_client=client, disable_spend_logs=disabled + payload={**completed, "startTime": "2026-11-05T12:00:00+00:00"}, + prisma_client=client, + disable_spend_logs=disabled, ) assert (first, again) == (True, False) assert delivery_database.execute('SELECT spend, metadata FROM "LiteLLM_SpendLogs"').fetchall() == ( diff --git a/tests/integration/database/test_spendlogs_ledger_parity.py b/tests/integration/database/test_spendlogs_ledger_parity.py index e1894aa66bf..60f3c074cca 100644 --- a/tests/integration/database/test_spendlogs_ledger_parity.py +++ b/tests/integration/database/test_spendlogs_ledger_parity.py @@ -63,10 +63,10 @@ _DAILY_TABLES: Final = ( ) -@pytest.mark.parametrize("batch_size,row_count,expected,distinct_models", [(None, 251, 251, False), ("200", 251, 200, False), (None, 10001, 10000, False), (None, 3001, 3001, True)]) +@pytest.mark.parametrize("batch_size,row_count,transactions,distinct_models", [(None, 251, 1, False), (None, 1001, 3, False), ("200", 251, 2, False), ("3001", 3001, 1, True)]) @pytest.mark.asyncio -async def test_worker_applies_the_configured_backlog_in_one_transaction( - parity_database: str, monkeypatch: pytest.MonkeyPatch, batch_size: str | None, row_count: int, expected: int, distinct_models: bool +async def test_worker_drains_a_backlog_in_one_run_one_configured_batch_per_transaction( + parity_database: str, monkeypatch: pytest.MonkeyPatch, batch_size: str | None, row_count: int, transactions: int, distinct_models: bool ): from litellm.accounting.worker import flush_accounting_deliveries @@ -92,13 +92,37 @@ async def test_worker_applies_the_configured_backlog_in_one_transaction( )) == row_count - 1 started = time.monotonic() await flush_accounting_deliveries(writer, client, proxy_server.proxy_logging_obj) - print(f"Ledger batch applied={expected} queued={row_count} seconds={time.monotonic() - started:.3f}") + print(f"Ledger backlog queued={row_count} seconds={time.monotonic() - started:.3f}") with psycopg.connect(parity_database) as connection: - assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'applied\'').fetchone() == (expected,) - assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'pending\'').fetchone() == (row_count - expected,) - assert connection.execute('SELECT COUNT(DISTINCT applied_at) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (1,) - assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (expected * 0.25,) - assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_DailyUserSpend"').fetchone() == (expected if distinct_models else 1,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'applied\'').fetchone() == (row_count,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE state=\'pending\'').fetchone() == (0,) + assert connection.execute('SELECT COUNT(DISTINCT applied_at) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (transactions,) + assert connection.execute('SELECT spend FROM "LiteLLM_UserTable" WHERE user_id=\'user\'').fetchone() == (row_count * 0.25,) + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_DailyUserSpend"').fetchone() == (row_count if distinct_models else 1,) + finally: + await client.db.disconnect() + + +@pytest.mark.asyncio +async def test_worker_removes_more_expired_receipts_than_one_cleanup_batch_in_one_run( + parity_database: str, monkeypatch: pytest.MonkeyPatch +): + from litellm.accounting.worker import flush_accounting_deliveries + + monkeypatch.setenv("DATABASE_URL", parity_database) + monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_BATCH_SIZE", "100") + client = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database, autocommit=True) as connection: + connection.execute( + 'INSERT INTO "LiteLLM_SpendLogDelivery" (request_id, accounting_payload, is_batch, state, applied_at) ' + "SELECT 'expired-' || n, '{}'::jsonb, false, 'applied', CURRENT_TIMESTAMP - INTERVAL '2 days' " + "FROM generate_series(1, 501) AS n" + ) + await client.db.connect() + try: + await flush_accounting_deliveries(DBSpendUpdateWriter(), client, proxy_server.proxy_logging_obj) + with psycopg.connect(parity_database) as connection: + assert connection.execute('SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery"').fetchone() == (0,) finally: await client.db.disconnect() @@ -404,6 +428,7 @@ async def _assert_legacy_database_parity( assert len(rows) == 1 if path == "ledger" and retention == "lean": assert rows[0]["accounting_log_disabled"] is True + assert not {"messages", "response", "proxy_server_request", "metadata"}.intersection(rows[0]) frozen = rows[0]["accounting_payload"] if isinstance(frozen, str): frozen = json.loads(frozen) diff --git a/tests/integration/management/test_skill_download_history.py b/tests/integration/management/test_skill_download_history.py new file mode 100644 index 00000000000..00445dd7e27 --- /dev/null +++ b/tests/integration/management/test_skill_download_history.py @@ -0,0 +1,163 @@ +import hashlib +import io +import os +import uuid +import zipfile +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace +from functools import partial +from pathlib import Path +from typing import Final + +import psycopg +import pytest +from pydantic import TypeAdapter + +from litellm.models.plugin_upload import PluginDownloadsResponse +from litellm.types.proxy.skill_hub import SkillHubList +from tests.integration._support.actors import GatewayClient, owned_settings_client +from tests.integration._support.client import Gateway, Scenario, object_value, string_value +from tests.integration._support.database import write_rows + +PLUGINS: Final = "/claude-code/plugins" +SKILLS: Final = "/skill-hub/skills" +HISTORY: Final = TypeAdapter(PluginDownloadsResponse) + + +@pytest.fixture(scope="module") +def gateway(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("skill-downloads") + with owned_settings_client(GatewayClient(session_gateway, {}), directory, {"public_skills_index": True}) as client: + with client.gateway.scenario() as publisher: + yield replace(client.gateway, key=publisher.key(user_id=publisher.user(user_role="proxy_admin"))) + + +def _delete(gateway: Gateway, name: str) -> None: + assert gateway.request("DELETE", f"{PLUGINS}/{name}").status_code in {200, 404} + + +def _upload(scenario: Scenario, name: str) -> str: + scenario.cleanups.callback(_delete, scenario.gateway, name) + content: Final = io.BytesIO() + with zipfile.ZipFile(content, "w") as archive: + archive.writestr("SKILL.md", f"---\nname: {name}\ndescription: Synthetic download regression\n---\n") + response: Final = scenario.gateway.request_multipart( + f"{SKILLS}/upload", {"version": "1.2.3"}, {"file": ("skill.zip", content.getvalue(), "application/zip")} + ) + assert response.status_code == 200, response.text + return string_value(object_value(scenario.gateway.get(f"{PLUGINS}/{name}")["source"])["url"]) + + +def _download(gateway: Gateway, key: str | None, url: str) -> str: + response: Final = gateway.client.get( + url, headers={"Authorization": f"Bearer {key}"} if key else {}, params={"user_id": "forged-user"} + ) + assert response.status_code == 200, response.text + return hashlib.sha256(response.content).hexdigest() + + +def _history(gateway: Gateway, name: str, **params: str) -> PluginDownloadsResponse: + history: Final = gateway.get(f"{SKILLS}/{name}/downloads", params) + assert gateway.get(f"{PLUGINS}/{name}/downloads", params) == history + return HISTORY.validate_python(history) + + +def test_attribution_survives_actor_deletion_and_name_reuse_starts_empty(gateway: Gateway) -> None: + with gateway.scenario() as skills: + name: Final = f"download-{uuid.uuid4().hex}" + legacy: Final = _upload(skills, name) + portable: Final = f"{SKILLS}/{name}/archive" + original_id: Final = string_value(gateway.get(f"{SKILLS}/{name}")["id"]) + assert gateway.client.get(f"{SKILLS}/{name}/content").status_code == 200 + assert _history(gateway, name).total_downloads == 0 + assert SkillHubList.model_validate(gateway.get(SKILLS, {"query": name})).skills[0].download_count == 0 + with gateway.scenario() as actors: + team: Final = actors.team(team_alias=name) + email: Final = f"{name}@example.test" + user: Final = actors.user(user_role="internal_user", user_email=email) + gateway.post("/team/member_add", {"team_id": team, "member": {"role": "user", "user_id": user}}) + key: Final = actors.key(user_id=user, team_id=team, key_alias=name) + digests: Final = tuple(_download(gateway, key, url) for url in (legacy, portable)) + for url in (legacy, portable): + _download(gateway, None, url) + assert gateway.request("GET", url, key="sk-invalid").status_code == 401 + recorded: Final = _history(gateway, name) + assert (recorded.total_downloads, recorded.anonymous_downloads) == (4, 2) + assert SkillHubList.model_validate(gateway.get(SKILLS, {"query": name})).skills[0].download_count == 4 + assert tuple((row.version, row.authenticated, row.user_id, row.user_email, row.key_alias, + row.team_id, row.team_alias) for row in recorded.downloads) == ( + ("1.2.3", True, user, email, name, team, name), + ) * 2 + assert tuple(row.digest for row in recorded.downloads) == digests[::-1] + for prefix in (PLUGINS, SKILLS): + assert gateway.request("GET", f"{prefix}/{name}/downloads", key=key).status_code == 403 + assert gateway.client.get(f"{prefix}/{name}/downloads").status_code == 401 + assert _history(gateway, name) == recorded + _delete(gateway, name) + assert gateway.rows('SELECT id FROM "LiteLLM_ClaudeCodePluginDownload" WHERE plugin_id = %s', (original_id,)) == [] + _upload(skills, name) + assert gateway.get(f"{SKILLS}/{name}")["id"] != original_id + assert _history(gateway, name) == PluginDownloadsResponse(0, 0, (), None) + + +def test_concurrent_counts_and_keyset_pages_preserve_scope_and_deleted_cursors(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + name: Final = f"download-{uuid.uuid4().hex}" + legacy: Final = _upload(scenario, name) + portable: Final = f"{SKILLS}/{name}/archive" + key: Final = scenario.key() + with ThreadPoolExecutor(max_workers=4) as pool: + tuple(pool.map(partial(_download, gateway, key), (legacy, portable) * 4)) + initial: Final = _history(gateway, name) + ordered: Final = tuple(row.id for row in initial.downloads) + assert initial.total_downloads == len(frozenset(ordered)) == 8 + assert ordered == tuple(sorted(ordered, key=int, reverse=True)) + timestamp: Final = gateway.get(f"{PLUGINS}/{name}")["updated_at"] + with ThreadPoolExecutor(max_workers=4) as pool: + tuple(pool.map(partial(_download, gateway, None), (legacy, portable) * 8)) + assert gateway.get(f"{PLUGINS}/{name}")["updated_at"] == timestamp + assert _history(gateway, name) == replace(initial, total_downloads=24, anonymous_downloads=16) + first: Final = _history(gateway, name, limit="3") + assert tuple(row.id for row in first.downloads) == ordered[:3] + assert first.next_cursor == ordered[2] + _download(gateway, key, portable) + other: Final = f"{name}-other" + _download(gateway, key, _upload(scenario, other)) + write_rows( + 'DELETE FROM "LiteLLM_ClaudeCodePluginDownload" WHERE id IN (%s::bigint, %s::bigint)', + (ordered[2], ordered[4]), database_url=gateway.database_url, + ) + assert first.next_cursor is not None + second: Final = _history(gateway, name, limit="3", cursor=first.next_cursor) + assert tuple(row.id for row in second.downloads) == (ordered[3], ordered[5], ordered[6]) + assert (second.total_downloads, second.anonymous_downloads) == (23, 16) + assert second.next_cursor is not None + last: Final = _history(gateway, name, limit="3", cursor=second.next_cursor) + assert tuple(row.id for row in last.downloads) == (ordered[7],) + assert last.next_cursor is None + assert _history(gateway, other, cursor=first.next_cursor) == PluginDownloadsResponse(1, 0, (), None) + catalog: Final = SkillHubList.model_validate(gateway.get(SKILLS, {"query": name})) + assert {skill.name: skill.download_count for skill in catalog.skills} == {name: 23, other: 1} + + +def test_repeat_migration_is_catalog_only_while_download_tables_are_locked(gateway: Gateway) -> None: + migration: Final = (Path(__file__).resolve().parents[3] / + "litellm-proxy-extras/litellm_proxy_extras/migrations/20261105094500_add_skill_download_history/migration.sql" + ).read_bytes() + with gateway.scenario() as scenario: + name: Final = f"download-{uuid.uuid4().hex}" + url: Final = _upload(scenario, name) + for key in (None, gateway.key): + _download(gateway, key, url) + before: Final = _history(gateway, name) + database_url: Final = gateway.database_url or os.environ["DATABASE_URL"] + with psycopg.connect(database_url) as locker, psycopg.connect(database_url, autocommit=True) as migrator: + locker.execute( + 'LOCK TABLE "LiteLLM_ClaudeCodePluginTable", "LiteLLM_ClaudeCodePluginDownload" IN ACCESS EXCLUSIVE MODE' + ) + for _ in range(2): + with migrator.transaction(): + migrator.execute("SET LOCAL statement_timeout = '500ms'") + migrator.execute(migration) + assert _history(gateway, name) == before diff --git a/tests/integration/security/test_skill_scanner.py b/tests/integration/security/test_skill_scanner.py new file mode 100644 index 00000000000..e18378463ee --- /dev/null +++ b/tests/integration/security/test_skill_scanner.py @@ -0,0 +1,58 @@ +import asyncio +import sys +from datetime import datetime, timedelta, timezone +from pathlib import Path +from typing import Final + +import psutil +import pytest +from fastapi import HTTPException + +from cookbook.skill_scanner import server +from tests.integration._support.process import _wait_until_gone, group_members, stop_processes + + +@pytest.mark.parametrize("mode", ("deadline", "cap", "complete")) +async def test_adapter_bounds_worker_lifetime_and_reaps_its_process_group( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, mode: str, +) -> None: + script: Final = ( + "import subprocess, sys, time\n" + "child = subprocess.Popen([sys.executable, '-c', 'import time; time.sleep(30)'], " + "stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)\n" + "print(child.pid, flush=True)\n" + + ("pass\n" if mode == "complete" else "time.sleep(30)\n") + ) + child: Final = await asyncio.create_subprocess_exec( + sys.executable, "-c", script, stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.DEVNULL, start_new_session=True, + ) + try: + assert child.stdout is not None + descendant: Final = psutil.Process(int(await asyncio.wait_for(child.stdout.readline(), 5))) + if mode == "complete": + await asyncio.wait_for(child.wait(), 5) + output: Final = tmp_path / "report.json" + output.write_text('{"complete": true}') + + async def launch(*args: object, **kwargs: object) -> asyncio.subprocess.Process: + return child + + monkeypatch.setattr(server.asyncio, "create_subprocess_exec", launch) + monkeypatch.setattr(server, "SCAN_TIMEOUT", 0.05 if mode == "cap" else 30) + deadline: Final = ( + None if mode == "cap" + else datetime.now(timezone.utc) + timedelta(seconds=0.05 if mode == "deadline" else 30) + ) + result: Final = server.run_worker("unused.zip", output, False, deadline) + if mode == "complete": + assert await asyncio.wait_for(result, 5) == {"complete": True} + else: + with pytest.raises(HTTPException) as raised: + await asyncio.wait_for(result, 5) + assert raised.value.status_code == 504 + assert child.returncode is not None + _wait_until_gone(descendant, 3) + finally: + await asyncio.to_thread(stop_processes, group_members(child.pid), 0.1) + await child.wait() diff --git a/tests/integration/spend/golden/daily_activity_user_aggregated.json b/tests/integration/spend/golden/daily_activity_user_aggregated.json index 185d333965a..3189c000035 100644 --- a/tests/integration/spend/golden/daily_activity_user_aggregated.json +++ b/tests/integration/spend/golden/daily_activity_user_aggregated.json @@ -1,5 +1,5 @@ { - "metadata": {"entity_total_api_keys":null, + "metadata": {"entity_total_api_keys":{"user-1": 5}, "api_key_limit": 100, "has_more": false, "page": 1, @@ -346,7 +346,181 @@ } } }, - "entities": {}, + "entities": { + "user-1": { + "api_key_breakdown": { + "key-a": { + "metadata": { + "key_alias": "alias-a", + "key_exists": true, + "team_id": "team-1", + "user_email": "user@example.com", + "user_id": "user-1" + }, + "metrics": { + "api_requests": 1, + "autorouter_savings_spend": 0.0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 1, + "completion_tokens": 2, + "compression_saved_tokens": 0, + "compression_savings_spend": 0.0, + "failed_requests": 0, + "flat_cost": 0.0, + "ptu_hours": 0.0, + "gateway_injected_caching_savings_spend": 0.0, + "prompt_caching_savings_spend": 0.0, + "prompt_tokens": 2, + "spend": 100.0, + "successful_requests": 1, + "timed_requests": 0, + "total_response_time_ms": 0, + "total_tokens": 4 + } + }, + "key-b": { + "metadata": { + "key_alias": "alias-b", + "key_exists": true, + "team_id": "team-1", + "user_email": "user@example.com", + "user_id": "user-1" + }, + "metrics": { + "api_requests": 1, + "autorouter_savings_spend": 0.0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 1, + "completion_tokens": 2, + "compression_saved_tokens": 0, + "compression_savings_spend": 0.0, + "failed_requests": 0, + "flat_cost": 0.0, + "ptu_hours": 0.0, + "gateway_injected_caching_savings_spend": 0.0, + "prompt_caching_savings_spend": 0.0, + "prompt_tokens": 3, + "spend": 90.0, + "successful_requests": 1, + "timed_requests": 0, + "total_response_time_ms": 0, + "total_tokens": 5 + } + }, + "key-c": { + "metadata": { + "key_alias": "alias-c", + "key_exists": true, + "team_id": "team-1", + "user_email": "user@example.com", + "user_id": "user-1" + }, + "metrics": { + "api_requests": 1, + "autorouter_savings_spend": 0.0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 1, + "completion_tokens": 2, + "compression_saved_tokens": 0, + "compression_savings_spend": 0.0, + "failed_requests": 0, + "flat_cost": 0.0, + "ptu_hours": 0.0, + "gateway_injected_caching_savings_spend": 0.0, + "prompt_caching_savings_spend": 0.0, + "prompt_tokens": 4, + "spend": 80.0, + "successful_requests": 1, + "timed_requests": 0, + "total_response_time_ms": 0, + "total_tokens": 6 + } + }, + "key-cache": { + "metadata": { + "key_alias": "alias-cache", + "key_exists": true, + "team_id": "team-1", + "user_email": "user@example.com", + "user_id": "user-1" + }, + "metrics": { + "api_requests": 1, + "autorouter_savings_spend": 0.0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 1, + "completion_tokens": 2, + "compression_saved_tokens": 0, + "compression_savings_spend": 0.0, + "failed_requests": 0, + "flat_cost": 0.0, + "ptu_hours": 0.0, + "gateway_injected_caching_savings_spend": 0.0, + "prompt_caching_savings_spend": 0.0, + "prompt_tokens": 1000, + "spend": 2.0, + "successful_requests": 1, + "timed_requests": 0, + "total_response_time_ms": 0, + "total_tokens": 1002 + } + }, + "key-target": { + "metadata": { + "key_alias": "deleted-target", + "key_exists": false, + "team_id": "team-1", + "user_email": "user@example.com", + "user_id": "user-1" + }, + "metrics": { + "api_requests": 1, + "autorouter_savings_spend": 0.0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 2, + "completion_tokens": 2, + "compression_saved_tokens": 0, + "compression_savings_spend": 0.0, + "failed_requests": 0, + "flat_cost": 0.0, + "ptu_hours": 0.0, + "gateway_injected_caching_savings_spend": 0.0, + "prompt_caching_savings_spend": 0.0, + "prompt_tokens": 5, + "spend": 1.0, + "successful_requests": 1, + "timed_requests": 0, + "total_response_time_ms": 0, + "total_tokens": 7 + } + } + }, + "metadata": { + "user_alias": null, + "user_email": "user@example.com" + }, + "metrics": { + "api_requests": 5, + "autorouter_savings_spend": 0.0, + "cache_creation_input_tokens": 0, + "cache_read_input_tokens": 6, + "completion_tokens": 10, + "compression_saved_tokens": 0, + "compression_savings_spend": 0.0, + "failed_requests": 0, + "flat_cost": 0.0, + "ptu_hours": 0.0, + "gateway_injected_caching_savings_spend": 0.0, + "prompt_caching_savings_spend": 0.0, + "prompt_tokens": 1014, + "spend": 1273.0, + "successful_requests": 5, + "timed_requests": 0, + "total_response_time_ms": 0, + "total_tokens": 1024 + } + } + }, "mcp_servers": {}, "model_groups": { "model-cache": { diff --git a/tests/unit/accounting/test_projection.py b/tests/unit/accounting/test_projection.py index 17186a89922..c78d6bec757 100644 --- a/tests/unit/accounting/test_projection.py +++ b/tests/unit/accounting/test_projection.py @@ -150,6 +150,7 @@ async def test_entity_projection_matches_real_request_queues(monkeypatch: pytest client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj) monkeypatch.setattr(proxy_server, "prisma_client", client) monkeypatch.setattr(proxy_server, "disable_spend_logs", False) + monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False) monkeypatch.setattr(client, "spend_log_transactions", []) monkeypatch.setattr(client, "model_usage_transactions", []) writer: Final = DBSpendUpdateWriter() @@ -521,7 +522,9 @@ def test_stored_payload_keeps_guardrail_billing_but_not_guardrail_content() -> N assert dict(daily_user.increments)["compression_saved_tokens"] == 7 -def test_caller_supplied_savings_that_cannot_be_computed_never_drop_the_charge(monkeypatch: pytest.MonkeyPatch) -> None: +def test_caller_supplied_savings_beyond_json_range_are_clamped_and_never_drop_the_charge( + monkeypatch: pytest.MonkeyPatch, +) -> None: import litellm monkeypatch.setitem( @@ -552,6 +555,8 @@ def test_caller_supplied_savings_that_cannot_be_computed_never_drop_the_charge(m ) restored: Final = _FIELDS.validate_json(json.dumps(accounting, default=str)) assert dict(restored["metadata"])["accounting_savings"]["compression"] == 0.0 - assert [dict(item.increments)["spend"] for item in project(restored) if item.table == "LiteLLM_VerificationToken"] == [ - 0.5 - ] + assert dict(restored["metadata"])["compression_savings"] == {"tokens_saved": 2**53 - 1} + assert dict(restored["metadata"])["compression_savings"] == {"tokens_saved": 2**53 - 1} + assert [ + dict(item.increments)["spend"] for item in project(restored) if item.table == "LiteLLM_VerificationToken" + ] == [0.5] diff --git a/tests/unit/integration_support/test_wire.py b/tests/unit/integration_support/test_wire.py new file mode 100644 index 00000000000..c786895720f --- /dev/null +++ b/tests/unit/integration_support/test_wire.py @@ -0,0 +1,289 @@ +from __future__ import annotations + +import datetime +import io +import socket +import ssl +import time +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from http.client import HTTPConnection, HTTPResponse, HTTPSConnection +from pathlib import Path +from typing import Final +from urllib.parse import urlsplit + +import pytest +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ec +from cryptography.x509.oid import NameOID + +from tests.integration._support.wire import Reply, Request, wire_server + +PAUSE_SECONDS: Final = 6 +HANDLER_READ_SECONDS: Final = 0.3 +FIRST: Final = b'{"call": "first"}' +LATE: Final = b'{"call": "late"}' +UNREAD: Final = b"x" * (32 * 1024 * 1024) + + +def _self_signed_server_context(directory: Path) -> ssl.SSLContext: + key: Final = ec.generate_private_key(ec.SECP256R1()) + name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "wire peer")]) + now: Final = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc) + certificate: Final = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(1) + .not_valid_before(now) + .not_valid_after(now + datetime.timedelta(days=36500)) + .sign(key, hashes.SHA256()) + ) + (directory / "peer.pem").write_bytes( + certificate.public_bytes(serialization.Encoding.PEM) + + key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()) + ) + context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(directory / "peer.pem") + return context + + +def _trusting() -> ssl.SSLContext: + context: Final = ssl.create_default_context() + context.check_hostname = False + context.verify_mode = ssl.CERT_NONE + return context + + +def _client(url: str) -> HTTPConnection: + address: Final = urlsplit(url) + if address.scheme == "http": + return HTTPConnection(address.hostname or "", address.port) + return HTTPSConnection(address.hostname or "", address.port, context=_trusting()) + + +def _tls_socket(url: str) -> ssl.SSLSocket: + address: Final = urlsplit(url) + return _trusting().wrap_socket(socket.create_connection((address.hostname or "", address.port or 0))) + + +def _echo(request: Request) -> Reply: + return Reply(body=request.body) + + +def _answer_slowly(request: Request) -> Reply: + time.sleep(HANDLER_READ_SECONDS) + return Reply(body=request.body) + + +def _answer_in_paused_chunks(request: Request) -> Reply: + return Reply(chunks=(request.body[:1], request.body[1:]), pause_between_chunks=HANDLER_READ_SECONDS) + + +def _hold_first(request: Request) -> Reply: + if request.target == "/first": + time.sleep(HANDLER_READ_SECONDS) + return Reply(body=request.body) + + +def _paced_pieces(client: socket.socket) -> Iterator[bytes]: + while piece := client.recv(256 * 1024): + yield piece + time.sleep(0.02) + + +def _read_slowly(client: socket.socket) -> bytes: + return b"".join(_paced_pieces(client)) + + +def _read_after(client: socket.socket, delay: float) -> bytes: + time.sleep(delay) + with client.makefile("rb") as stream: + return stream.read() + + +def _post(target: str, body: bytes) -> bytes: + return b"POST %s HTTP/1.1\r\nhost: peer\r\ncontent-length: %d\r\n\r\n%s" % (target.encode(), len(body), body) + + +def _pause_then_send(connection: HTTPConnection, body: bytes) -> HTTPResponse: + time.sleep(PAUSE_SECONDS) + connection.request("POST", "/late", body) + return connection.getresponse() + + +def _send_headers_then_pause_before_the_body(connection: HTTPConnection, body: bytes) -> HTTPResponse: + connection.putrequest("POST", "/late") + connection.putheader("content-length", str(len(body))) + connection.endheaders() + time.sleep(PAUSE_SECONDS) + connection.send(body) + return connection.getresponse() + + +@pytest.mark.parametrize( + "send_late", + (_pause_then_send, _send_headers_then_pause_before_the_body), + ids=("pause-before-the-request", "pause-before-the-body"), +) +def test_a_kept_alive_connection_answers_a_request_whose_bytes_arrive_after_a_long_pause( + send_late: Callable[[HTTPConnection, bytes], HTTPResponse], +) -> None: + with wire_server(_echo, keep_alive=True) as wire: + address: Final = urlsplit(wire.url) + connection: Final = HTTPConnection(address.hostname or "", address.port) + connection.request("POST", "/first", FIRST) + first: Final = connection.getresponse().read() + late: Final = send_late(connection, LATE) + assert (first, late.status, late.read()) == (FIRST, 200, LATE) + assert tuple(request.body for request in wire.drain()) == (FIRST, LATE) + assert wire.connections() == 1 + assert connection.sock.recv(1) == b"" + connection.close() + + +def test_a_request_cut_off_by_teardown_never_reaches_respond() -> None: + with wire_server(_echo, keep_alive=True) as wire: + address: Final = urlsplit(wire.url) + connection: Final = HTTPConnection(address.hostname or "", address.port) + connection.request("POST", "/first", FIRST) + first: Final = connection.getresponse().read() + connection.sock.sendall(b"POST /cut HTTP/1.1\r\ncontent-type: application/json\r\ncontent-len") + time.sleep(HANDLER_READ_SECONDS) + assert (first, tuple(request.target for request in wire.drain())) == (FIRST, ("/first",)) + connection.close() + + +@pytest.mark.timeout(60) +def test_a_reply_its_client_stopped_reading_fails_teardown_instead_of_hanging() -> None: + peer: Final = ExitStack() + wire: Final = peer.enter_context(wire_server(lambda _: Reply(body=UNREAD), keep_alive=True, teardown_grace=0.5)) + address: Final = urlsplit(wire.url) + with socket.create_connection((address.hostname or "", address.port or 0)) as client: + client.sendall(b"GET /large HTTP/1.1\r\nhost: peer\r\n\r\n") + while wire.received.empty(): + time.sleep(0.01) + with pytest.raises(AssertionError, match="no progress"): + peer.close() + + +@pytest.mark.timeout(60) +def test_a_long_reply_its_client_is_still_reading_at_teardown_is_not_cut() -> None: + peer: Final = ExitStack() + wire: Final = peer.enter_context(wire_server(lambda _: Reply(body=UNREAD), keep_alive=True, teardown_grace=0.5)) + address: Final = urlsplit(wire.url) + with socket.create_connection((address.hostname or "", address.port or 0)) as client, ThreadPoolExecutor(1) as pool: + client.sendall(b"GET /large HTTP/1.1\r\nhost: peer\r\n\r\n") + while wire.received.empty(): + time.sleep(0.01) + reply: Final = pool.submit(_read_slowly, client) + time.sleep(1) + peer.close() + received: Final = reply.result() + assert (received[: len(b"HTTP/1.1 200")], received[-len(UNREAD) :] == UNREAD) == (b"HTTP/1.1 200", True) + + +@pytest.mark.timeout(60) +def test_a_reply_its_client_starts_reading_after_teardown_starts_still_arrives_whole() -> None: + peer: Final = ExitStack() + wire: Final = peer.enter_context(wire_server(lambda _: Reply(body=UNREAD), keep_alive=True, teardown_grace=2)) + address: Final = urlsplit(wire.url) + with socket.create_connection((address.hostname or "", address.port or 0)) as client, ThreadPoolExecutor(1) as pool: + client.sendall(b"GET /large HTTP/1.1\r\nhost: peer\r\n\r\n") + while wire.received.empty(): + time.sleep(0.01) + time.sleep(2.5) + reply: Final = pool.submit(_read_after, client, 0.3) + peer.close() + received: Final = reply.result() + assert (received[: len(b"HTTP/1.1 200")], received[-len(UNREAD) :] == UNREAD) == (b"HTTP/1.1 200", True) + + +def test_a_complete_request_sent_as_teardown_starts_is_still_answered() -> None: + with wire_server(_echo, keep_alive=True) as wire: + connection: Final = _client(wire.url) + connection.request("POST", "/first", FIRST) + first: Final = connection.getresponse().read() + connection.request("POST", "/complete", LATE) + late: Final = connection.getresponse() + assert (first, late.status, late.read(), tuple(request.target for request in wire.drain())) == ( + FIRST, + 200, + LATE, + ("/first", "/complete"), + ) + connection.close() + + +@pytest.mark.parametrize("secure", (False, True), ids=("plain", "tls")) +def test_a_reply_still_being_written_at_teardown_reaches_its_client_whole(secure: bool, tmp_path: Path) -> None: + slow: Final = Reply(chunks=(b"first", b"second"), pause_between_chunks=HANDLER_READ_SECONDS) + tls: Final = _self_signed_server_context(tmp_path) if secure else None + with wire_server(lambda _: slow, tls=tls, keep_alive=True) as wire: + connection: Final = _client(wire.url) + connection.request("GET", "/slow") + response: Final = connection.getresponse() + head: Final = response.read(len(b"first")) + assert head + response.read() == b"firstsecond" + connection.close() + + +def test_a_complete_request_queued_behind_a_busy_handler_is_still_answered() -> None: + with wire_server(_hold_first, keep_alive=True) as wire: + address: Final = urlsplit(wire.url) + client: Final = socket.create_connection((address.hostname or "", address.port or 0)) + client.sendall(_post("/first", FIRST)) + while wire.received.empty(): + time.sleep(0.01) + client.sendall(_post("/second", LATE)) + with client, client.makefile("rb") as stream: + replies: Final = stream.read() + assert ( + tuple(request.target for request in wire.drain()), + replies.count(b"HTTP/1.1 200"), + replies[-len(LATE) :], + ) == ( + ("/first", "/second"), + 2, + LATE, + ) + + +@pytest.mark.parametrize("respond", (_answer_slowly, _answer_in_paused_chunks), ids=("slow-respond", "paused-chunks")) +def test_a_reply_that_takes_longer_than_the_teardown_grace_still_reaches_its_client( + respond: Callable[[Request], Reply], +) -> None: + with wire_server(respond, keep_alive=True, teardown_grace=HANDLER_READ_SECONDS / 2) as wire: + connection: Final = _client(wire.url) + connection.request("POST", "/slow", FIRST) + while wire.received.empty(): + time.sleep(0.01) + response: Final = connection.getresponse() + assert (response.status, response.read()) == (200, FIRST) + connection.close() + + +@pytest.mark.timeout(60) +def test_a_tls_request_larger_than_the_read_buffer_is_answered(tmp_path: Path) -> None: + large: Final = b"x" * (3 * io.DEFAULT_BUFFER_SIZE // 2) + with wire_server(_echo, tls=_self_signed_server_context(tmp_path), keep_alive=True) as wire: + client: Final = _tls_socket(wire.url) + client.sendall(_post("/large", large)) + response: Final = HTTPResponse(client) + response.begin() + assert (response.status, response.read()) == (200, large) + client.close() + + +@pytest.mark.timeout(60) +def test_a_tls_record_cut_off_by_teardown_closes_its_connection_instead_of_hanging(tmp_path: Path) -> None: + tls: Final = _self_signed_server_context(tmp_path) + with wire_server(_echo, tls=tls, keep_alive=True, teardown_grace=HANDLER_READ_SECONDS) as wire: + client: Final = _tls_socket(wire.url) + socket.socket.sendall(client, b"\x17\x03\x03\x40\x00" + b"x" * 16) + time.sleep(HANDLER_READ_SECONDS) + assert wire.drain() == () + client.close() diff --git a/tests/unit/integrations/email_alerting/__init__.py b/tests/unit/integrations/email_alerting/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/integrations/email_alerting/test_model_deprecation_email.py b/tests/unit/integrations/email_alerting/test_model_deprecation_email.py new file mode 100644 index 00000000000..54dbe342148 --- /dev/null +++ b/tests/unit/integrations/email_alerting/test_model_deprecation_email.py @@ -0,0 +1,298 @@ +"""Retirement reminder schedule, durable dedupe, ownership scope and safe rendering.""" + +import asyncio +from dataclasses import replace +from datetime import date, datetime, timedelta, timezone +from types import MappingProxyType +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.integrations import model_deprecation_email as module +from litellm.integrations.email_templates.model_deprecation_email import ( + render_model_deprecation_email, +) +from litellm.models.team import LiteLLM_TeamTable +from litellm.models.user import LiteLLM_UserTable +from litellm.proxy._types import LiteLLM_VerificationToken +from litellm.repositories.config_repository import ConfigParam +from litellm.types.integrations.model_deprecation_email import ( + DeprecationEmailMilestone, + ModelDeprecationEmail, +) +from litellm.types.proxy.model_deprecation import ( + ModelDeprecationInfo, + ModelDeprecationResponse, + SuccessorModelSource, +) + +NOW = datetime(2026, 11, 1, tzinfo=timezone.utc) + + +def info(days=30, model="gemini/gemini-2.5-pro"): + return ModelDeprecationInfo( + model_name="research", + litellm_model=model, + deprecation_date=NOW.date() + timedelta(days=days), + days_until_deprecation=days, + status="imminent", + successor_model="gemini/gemini-3.8-flash", + successor_model_source=SuccessorModelSource.MODEL_COST_MAP, + deprecation_source="https://ai.google.dev/gemini-api/docs/deprecations", + ) + + +class MemoryConfig: + def __init__(self): + self.rows = {"general_settings": {"email_settings": {"Model Deprecations": True}}} + + async def get_param(self, key): + return ConfigParam(key, self.rows[key]) if key in self.rows else None + + async def set_param_if_revision(self, key, value, revision): + current = self.rows.get(key, {}).get("revision", 0) + if current != revision: + return False + self.rows[key] = value + return True + + async def delete_params_older_than(self, prefix, before): + return 0 + + +@pytest.fixture() +def sender(monkeypatch): + monkeypatch.setenv("PROXY_BASE_URL", "https://gateway.example") + instance = module.ModelDeprecationEmailSender(MagicMock()) + instance.config = MemoryConfig() + instance._impacts = AsyncMock( + return_value=MappingProxyType( + {"owner@example.com": module._Impact(key_count=3, requests=12, teams=frozenset({"Research"}))} + ) + ) + send = AsyncMock(return_value=True) + monkeypatch.setattr("litellm.proxy.utils.send_email", send) + return instance, send + + +def schedule(monkeypatch, models): + snapshot = ModelDeprecationResponse( + deprecated=[], imminent=models, upcoming=[], warn_within_days=30, checked_at=NOW + ) + monkeypatch.setattr(module, "collect_model_deprecations", lambda **kwargs: snapshot) + + +@pytest.mark.asyncio +async def test_should_send_only_the_two_milestones_and_keep_receipts_across_restart(sender, monkeypatch): + instance, send = sender + schedule(monkeypatch, [info()]) + assert await instance.run(MagicMock(), now=NOW) == 1 + assert await instance.run(MagicMock(), now=NOW) == 0 + restarted = module.ModelDeprecationEmailSender(MagicMock()) + restarted.config = instance.config + restarted._impacts = instance._impacts + assert await restarted.run(MagicMock(), now=NOW) == 0 + final_day = NOW + timedelta(days=29) + schedule(monkeypatch, [info().model_copy(update={"days_until_deprecation": 1})]) + assert await restarted.run(MagicMock(), now=final_day) == 1 + assert await restarted.run(MagicMock(), now=final_day) == 0 + assert send.await_count == 2 + assert "in 30 days" in send.await_args_list[0].kwargs["subject"] + assert "tomorrow" in send.await_args_list[1].kwargs["subject"] + assert send.await_args_list[0].kwargs["message_id"] != send.await_args_list[1].kwargs["message_id"] + assert "alert("x")', + "deprecation_source": "javascript:alert(1)", + } + ) + email = ModelDeprecationEmail( + model=model, + usage_model=model.litellm_model, + recipient_email="owner@example.com", + milestone=DeprecationEmailMilestone.THIRTY_DAYS, + key_count=3, + requests=12, + team_names=("Research",), + start_date=date(2026, 10, 3), + end_date=NOW.date(), + ) + rendered = render_model_deprecation_email(email, base_url="https://gateway.example") + assert "/public/assets/logos/litellm_logo.png" in rendered.html + assert "litellm-listing.s3.amazonaws.com" not in rendered.html + assert "