scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 353

Merge remote-tracking branch 'origin/main' into litellm_/i18n-left-shell-cca63a

Browse source 
# Conflicts:
#	ui/litellm-dashboard/src/components/leftnav.tsx
This commit is contained in:
Yuneng Jiang 2026-11-05 09:52:57 -08:00
commit cbbfe9e17d
No known key found for this signature in database
160 changed files with 17403 additions and 2830 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
.github/PULL_REQUEST_TEMPLATE/general.md vendored
View file

@ -6,19 +6,12 @@
## Handwritten 1-3 sentence description
<!-- The author writes a plain summary in their own words before requesting review.
Wrap newly AI-written summary text in the make-me-human-readable tags below, including backticks.
The author reviews or rewrites the draft and removes the tags.
When updating an existing PR, preserve the author's wording and removed tags.
Changes elsewhere in the body do not make an unchanged summary an AI draft again.
Add tags again only around newly AI-written replacement summary text, never merely to match this template.
This preserves author review across later updates (2026-11-02, #49779).
<!-- The author writes this in their own words, and it can be left empty. AI agents never write here and leave the
placeholder line below as is, and when updating an existing PR they keep whatever the author wrote. Keep this heading
Say what the PR does for users and why, then anything that keeps it working over time (a sync job, a config someone owns)
No bullets, tables, or code names here. Everything else goes inside the Details block below -->
`<make-me-human-readable>`
<blah>
`</make-me-human-readable>`
[insert handwritten description here]
<details>
<summary>Details</summary>

60
cookbook/skill_scanner/README.md Normal file
View file

@ -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

172
cookbook/skill_scanner/server.py Normal file
View file

@ -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")))

28
enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py
View file

@ -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()
return cls().to_dict()

16
litellm-proxy-extras/litellm_proxy_extras/migrations/20261105000000_skill_version_history/migration.sql Normal file
View file

@ -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");

27
litellm-proxy-extras/litellm_proxy_extras/migrations/20261105010000_skill_submissions/migration.sql Normal file
View file

@ -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);

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20261105020000_skill_scanning/migration.sql Normal file
View file

@ -0,0 +1,3 @@
SET LOCAL lock_timeout = '10s';
ALTER TABLE "LiteLLM_SkillSubmission" ADD COLUMN IF NOT EXISTS "scan_json" TEXT;

37
litellm-proxy-extras/litellm_proxy_extras/migrations/20261105094500_add_skill_download_history/migration.sql Normal file
View file

@ -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 $$;

59
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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

2
litellm/accounting/application.py
View file

@ -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)
)

2
litellm/accounting/snapshot.py
View file

@ -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,
)
)

35
litellm/accounting/worker.py
View file

@ -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(

2
litellm/constants.py
View file

@ -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"

181
litellm/integrations/email_templates/model_deprecation_email.py Normal file
View file

@ -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'<td width="38" valign="middle"><img src="{_LOGO_BASE_URL}/{brand[1]}" width="24" height="24" alt="{escape(brand[0], quote=True)}" style="display:block;object-fit:contain;border:0;"></td>'
if brand
else ""
)
return (
f'<table role="presentation" width="100%" cellspacing="0" cellpadding="0"><tr>{logo}'
f'<td><div style="font-size:11px;line-height:18px;color:#526079;">{escape(_provider_name(provider))}</div>'
f"<div style=\"font:14px 'Courier New',monospace;line-height:22px;word-break:break-word;overflow-wrap:anywhere;color:#17213b;\">{escape(model)}</div></td></tr></table>"
)
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 '<p style="margin:0;font-size:14px;line-height:22px;color:#17213b;">No recommended model listed. Choose a supported replacement.</p>'
)
source: Final = _recommendation_source(email.model.successor_model_source)
notice: Final = _http_url(email.model.deprecation_source)
source_link: Final = (
f'<a href="{escape(notice, quote=True)}" style="color:#263cdb;text-decoration:underline;">Read the provider notice &#8599;</a>'
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"""<!doctype html>
<html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
<title>{escape(subject)}</title>
<style>@import url('https://fonts.googleapis.com/css2?family=DM+Sans:wght@400;500;600;700&display=swap');
@media only screen and (max-width:620px){{.pad{{padding:24px!important}}.headline{{font-size:30px!important;line-height:36px!important}}.brand{{padding:20px 24px!important}}.brand img{{width:96px!important}}.comparison-half{{display:block!important;width:auto!important;border-right:0!important}}.comparison-after{{border-top:1px solid #e2e6ee!important}}}}
</style></head>
<body style="margin:0;padding:0;background:#f7f8fa;color:#17213b;font-family:'DM Sans',Arial,sans-serif;">
<div style="display:none;max-height:0;overflow:hidden;mso-hide:all;">{escape(preheader)}</div>
<table role="presentation" width="100%" cellspacing="0" cellpadding="0" border="0"><tr><td style="padding:0;">
<table role="presentation" width="100%" cellspacing="0" cellpadding="0" border="0" style="width:100%;background:#ffffff;">
<tr><td class="brand" style="padding:24px 36px;border-bottom:1px solid #e2e6ee;">
<table role="presentation" width="100%" cellspacing="0" cellpadding="0"><tr><td><img src="{_EMAIL_LOGO_URL}" width="112" alt="LiteLLM" style="display:block;border:0;height:auto;"></td>
<td align="right" style="font:11px 'Courier New',monospace;letter-spacing:1.5px;color:#526079;">/ MODEL DEPRECATIONS</td></tr></table></td></tr>
<tr><td class="pad" style="padding:32px 36px;background:#263cdb;color:#ffffff;">
<table role="presentation" width="100%" cellspacing="0" cellpadding="0"><tr><td style="font:11px 'Courier New',monospace;letter-spacing:1.8px;color:#ffffff;">{label}</td>
<td align="right" style="font:11px 'Courier New',monospace;color:#ffffff;">{days:02d} {day_label} UNTIL DEPRECATION</td></tr></table>
<h1 class="headline" style="margin:22px 0 16px;font-size:36px;line-height:42px;font-weight:600;letter-spacing:-1px;color:#ffffff;">Model being deprecated</h1>
<p style="margin:0;font-size:14px;line-height:22px;color:#ffffff;">{escape(announcement)}</p>
</td></tr>
<tr><td class="pad" style="padding:28px 36px;">
<table role="presentation" width="100%" cellspacing="0" cellpadding="0" style="border:1px solid #e2e6ee;border-collapse:collapse;">
<tr><td class="comparison-half" width="50%" valign="top" style="padding:16px 18px;border-right:1px solid #e2e6ee;">
<div style="margin-bottom:8px;font:10px 'Courier New',monospace;letter-spacing:1px;color:#526079;">BEFORE &middot; MODEL TO REPLACE</div>
{_model_html(email.usage_model, provider)}
</td>
<td class="comparison-half comparison-after" width="50%" valign="top" style="padding:16px 18px;background:#f5f7ff;">
<div style="margin-bottom:8px;font:10px 'Courier New',monospace;letter-spacing:1px;color:#263cdb;">AFTER &middot; RECOMMENDED MODEL</div>
{after}
</td></tr></table>
<p style="margin:10px 0 24px;font-size:12px;line-height:20px;color:#526079;">Recommendation source: {escape(source)}<br>{source_link}</p>
<h2 style="margin:0 0 12px;font-size:17px;font-weight:600;color:#17213b;">Your affected usage</h2>
<table role="presentation" width="100%" cellspacing="0" cellpadding="0" style="border-top:1px solid #e2e6ee;border-bottom:1px solid #e2e6ee;">
<tr><td width="50%" style="padding:14px 0;"><div style="font:24px 'Courier New',monospace;color:#17213b;">{email.key_count:,}</div><div style="margin-top:5px;font-size:12px;color:#526079;">Affected keys</div></td>
<td width="50%" style="padding:14px 0 14px 20px;border-left:1px solid #e2e6ee;"><div style="font:24px 'Courier New',monospace;color:#17213b;">{email.requests:,}</div><div style="margin-top:5px;font-size:12px;color:#526079;">Requests &middot; last 30 days</div></td></tr></table>
<p style="margin:10px 0 24px;font-size:12px;line-height:20px;color:#526079;">{escape(team_names + more_teams)}<br>Your keys and teams you administer.</p>
<h2 style="margin:0 0 10px;font-size:17px;font-weight:600;color:#17213b;">What to do</h2>
<ol style="margin:0 0 24px;padding-left:20px;font-size:13px;line-height:22px;color:#526079;">
<li style="padding:0 0 8px 4px;">Review the affected keys in your usage dashboard.</li>
<li style="padding:0 0 8px 4px;">{escape(test_step)}</li>
<li style="padding-left:4px;">{escape(switch_step)}</li></ol>
<table role="presentation" cellspacing="0" cellpadding="0"><tr><td bgcolor="#263cdb" style="border-radius:4px;mso-padding-alt:14px 20px;">
<a href="{escape(review_url, quote=True)}" style="display:inline-block;padding:14px 20px;color:#ffffff;font-size:13px;font-weight:600;line-height:18px;text-decoration:none;">Review affected usage &#8599;</a></td></tr></table>
</td></tr>
<tr><td class="pad" style="padding:20px 36px;background:#fafbfc;border-top:1px solid #e2e6ee;font-size:11px;line-height:18px;color:#526079;">
Your gateway admin enabled these emails. Reminder schedule: 30 days and 1 day before deprecation.<br><br>LiteLLM &middot; AI Gateway</td></tr>
</table></td></tr></table></body></html>"""
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)

267
litellm/integrations/model_deprecation_email.py Normal file
View file

@ -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"<litellm-deprecation-{key.removeprefix(_RECEIPT_PREFIX)}@alerts.litellm.ai>",
):
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)

3
litellm/integrations/otel/model/spans.py
View file

@ -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",

11
litellm/litellm_core_utils/safe_json_dumps.py
View file

@ -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"

225
litellm/llms/github_copilot/authenticator.py
View file

@ -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):

6
litellm/llms/github_copilot/chat/transformation.py
View file

@ -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(

12
litellm/llms/github_copilot/common_utils.py
View file

@ -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

4
litellm/llms/github_copilot/embedding/transformation.py
View file

@ -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",

11
litellm/llms/github_copilot/github_host.py
View file

@ -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}",

4
litellm/llms/github_copilot/messages/transformation.py
View file

@ -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",

146
litellm/llms/github_copilot/per_user_auth.py
View file

@ -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,

4
litellm/llms/github_copilot/responses/transformation.py
View file

@ -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",

180
litellm/model_prices_and_context_window_backup.json
View file

@ -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
}
}

140
litellm/models/plugin_upload.py
View file

@ -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]: ...

110
litellm/models/skill_submission.py Normal file
View file

@ -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: ...

15
litellm/proxy/_lazy_features.py
View file

@ -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",

2859
litellm/proxy/_lazy_openapi_snapshot.json
View file

File diff suppressed because it is too large Load diff

33
litellm/proxy/_types.py
View file

@ -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']).",

210
litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py
View file

@ -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,

5
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -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)

1
litellm/proxy/config_resolvers/alerting.py
View file

@ -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"),

266
litellm/proxy/db/db_spend_update_writer.py
View file

@ -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::<value>::user_id::<value>"; 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")

11
litellm/proxy/deps/repositories.py
View file

@ -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)
]

208
litellm/proxy/discovery_endpoints/agent_skills_archive.py
View file

@ -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):

71
litellm/proxy/discovery_endpoints/agent_skills_endpoints.py
View file

@ -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:

682
litellm/proxy/discovery_endpoints/skill_hub_endpoints.py Normal file
View file

@ -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,
)

172
litellm/proxy/discovery_endpoints/skill_hub_mutations.py Normal file
View file

@ -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)

192
litellm/proxy/discovery_endpoints/skill_hub_submissions.py Normal file
View file

@ -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)

126
litellm/proxy/discovery_endpoints/skill_import.py Normal file
View file

@ -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/<owner>/<repository>/<skill> 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

186
litellm/proxy/discovery_endpoints/skill_scanning.py Normal file
View file

@ -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)

91
litellm/proxy/management_endpoints/access_group_endpoints.py
View file

@ -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)

26
litellm/proxy/management_endpoints/common_daily_activity.py
View file

@ -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))

2
litellm/proxy/management_endpoints/daily_activity_scopes.py
View file

@ -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(
{

24
litellm/proxy/management_endpoints/internal_user_endpoints.py
View file

@ -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(

29
litellm/proxy/management_endpoints/organization_endpoints.py
View file

@ -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,

10
litellm/proxy/management_helpers/access_group_model_sync.py
View file

@ -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

62
litellm/proxy/management_helpers/model_allowlist_rename_sync.py
View file

@ -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(

16
litellm/proxy/middleware/request_size_limit_middleware.py
View file

@ -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

6
litellm/proxy/proxy_server.py
View file

@ -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,
}
)

59
litellm/proxy/schema.prisma
View file

@ -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

28
litellm/proxy/spend_tracking/spend_management_endpoints.py
View file

@ -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",

68
litellm/proxy/utils.py
View file

@ -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):

52
litellm/repositories/accounting_delivery_repository.py
View file

@ -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

9
litellm/repositories/config_repository.py
View file

@ -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}}
)

175
litellm/repositories/plugin_publication.py Normal file
View file

@ -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)

265
litellm/repositories/plugin_upload_repository.py
View file

@ -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,
)

368
litellm/repositories/skill_submission_repository.py Normal file
View file

@ -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,
)

4
litellm/repositories/table_repositories.py
View file

@ -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"

42
litellm/repositories/unit_of_work.py
View file

@ -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
),

35
litellm/types/integrations/model_deprecation_email.py Normal file
View file

@ -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

13
litellm/types/proxy/claude_code_endpoints.py
View file

@ -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):

165
litellm/types/proxy/skill_hub.py Normal file
View file

@ -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)

91
litellm/types/proxy/skill_scanning.py Normal file
View file

@ -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

180
model_prices_and_context_window.json
View file

@ -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
}
}

2
pyproject.toml
View file

@ -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",

59
schema.prisma
View file

@ -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

1
tests/code_coverage_tests/recursive_detector.py
View file

@ -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.
]

126
tests/integration/_support/wire.py
View file

@ -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}"

201
tests/integration/database/test_skill_publication.py Normal file
View file

@ -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

454
tests/integration/database/test_skill_submissions.py Normal file
View file

@ -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

197
tests/integration/database/test_spend_lock_order.py Normal file
View file

@ -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"])

63
tests/integration/database/test_spendlogs_delivery.py
View file

@ -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() == (

43
tests/integration/database/test_spendlogs_ledger_parity.py
View file

@ -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)

163
tests/integration/management/test_skill_download_history.py Normal file
View file

@ -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

58
tests/integration/security/test_skill_scanner.py Normal file
View file

@ -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()

178
tests/integration/spend/golden/daily_activity_user_aggregated.json
View file

@ -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": {

13
tests/unit/accounting/test_projection.py
View file

@ -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]

289
tests/unit/integration_support/test_wire.py Normal file
View file

@ -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()

0
tests/unit/integrations/email_alerting/__init__.py Normal file
View file

298
tests/unit/integrations/email_alerting/test_model_deprecation_email.py Normal file
View file

@ -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 "<html" in send.await_args_list[0].kwargs["html"]
assert "3 keys | 12 requests in the last 30 days" in send.await_args_list[0].kwargs["text"]
@pytest.mark.asyncio
@pytest.mark.parametrize("days", [31, 29, 2, 0, -1])
async def test_should_not_send_on_other_days(sender, monkeypatch, days):
instance, send = sender
schedule(monkeypatch, [info(days)])
assert await instance.run(MagicMock(), now=NOW) == 0
send.assert_not_awaited()
instance._impacts.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"settings",
[{}, {"email_settings": {}}, {"email_settings": {"Model Deprecations": False}}],
)
async def test_should_require_explicit_opt_in(sender, settings):
instance, send = sender
instance.config.rows["general_settings"] = settings
assert await instance.run(MagicMock(), now=NOW) == 0
send.assert_not_awaited()
instance._impacts.assert_not_awaited()
@pytest.mark.asyncio
async def test_should_claim_delivery_once_across_concurrent_workers(sender, monkeypatch):
instance, send = sender
schedule(monkeypatch, [info()])
second = module.ModelDeprecationEmailSender(MagicMock())
second.config = instance.config
second._impacts = instance._impacts
assert sum(await asyncio.gather(instance.run(MagicMock(), now=NOW), second.run(MagicMock(), now=NOW))) == 1
send.assert_awaited_once()
@pytest.mark.asyncio
async def test_should_retry_failed_delivery_after_the_lease_expires(sender, monkeypatch):
instance, send = sender
schedule(monkeypatch, [info()])
send.side_effect = [False, True]
assert await instance.run(MagicMock(), now=NOW) == 0
assert await instance.run(MagicMock(), now=NOW + timedelta(minutes=1)) == 0
assert await instance.run(MagicMock(), now=NOW + timedelta(minutes=6)) == 1
assert await instance.run(MagicMock(), now=NOW + timedelta(minutes=7)) == 0
assert send.await_count == 2
assert send.await_args_list[0].kwargs["message_id"] == send.await_args_list[1].kwargs["message_id"]
@pytest.mark.asyncio
async def test_should_dedupe_model_aliases_and_skip_models_with_no_observed_owners(sender, monkeypatch):
instance, send = sender
schedule(monkeypatch, [info(), info().model_copy(update={"model_name": "second-alias"})])
assert await instance.run(MagicMock(), now=NOW) == 1
send.assert_awaited_once()
instance.config.rows = {"general_settings": {"email_settings": {"Model Deprecations": True}}}
instance._impacts.return_value = {}
assert await instance.run(MagicMock(), now=NOW) == 0
assert send.await_count == 1
@pytest.mark.asyncio
async def test_should_scope_usage_to_current_owners_and_admins_excluding_former_admins_across_pages(
monkeypatch,
):
instance = module.ModelDeprecationEmailSender(MagicMock())
monkeypatch.setattr(module, "_PAGE_SIZE", 2)
grouped = AsyncMock(
side_effect=[
[
{"api_key": "one", "_sum": {"api_requests": 10}},
{"api_key": "two", "_sum": {"api_requests": 20}},
],
[{"api_key": "three", "_sum": {"api_requests": 30}}],
]
)
monkeypatch.setattr(
module.DailyUserSpendRepository,
"table",
property(lambda _: MagicMock(group_by=grouped)),
)
keys = [
LiteLLM_VerificationToken(token="one", user_id="alice", team_id="research"),
LiteLLM_VerificationToken(token="two", user_id="alice"),
LiteLLM_VerificationToken(token="three", user_id="bob", team_id="research"),
]
async def key_rows(_, field, values):
ids = set(values)
return [key for key in keys if key.token in ids]
monkeypatch.setattr(module.VerificationTokenRepository, "find_many_in", key_rows)
team = LiteLLM_TeamTable(
team_id="research",
team_alias="Research",
admins=["bob", "mallory", "removed-admin"],
members_with_roles=[
{"user_id": "bob", "user_email": "BOB@example.com", "role": "admin"},
{"user_id": "mallory", "user_email": "mallory@example.com", "role": "user"},
],
)
async def team_rows(_, field, values):
return [team] if "research" in values else []
monkeypatch.setattr(module.TeamRepository, "find_many_in", team_rows)
async def user_rows(_, field, values):
return [LiteLLM_UserTable(user_id=user_id, user_email=f"{user_id}@example.com") for user_id in values]
monkeypatch.setattr(module.UserRepository, "find_many_in", user_rows)
impacts = await instance._impacts("gemini/gemini-2.5-pro", date(2026, 10, 3), NOW.date())
assert set(impacts) == {"alice@example.com", "bob@example.com"}
assert impacts["alice@example.com"].requests == 30
assert impacts["alice@example.com"].key_count == 2
assert impacts["bob@example.com"].requests == 40
assert impacts["bob@example.com"].key_count == 2
assert impacts["bob@example.com"].teams == frozenset({"Research"})
assert grouped.await_args_list[0].kwargs["where"] == {
"model": "gemini/gemini-2.5-pro",
"date": {"gte": "2026-10-03", "lte": "2026-11-01"},
"api_requests": {"gt": 0},
}
def test_should_render_the_current_ui_logo_safe_model_text_and_actionable_links():
model = info().model_copy(
update={
"successor_model": '<script>alert("x")</script>',
"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 "<script>" not in rendered.html
assert "javascript:" not in rendered.html
assert "&lt;script&gt;" in rendered.html
assert "/ui/usage?tab=models&amp;from=2026-10-03&amp;to=2026-11-01" in rendered.html
assert "Recommendation source: LiteLLM model cost map" in rendered.text
assert "Google is deprecating this model on 01 Dec 2026 (UTC)." in rendered.html
assert "Google is deprecating this model on 01 Dec 2026 (UTC)." in rendered.text
assert "google_pse.png" in rendered.html
assert 'alt="Google"' in rendered.html
assert "BEFORE &middot; MODEL TO REPLACE" in rendered.html
assert "AFTER &middot; RECOMMENDED MODEL" in rendered.html
assert "What to do" in rendered.html
assert "check that requests to the old model reach zero" in rendered.text
missing = render_model_deprecation_email(
replace(email, model=model.model_copy(update={"successor_model": None})),
base_url="https://gateway.example",
)
assert "No recommended model listed. Choose a supported replacement." in missing.html
assert "Choose a supported replacement and test it with your workload." in missing.text
tomorrow = render_model_deprecation_email(
replace(email, milestone=DeprecationEmailMilestone.ONE_DAY),
base_url="https://gateway.example",
)
assert "01 DAY UNTIL DEPRECATION" in tomorrow.html
assert "Model being deprecated" in rendered.html
assert "Model being deprecated" in rendered.subject
assert "Model being deprecated" in tomorrow.html
assert "Model being deprecated" in tomorrow.subject
assert "FINAL REMINDER" in tomorrow.html
assert "tomorrow" in tomorrow.subject
@pytest.mark.asyncio
async def test_should_keep_the_email_loop_running_after_a_delivery_error(monkeypatch):
from litellm.proxy import proxy_server
router = MagicMock()
database = MagicMock()
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "prisma_client", database)
runner = MagicMock()
runner.run = AsyncMock(side_effect=[RuntimeError("temporary SMTP failure"), 0])
monkeypatch.setattr(module, "ModelDeprecationEmailSender", lambda _: runner)
monkeypatch.setattr(module.asyncio, "sleep", AsyncMock(side_effect=[None, asyncio.CancelledError()]))
with pytest.raises(asyncio.CancelledError):
await module.run_scheduled_model_deprecation_emails()
assert runner.run.await_count == 2
assert all(call.args == (router,) for call in runner.run.await_args_list)

18
tests/unit/litellm_core_utils/llm_response_utils/test_get_api_base.py
View file

@ -1,6 +1,5 @@
import json
import pytest
import respx
import litellm
from litellm.litellm_core_utils.llm_response_utils import get_api_base as get_api_base_module
@ -53,17 +52,20 @@ class TestDeclaredAuthenticatingProvider:
assert resolution_lookups == []
assert api_base == expected
def test_copilot_keeps_the_enterprise_endpoint_from_disk(self, isolated_token_dirs, resolution_lookups):
def test_copilot_uses_the_user_endpoint_for_an_enterprise_api_base(self, isolated_token_dirs, resolution_lookups):
token_dir = isolated_token_dirs / "github_copilot"
token_dir.mkdir()
(token_dir / "api-key.json").write_text(
json.dumps({"endpoints": {"api": "https://api.enterprise.githubcopilot.com"}})
)
api_base = litellm.get_api_base(model="github_copilot/gpt-4o", optional_params={})
token = f"gho_{isolated_token_dirs.name}"
(token_dir / "access-token").write_text(token)
with respx.mock(assert_all_called=True) as router:
route = router.get("https://api.github.com/copilot_internal/user").respond(
json={"endpoints": {"api": "https://api.enterprise.githubcopilot.com"}}
)
api_base = litellm.get_api_base(model="github_copilot/gpt-4o", optional_params={})
assert resolution_lookups == []
assert api_base == "https://api.enterprise.githubcopilot.com"
assert route.calls.last.request.headers["authorization"] == f"token {token}"
def test_explicit_api_base_still_wins(self, isolated_token_dirs, resolution_lookups):
api_base = litellm.get_api_base(

10
tests/unit/litellm_core_utils/test_safe_json_dumps.py
View file

@ -230,9 +230,7 @@ def test_pydantic_base_model():
inner: InnerModel
tags: list
outer = OuterModel(
name="test", inner=InnerModel(value=42, label="hello"), tags=["a", "b"]
)
outer = OuterModel(name="test", inner=InnerModel(value=42, label="hello"), tags=["a", "b"])
# Test a pydantic model at the top level
result = json.loads(safe_dumps(outer))
@ -261,3 +259,9 @@ def test_safe_json_structure_keeps_tuples_and_drops_non_string_keys():
assert structure == {"models": ("A", "B"), "tags": ["X", "Y"], "nested": {"deep": ("C",)}}
assert type(structure["models"]) is tuple
assert json.loads(safe_dumps(data)) == {"models": ["a", "b"], "tags": ["x", "y"], "nested": {"deep": ["c"]}}
def test_integer_limit_clamps_both_signs_and_leaves_bools_and_floats_alone() -> None:
assert safe_json_structure(
{"big": 10**400, "small": -(10**400), "inside": 7, "flag": True, "ratio": 1e300}, integer_limit=100
) == {"big": 100, "small": -100, "inside": 7, "flag": True, "ratio": 1e300}

7
tests/unit/llms/github_copilot/embedding/test_github_copilot_embedding_transformation.py
View file

@ -1,15 +1,14 @@
from unittest.mock import MagicMock, patch
from unittest.mock import MagicMock
import httpx
import pytest
from pydantic import ValidationError
from litellm.exceptions import AuthenticationError
from litellm.llms.github_copilot.common_utils import GetAccessTokenError
from litellm.llms.github_copilot.embedding.transformation import (
GithubCopilotEmbeddingConfig,
)
from litellm.llms.github_copilot.common_utils import GetAPIKeyError
from litellm.types.utils import EmbeddingResponse
@ -41,7 +40,7 @@ def test_github_copilot_embedding_config_validate_environment():
assert "x-request-id" in validated_headers
# Test with authentication failure
config.authenticator.get_api_key.side_effect = GetAPIKeyError(
config.authenticator.get_api_key.side_effect = GetAccessTokenError(
message="Failed to get API key",
status_code=401,
)

7
tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py
View file

@ -4,9 +4,8 @@ from unittest.mock import MagicMock
import httpx
import pytest
from litellm.exceptions import AuthenticationError
from litellm.llms.github_copilot.common_utils import GetAPIKeyError
from litellm.llms.github_copilot.common_utils import GetAccessTokenError
from litellm.llms.github_copilot.messages.transformation import (
GithubCopilotAnthropicMessagesConfig,
)
@ -216,7 +215,9 @@ def test_github_copilot_anthropic_messages_validate_environment_auth_error():
# Mock the authenticator to raise an error
config.authenticator = MagicMock()
config.authenticator.get_api_key.side_effect = GetAPIKeyError(status_code=401, message="No valid API key found")
config.authenticator.get_api_key.side_effect = GetAccessTokenError(
status_code=401, message="No valid API key found"
)
with pytest.raises(AuthenticationError):
config.validate_anthropic_messages_environment(

807
tests/unit/llms/github_copilot/test_github_copilot_authenticator.py
View file

@ -1,559 +1,384 @@
import json
import logging
import os
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timedelta
from pathlib import Path
from types import MappingProxyType
from typing import Final
from unittest.mock import MagicMock, mock_open, patch
from unittest.mock import MagicMock, patch
import pytest
import respx
from litellm.llms.github_copilot.authenticator import Authenticator
from litellm.llms.github_copilot.common_utils import (
APIKeyExpiredError,
GetAccessTokenError,
GetAPIKeyError,
GetDeviceCodeError,
RefreshAPIKeyError,
)
from litellm.llms.github_copilot.authenticator import Authenticator, github_api_headers
from litellm.llms.github_copilot.common_utils import GetAccessTokenError, GetDeviceCodeError
API_KEY_URL: Final = "https://copilot-key.test/copilot_internal/v2/token"
REFRESHED_KEY: Final = MappingProxyType({"token": "new-api-key", "expires_at": 4102444800})
COPILOT_USER_URL: Final = "https://api.github.com/copilot_internal/user"
DEVICE_CODE_URL: Final = "https://github.com/login/device/code"
TENANT_API: Final = "https://api.tenant.githubcopilot.com"
REPAIRED_API: Final = "https://api.repaired.githubcopilot.com"
TRUNCATED_UTF_16: Final = b"\xff\xfe\x00\xd8"
@pytest.fixture
def logged_in_token_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
def logged_in_token_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> tuple[Path, str]:
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
monkeypatch.setenv("GITHUB_COPILOT_API_KEY_URL", API_KEY_URL)
(tmp_path / "access-token").write_text("gho_logged_in")
return tmp_path
token: Final = f"gho_{tmp_path.name}"
(tmp_path / "access-token").write_text(token)
return tmp_path, token
class TestGitHubCopilotAuthenticator:
@pytest.fixture
def authenticator(self):
with (
patch("os.path.exists", return_value=False),
patch("os.makedirs") as mock_makedirs,
):
auth = Authenticator()
mock_makedirs.assert_called_once()
return auth
def test_init_uses_access_token_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
auth: Final = Authenticator()
assert auth.token_dir == str(tmp_path)
assert auth.access_token_file == str(tmp_path / "access-token")
@pytest.fixture
def mock_http_client(self):
mock_client = MagicMock()
mock_response = MagicMock()
mock_client.get.return_value = mock_response
mock_client.post.return_value = mock_response
mock_response.raise_for_status.return_value = None
return mock_client, mock_response
def test_init(self):
"""Test the initialization of the authenticator."""
with (
patch("os.path.exists", return_value=False),
patch("os.makedirs") as mock_makedirs,
):
auth = Authenticator()
assert auth.token_dir.endswith("/github_copilot")
assert auth.access_token_file.endswith("/access-token")
assert auth.api_key_file.endswith("/api-key.json")
mock_makedirs.assert_called_once()
def test_init_uses_default_access_token_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("HOME", str(tmp_path))
monkeypatch.delenv("GITHUB_COPILOT_TOKEN_DIR", raising=False)
def test_ensure_token_dir(self):
"""Test that the token directory is created if it doesn't exist."""
with (
patch("os.path.exists", return_value=False),
patch("os.makedirs") as mock_makedirs,
):
auth = Authenticator()
mock_makedirs.assert_called_once_with(auth.token_dir, exist_ok=True)
auth: Final = Authenticator()
def test_get_github_headers(self, authenticator):
"""Test that GitHub headers are correctly generated."""
headers = authenticator._get_github_headers()
assert "accept" in headers
assert "editor-version" in headers
assert "user-agent" in headers
assert "content-type" in headers
assert auth.token_dir == str(tmp_path / ".config/litellm/github_copilot")
assert auth.access_token_file == str(tmp_path / ".config/litellm/github_copilot/access-token")
headers_with_token = authenticator._get_github_headers("test-token")
assert headers_with_token["authorization"] == "token test-token"
def test_get_access_token_from_file(self, logged_in_token_dir: Path) -> None:
"""Test retrieving an access token from a file."""
assert Authenticator().get_access_token() == "gho_logged_in"
def test_ensure_token_dir_creates_missing_directory(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
token_dir: Final = str(tmp_path / "github_copilot")
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", token_dir)
def test_get_access_token_login(self, authenticator):
"""Test logging in to get an access token."""
mock_token = "mock-access-token"
auth: Final = Authenticator()
with (
patch.object(authenticator, "_login", return_value=mock_token),
patch("builtins.open", mock_open()),
patch("builtins.open", side_effect=IOError) as mock_read,
):
token = authenticator.get_access_token()
assert token == mock_token
authenticator._login.assert_called_once()
assert Path(auth.token_dir).is_dir()
@pytest.mark.asyncio
async def test_get_access_token_refuses_device_code_login_in_event_loop(self, authenticator):
with (
patch("builtins.open", side_effect=FileNotFoundError),
patch.object(authenticator, "_login") as mock_login,
):
with pytest.raises(GetAccessTokenError) as exc:
authenticator.get_access_token()
assert exc.value.status_code == 401
assert "event loop" in str(exc.value)
assert authenticator.access_token_file not in str(exc.value)
mock_login.assert_not_called()
def test_get_github_headers_include_access_token() -> None:
headers: Final = github_api_headers("github-token")
def test_get_access_token_refuses_device_code_login_in_worker_thread(self, authenticator):
with (
patch("builtins.open", side_effect=FileNotFoundError),
patch.object(authenticator, "_login") as mock_login,
):
with ThreadPoolExecutor(max_workers=1) as pool:
with pytest.raises(GetAccessTokenError) as exc:
pool.submit(authenticator.get_access_token).result()
assert headers == {
"accept": "application/json",
"editor-version": "vscode/1.85.1",
"editor-plugin-version": "copilot/1.155.0",
"user-agent": "GithubCopilot/1.155.0",
"accept-encoding": "gzip,deflate,br",
"content-type": "application/json",
"authorization": "token github-token",
}
assert exc.value.status_code == 401
assert "worker thread" in str(exc.value)
assert authenticator.access_token_file not in str(exc.value)
mock_login.assert_not_called()
@pytest.mark.parametrize("encoding", ["utf-16", "utf-16-le", "utf-8-sig"])
def test_saved_access_token_in_another_encoding_is_sent_on_refresh(
self, logged_in_token_dir: Path, encoding: str
) -> None:
"""An access-token file a Windows editor saved as UTF-16 or with a BOM still refreshes the key."""
(logged_in_token_dir / "access-token").write_bytes("gho_logged_in\n".encode(encoding))
def test_get_access_token_reads_saved_token(logged_in_token_dir: tuple[Path, str]) -> None:
_, token = logged_in_token_dir
assert Authenticator().get_access_token() == token
with respx.mock(assert_all_called=True) as router:
route: Final = router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY))
assert Authenticator().get_api_key() == "new-api-key"
assert route.calls.last.request.headers["authorization"] == "token gho_logged_in"
def test_get_access_token_logs_in_and_saves_token(
logged_in_token_dir: tuple[Path, str],
) -> None:
token_dir, _ = logged_in_token_dir
access_token_file: Final = token_dir / "access-token"
access_token_file.unlink()
auth: Final = Authenticator()
def test_undecodable_access_token_is_refused_like_a_missing_one(self, logged_in_token_dir: Path) -> None:
"""A proxy worker gets the login hint for an access-token file it cannot decode, and the file is kept."""
access_token_file: Final = logged_in_token_dir / "access-token"
access_token_file.write_bytes(TRUNCATED_UTF_16)
with patch.object(auth, "_login", return_value="github-access-token") as login:
assert auth.get_access_token() == "github-access-token"
with ThreadPoolExecutor(max_workers=1) as pool, pytest.raises(GetAccessTokenError) as exc:
pool.submit(Authenticator().get_access_token).result()
login.assert_called_once()
assert access_token_file.read_text() == "github-access-token"
assert "worker thread" in str(exc.value)
assert access_token_file.read_bytes() == TRUNCATED_UTF_16
def test_get_access_token_failure(self, authenticator):
"""Test that an exception is raised after multiple login failures."""
with (
patch.object(
authenticator,
"_login",
side_effect=GetDeviceCodeError(message="Test error", status_code=400),
),
patch("builtins.open", side_effect=IOError),
):
with pytest.raises(GetAccessTokenError):
authenticator.get_access_token()
assert authenticator._login.call_count == 3
def test_get_api_key_returns_github_token_without_http(
logged_in_token_dir: tuple[Path, str],
) -> None:
_, token = logged_in_token_dir
with patch("litellm.llms.github_copilot.authenticator.get_httpx_client") as get_client:
assert Authenticator().get_api_key() == token
get_client.assert_not_called()
def test_get_api_key_from_file(self, authenticator):
"""Test retrieving an API key from a file."""
future_time = (datetime.now() + timedelta(hours=1)).timestamp()
mock_api_key_data = json.dumps(
{"token": "mock-api-key", "expires_at": future_time}
).encode()
with patch("builtins.open", mock_open(read_data=mock_api_key_data)):
api_key = authenticator.get_api_key()
assert api_key == "mock-api-key"
def test_get_api_key_expired(self, authenticator):
"""Test refreshing an expired API key."""
past_time = (datetime.now() - timedelta(hours=1)).timestamp()
mock_expired_data = json.dumps(
{"token": "expired-api-key", "expires_at": past_time}
).encode()
mock_new_data = {
"token": "new-api-key",
"expires_at": (datetime.now() + timedelta(hours=1)).timestamp(),
}
with (
patch("builtins.open", mock_open(read_data=mock_expired_data)),
patch.object(authenticator, "_refresh_api_key", return_value=mock_new_data),
patch("json.dump") as mock_json_dump,
):
api_key = authenticator.get_api_key()
assert api_key == "new-api-key"
authenticator._refresh_api_key.assert_called_once()
@pytest.mark.parametrize(
"malformed",
[
b"[]",
json.dumps({"token": "tid=abc", "expires_at": "soon"}).encode(),
TRUNCATED_UTF_16,
json.dumps({"token": "tid=abc", "expires_at": 4102444800, "endpoints": []}).encode(),
],
ids=["json-list", "text-expires-at", "truncated-utf-16", "list-endpoints-beside-live-key"],
)
def test_get_api_key_malformed_file_refreshes_like_missing(
self, logged_in_token_dir: Path, malformed: bytes
) -> None:
"""A file that is not a Copilot key object is ignored, refreshed over, never answered with a parse error."""
api_key_file: Final = logged_in_token_dir / "api-key.json"
api_key_file.write_bytes(malformed)
with respx.mock(assert_all_called=True) as router:
route: Final = router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY))
assert Authenticator().get_api_key() == "new-api-key"
assert route.call_count == 1
assert route.calls.last.request.headers["authorization"] == "token gho_logged_in"
assert json.loads(api_key_file.read_text()) == REFRESHED_KEY
@pytest.mark.parametrize(
"malformed",
[
b"[]",
json.dumps({"token": "tid=abc", "expires_at": "soon"}).encode(),
TRUNCATED_UTF_16,
json.dumps({"token": "tid=abc", "expires_at": 4102444800, "endpoints": []}).encode(),
],
ids=["json-list", "text-expires-at", "truncated-utf-16", "list-endpoints-beside-live-key"],
)
def test_get_api_key_malformed_file_is_kept_when_refresh_fails(
self, logged_in_token_dir: Path, malformed: bytes
) -> None:
"""Ignoring the file never deletes it, so a failed refresh leaves the operator's file for inspection."""
api_key_file: Final = logged_in_token_dir / "api-key.json"
api_key_file.write_bytes(malformed)
with respx.mock(assert_all_called=True) as router:
router.get(API_KEY_URL).respond(401)
with pytest.raises(GetAPIKeyError):
Authenticator().get_api_key()
assert api_key_file.read_bytes() == malformed
def test_refreshed_key_that_cannot_be_saved_is_reused_with_its_endpoint(self, logged_in_token_dir: Path) -> None:
unwritable_key_file: Final = logged_in_token_dir / "api-key.json"
unwritable_key_file.mkdir()
first: Final = Authenticator()
second: Final = Authenticator()
with respx.mock(assert_all_called=True) as router:
route: Final = router.get(API_KEY_URL).respond(json={**REFRESHED_KEY, "endpoints": {"api": TENANT_API}})
served: Final = (first.get_api_key(), second.get_api_key(), first.get_api_key())
assert served == ("new-api-key",) * 3
assert (first.get_api_base(), second.get_api_base()) == (TENANT_API, TENANT_API)
assert route.call_count == 1
assert unwritable_key_file.is_dir()
def test_refreshed_key_that_cannot_be_saved_and_has_expired_is_refreshed_again(
self, logged_in_token_dir: Path
) -> None:
(logged_in_token_dir / "api-key.json").mkdir()
def test_get_api_base_uses_user_endpoint_and_memoizes_per_access_token(
logged_in_token_dir: tuple[Path, str],
) -> None:
_, token = logged_in_token_dir
with respx.mock(assert_all_called=True) as router:
route: Final = router.get(COPILOT_USER_URL).respond(json={"endpoints": {"api": TENANT_API}})
auth: Final = Authenticator()
assert auth.get_api_base() == TENANT_API
assert Authenticator().get_api_base() == TENANT_API
with respx.mock(assert_all_called=True) as router:
route: Final = router.get(API_KEY_URL).respond(json={"token": "expired-api-key", "expires_at": 1})
served: Final = (auth.get_api_key(), auth.get_api_key())
assert route.call_count == 1
assert route.calls.last.request.headers["authorization"] == f"token {token}"
assert served == ("expired-api-key",) * 2
assert route.call_count == 2
def test_kept_key_and_its_endpoint_stay_paired_after_the_key_file_is_repaired(
self, logged_in_token_dir: Path
) -> None:
key_file: Final = logged_in_token_dir / "api-key.json"
key_file.mkdir()
auth: Final = Authenticator()
with respx.mock(assert_all_called=True) as router:
router.get(API_KEY_URL).respond(json={**REFRESHED_KEY, "endpoints": {"api": TENANT_API}})
auth.get_api_key()
key_file.rmdir()
key_file.write_text(
json.dumps({"token": "repaired-api-key", "expires_at": 9999999999, "endpoints": {"api": REPAIRED_API}})
)
assert (auth.get_api_key(), auth.get_api_base()) in (
("new-api-key", TENANT_API),
("repaired-api-key", REPAIRED_API),
)
def test_unsaved_key_is_not_reused_after_the_access_token_changes(self, logged_in_token_dir: Path) -> None:
(logged_in_token_dir / "api-key.json").mkdir()
with respx.mock(assert_all_called=True) as router:
route: Final = router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY))
Authenticator().get_api_key()
(logged_in_token_dir / "access-token").write_text("gho_logged_in_again")
Authenticator().get_api_key()
assert [call.request.headers["authorization"] for call in route.calls] == [
"token gho_logged_in",
"token gho_logged_in_again",
]
@pytest.mark.parametrize(
"malformed",
[
b"[]",
json.dumps({"endpoints": {"api": 5}}).encode(),
json.dumps({"token": "t", "expires_at": 1, "endpoints": []}).encode(),
TRUNCATED_UTF_16,
],
ids=["json-list", "number-api", "list-endpoints", "truncated-utf-16"],
)
def test_get_api_base_malformed_file_returns_none(
self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, malformed: bytes
) -> None:
"""get_api_base runs before get_api_key on every request, so it must survive malformed shapes too."""
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
(tmp_path / "api-key.json").write_bytes(malformed)
@pytest.mark.parametrize("status_code", [401, 403, 500])
def test_get_api_base_returns_none_when_user_lookup_fails(
logged_in_token_dir: tuple[Path, str], status_code: int
) -> None:
with respx.mock(assert_all_called=True) as router:
router.get(COPILOT_USER_URL).respond(status_code=status_code)
assert Authenticator().get_api_base() is None
def test_get_api_base_keeps_endpoints_when_the_key_is_malformed(self, logged_in_token_dir: Path) -> None:
"""The endpoint is read on its own, so a bad key never sends the refreshed key to the default host."""
(logged_in_token_dir / "api-key.json").write_text(
json.dumps({"token": "tid=abc", "expires_at": "soon", "endpoints": {"api": TENANT_API}})
)
auth: Final = Authenticator()
def test_get_api_base_retries_after_failed_user_lookup(logged_in_token_dir: tuple[Path, str]) -> None:
auth: Final = Authenticator()
with respx.mock(assert_all_called=True) as router:
failed_lookup: Final = router.get(COPILOT_USER_URL).respond(status_code=503)
assert auth.get_api_base() is None
assert failed_lookup.call_count == 1
with respx.mock(assert_all_called=True) as router:
successful_lookup: Final = router.get(COPILOT_USER_URL).respond(json={"endpoints": {"api": TENANT_API}})
assert auth.get_api_base() == TENANT_API
with respx.mock(assert_all_called=True) as router:
router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY))
assert auth.get_api_key() == "new-api-key"
assert successful_lookup.call_count == 1
@pytest.mark.parametrize("encoding", ["utf-16", "utf-16-le", "utf-8-sig"])
def test_live_key_file_in_another_json_encoding_is_used_without_a_refresh(
self, logged_in_token_dir: Path, encoding: str
) -> None:
"""A key file a Windows editor saved as UTF-16 or with a BOM still holds a live key and tenant."""
live_key: Final = json.dumps({"token": "tid=abc", "expires_at": 4102444800, "endpoints": {"api": TENANT_API}})
api_key_file: Final = logged_in_token_dir / "api-key.json"
api_key_file.write_bytes(live_key.encode(encoding))
auth: Final = Authenticator()
with respx.mock:
assert auth.get_api_base() == TENANT_API
assert auth.get_api_key() == "tid=abc"
def test_get_api_base_returns_none_when_user_lookup_response_has_no_endpoint(
logged_in_token_dir: tuple[Path, str],
) -> None:
with respx.mock(assert_all_called=True) as router:
router.get(COPILOT_USER_URL).respond(json={"endpoints": {}})
assert Authenticator().get_api_base() is None
assert api_key_file.read_bytes() == live_key.encode(encoding)
def test_ignored_key_file_is_logged_without_its_contents(
self, logged_in_token_dir: Path, caplog: pytest.LogCaptureFixture
) -> None:
"""The warning names the file but never echoes a token saved in the wrong shape."""
(logged_in_token_dir / "api-key.json").write_text(
json.dumps({"token": ["tid=secret-value"], "expires_at": 4102444800})
)
def test_get_api_base_returns_none_when_user_lookup_endpoints_are_not_a_mapping(
logged_in_token_dir: tuple[Path, str],
) -> None:
with respx.mock(assert_all_called=True) as router:
router.get(COPILOT_USER_URL).respond(json={"endpoints": []})
assert Authenticator().get_api_base() is None
with caplog.at_level(logging.WARNING, logger="LiteLLM"), respx.mock(assert_all_called=True) as router:
router.get(API_KEY_URL).respond(json=dict(REFRESHED_KEY))
assert Authenticator().get_api_key() == "new-api-key"
assert "not a Copilot API key file" in caplog.text
assert "tid=secret-value" not in caplog.text
def test_get_api_base_rejects_untrusted_user_endpoint(
logged_in_token_dir: tuple[Path, str],
) -> None:
with respx.mock(assert_all_called=True) as router:
router.get(COPILOT_USER_URL).respond(json={"endpoints": {"api": "https://attacker.example"}})
assert Authenticator().get_api_base() is None
def test_refresh_api_key(self, authenticator, mock_http_client):
"""Test refreshing an API key."""
mock_client, mock_response = mock_http_client
mock_token = "mock-access-token"
mock_api_key_data = {"token": "new-api-key", "expires_at": 12345}
with (
patch.object(authenticator, "get_access_token", return_value=mock_token),
patch(
"litellm.llms.github_copilot.authenticator.get_httpx_client",
return_value=mock_client,
),
patch.object(mock_response, "json", return_value=mock_api_key_data),
):
result = authenticator._refresh_api_key()
assert result == mock_api_key_data
mock_client.get.assert_called_once()
authenticator.get_access_token.assert_called_once()
def test_get_api_base_does_not_trigger_device_login_without_saved_token(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
with patch("litellm.llms.github_copilot.authenticator.get_httpx_client") as get_client:
assert Authenticator().get_api_base() is None
get_client.assert_not_called()
def test_refresh_api_key_failure(self, authenticator, mock_http_client):
"""Test failure to refresh an API key."""
mock_client, mock_response = mock_http_client
mock_token = "mock-access-token"
with (
patch.object(authenticator, "get_access_token", return_value=mock_token),
patch(
"litellm.llms.github_copilot.authenticator.get_httpx_client",
return_value=mock_client,
),
patch.object(mock_response, "json", return_value={}),
):
with pytest.raises(RefreshAPIKeyError):
authenticator._refresh_api_key()
assert mock_client.get.call_count == 3
@pytest.mark.asyncio
async def test_get_access_token_refuses_device_code_login_in_event_loop() -> None:
auth: Final = Authenticator()
with (
patch("builtins.open", side_effect=FileNotFoundError),
patch.object(auth, "_login") as login,
):
with pytest.raises(GetAccessTokenError) as exc:
auth.get_access_token()
def test_get_device_code(self, authenticator, mock_http_client):
"""Test getting a device code."""
mock_client, mock_response = mock_http_client
mock_device_code_data = {
"device_code": "mock-device-code",
"user_code": "ABCD-EFGH",
"verification_uri": "https://github.com/login/device",
}
assert exc.value.status_code == 401
assert "event loop" in str(exc.value)
login.assert_not_called()
with (
patch(
"litellm.llms.github_copilot.authenticator.get_httpx_client",
return_value=mock_client,
),
patch.object(mock_response, "json", return_value=mock_device_code_data),
):
result = authenticator._get_device_code()
assert result == mock_device_code_data
mock_client.post.assert_called_once()
def test_poll_for_access_token(self, authenticator, mock_http_client):
"""Test polling for an access token."""
mock_client, mock_response = mock_http_client
mock_token_data = {"access_token": "mock-access-token"}
def test_get_access_token_refuses_device_code_login_in_worker_thread() -> None:
auth: Final = Authenticator()
with (
patch("builtins.open", side_effect=FileNotFoundError),
patch.object(auth, "_login") as login,
):
with ThreadPoolExecutor(max_workers=1) as pool:
with pytest.raises(GetAccessTokenError) as exc:
pool.submit(auth.get_access_token).result()
with (
patch(
"litellm.llms.github_copilot.authenticator.get_httpx_client",
return_value=mock_client,
),
patch.object(mock_response, "json", return_value=mock_token_data),
patch("time.sleep"),
):
result = authenticator._poll_for_access_token("mock-device-code")
assert result == "mock-access-token"
mock_client.post.assert_called_once()
assert exc.value.status_code == 401
assert "worker thread" in str(exc.value)
login.assert_not_called()
def test_login(self, authenticator):
"""Test the login process."""
mock_device_code_data = {
"device_code": "mock-device-code",
"user_code": "ABCD-EFGH",
"verification_uri": "https://github.com/login/device",
}
mock_token = "mock-access-token"
with (
patch.object(
authenticator, "_get_device_code", return_value=mock_device_code_data
),
patch.object(
authenticator, "_poll_for_access_token", return_value=mock_token
),
patch("builtins.print") as mock_print,
):
result = authenticator._login()
assert result == mock_token
authenticator._get_device_code.assert_called_once()
authenticator._poll_for_access_token.assert_called_once_with(
"mock-device-code"
)
mock_print.assert_called_once()
@pytest.mark.parametrize("encoding", ["utf-16", "utf-16-le", "utf-8-sig"])
def test_get_access_token_reads_saved_token_with_supported_encoding(
logged_in_token_dir: tuple[Path, str], encoding: str
) -> None:
token_dir, token = logged_in_token_dir
(token_dir / "access-token").write_bytes(f"{token}\n".encode(encoding))
def test_get_api_base_from_file(self, authenticator):
"""Test retrieving the API base endpoint from a file."""
mock_api_key_data = json.dumps(
assert Authenticator().get_access_token() == token
def test_get_access_token_refuses_undecodable_saved_token(
logged_in_token_dir: tuple[Path, str],
) -> None:
token_dir, _ = logged_in_token_dir
access_token_file: Final = token_dir / "access-token"
access_token_file.write_bytes(TRUNCATED_UTF_16)
with ThreadPoolExecutor(max_workers=1) as pool, pytest.raises(GetAccessTokenError) as exc:
pool.submit(Authenticator().get_access_token).result()
assert exc.value.status_code == 401
assert "worker thread" in str(exc.value)
assert access_token_file.read_bytes() == TRUNCATED_UTF_16
def test_get_access_token_retries_device_flow_three_times(
logged_in_token_dir: tuple[Path, str],
) -> None:
token_dir, _ = logged_in_token_dir
(token_dir / "access-token").unlink()
auth: Final = Authenticator()
with (
patch.object(
auth,
"_login",
side_effect=GetDeviceCodeError(message="device flow failed", status_code=400),
) as login,
pytest.raises(GetAccessTokenError, match="after 3 attempts"),
):
auth.get_access_token()
assert login.call_count == 3
def test_get_access_token_returns_login_token_when_file_write_fails(
logged_in_token_dir: tuple[Path, str],
) -> None:
token_dir, _ = logged_in_token_dir
(token_dir / "access-token").unlink()
auth: Final = Authenticator()
with (
patch.object(auth, "_login", return_value="github-access-token"),
patch("builtins.open", side_effect=(FileNotFoundError, PermissionError)),
):
assert auth.get_access_token() == "github-access-token"
@pytest.mark.parametrize(
("environment", "url", "body"),
[
(
{"GITHUB_COPILOT_DEVICE_CODE_URL": "https://custom.example/device"},
"https://custom.example/device",
{"client_id": "Iv1.b507a08c87ecfe98", "scope": "read:user"},
),
(
{"GITHUB_COPILOT_CLIENT_ID": "custom-client"},
DEVICE_CODE_URL,
{"client_id": "custom-client", "scope": "read:user"},
),
],
)
def test_get_device_code_uses_runtime_url_and_client_id(
environment: dict[str, str], url: str, body: dict[str, str]
) -> None:
client: Final = MagicMock()
response: Final = MagicMock()
response.json.return_value = {
"device_code": "dc",
"user_code": "UC",
"verification_uri": "https://github.com/login/device",
}
response.raise_for_status.return_value = None
client.post.return_value = response
with (
patch.dict(os.environ, environment),
patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=client),
):
assert Authenticator()._get_device_code()["device_code"] == "dc"
client.post.assert_called_once()
assert client.post.call_args.args[0] == url
assert client.post.call_args.kwargs["json"] == body
def test_get_device_code_sends_device_flow_request(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("GITHUB_COPILOT_DEVICE_CODE_URL", raising=False)
monkeypatch.delenv("GITHUB_COPILOT_CLIENT_ID", raising=False)
client: Final = MagicMock()
response: Final = MagicMock()
device_code: Final = {
"device_code": "device-code",
"user_code": "ABCD-EFGH",
"verification_uri": "https://github.com/login/device",
}
response.json.return_value = device_code
response.raise_for_status.return_value = None
client.post.return_value = response
with patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=client):
assert Authenticator()._get_device_code() == device_code
client.post.assert_called_once_with(
DEVICE_CODE_URL,
headers=github_api_headers(),
json={"client_id": "Iv1.b507a08c87ecfe98", "scope": "read:user"},
)
def test_poll_for_access_token_sends_device_code_and_grant_type(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.delenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", raising=False)
monkeypatch.delenv("GITHUB_COPILOT_CLIENT_ID", raising=False)
client: Final = MagicMock()
response: Final = MagicMock()
response.json.return_value = {"access_token": "github-access-token"}
response.raise_for_status.return_value = None
client.post.return_value = response
with patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=client):
assert Authenticator()._poll_for_access_token("device-code") == "github-access-token"
client.post.assert_called_once_with(
"https://github.com/login/oauth/access_token",
headers=github_api_headers(),
json={
"client_id": "Iv1.b507a08c87ecfe98",
"device_code": "device-code",
"grant_type": "urn:ietf:params:oauth:grant-type:device_code",
},
)
def test_poll_for_access_token_uses_runtime_url_and_client_id() -> None:
client: Final = MagicMock()
response: Final = MagicMock()
response.json.return_value = {"access_token": "github-token"}
response.raise_for_status.return_value = None
client.post.return_value = response
with (
patch.dict(
os.environ,
{
"token": "mock-api-key",
"expires_at": (datetime.now() + timedelta(hours=1)).timestamp(),
"endpoints": {"api": "https://api.enterprise.githubcopilot.com"},
}
).encode()
with patch("builtins.open", mock_open(read_data=mock_api_key_data)):
api_base = authenticator.get_api_base()
assert api_base == "https://api.enterprise.githubcopilot.com"
"GITHUB_COPILOT_ACCESS_TOKEN_URL": "https://custom.example/token",
"GITHUB_COPILOT_CLIENT_ID": "custom-client",
},
),
patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=client),
patch("time.sleep"),
):
assert Authenticator()._poll_for_access_token("dc") == "github-token"
assert client.post.call_args.args[0] == "https://custom.example/token"
assert client.post.call_args.kwargs["json"]["client_id"] == "custom-client"
def test_get_device_code_with_custom_url(self, authenticator, mock_http_client):
"""GITHUB_COPILOT_DEVICE_CODE_URL env var must be used by _get_device_code at call time."""
mock_client, mock_response = mock_http_client
custom_url = "https://custom.example.com/device"
mock_response.json.return_value = {
"device_code": "dc",
"user_code": "UC",
"verification_uri": "https://example.com",
}
with (
patch.dict(os.environ, {"GITHUB_COPILOT_DEVICE_CODE_URL": custom_url}),
patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=mock_client),
):
authenticator._get_device_code()
assert mock_client.post.call_args[0][0] == custom_url
def test_get_device_code_with_custom_client_id(self, authenticator, mock_http_client):
"""GITHUB_COPILOT_CLIENT_ID env var must appear as client_id in the device-code request body."""
mock_client, mock_response = mock_http_client
custom_id = "custom_client_id"
mock_response.json.return_value = {
"device_code": "dc",
"user_code": "UC",
"verification_uri": "https://example.com",
}
with (
patch.dict(os.environ, {"GITHUB_COPILOT_CLIENT_ID": custom_id}),
patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=mock_client),
):
authenticator._get_device_code()
assert mock_client.post.call_args[1]["json"]["client_id"] == custom_id
def test_login_requests_device_code_and_polls_for_access_token() -> None:
auth: Final = Authenticator()
device_code: Final = {
"device_code": "device-code",
"user_code": "ABCD-EFGH",
"verification_uri": "https://github.com/login/device",
}
def test_poll_for_access_token_with_custom_url(self, authenticator, mock_http_client):
"""GITHUB_COPILOT_ACCESS_TOKEN_URL env var must be used by _poll_for_access_token at call time."""
mock_client, mock_response = mock_http_client
custom_url = "https://custom.example.com/token"
mock_response.json.return_value = {"access_token": "tok"}
with (
patch.dict(os.environ, {"GITHUB_COPILOT_ACCESS_TOKEN_URL": custom_url}),
patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=mock_client),
patch("time.sleep"),
):
authenticator._poll_for_access_token("dc")
assert mock_client.post.call_args[0][0] == custom_url
with (
patch.object(auth, "_get_device_code", return_value=device_code) as get_device_code,
patch.object(auth, "_poll_for_access_token", return_value="github-access-token") as poll,
patch("builtins.print") as print_message,
):
assert auth._login() == "github-access-token"
def test_poll_for_access_token_with_custom_client_id(self, authenticator, mock_http_client):
"""GITHUB_COPILOT_CLIENT_ID env var must appear as client_id in the polling request body."""
mock_client, mock_response = mock_http_client
custom_id = "custom_client_id"
mock_response.json.return_value = {"access_token": "tok"}
with (
patch.dict(os.environ, {"GITHUB_COPILOT_CLIENT_ID": custom_id}),
patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=mock_client),
patch("time.sleep"),
):
authenticator._poll_for_access_token("dc")
assert mock_client.post.call_args[1]["json"]["client_id"] == custom_id
def test_refresh_api_key_with_custom_url(self, authenticator, mock_http_client):
"""GITHUB_COPILOT_API_KEY_URL env var must be used by _refresh_api_key at call time."""
mock_client, mock_response = mock_http_client
custom_url = "https://custom.example.com/api-key"
mock_response.json.return_value = {"token": "api-tok", "expires_at": 9999999999}
with (
patch.dict(os.environ, {"GITHUB_COPILOT_API_KEY_URL": custom_url}),
patch("litellm.llms.github_copilot.authenticator.get_httpx_client", return_value=mock_client),
patch.object(authenticator, "get_access_token", return_value="access-tok"),
):
authenticator._refresh_api_key()
assert mock_client.get.call_args[0][0] == custom_url
get_device_code.assert_called_once()
poll.assert_called_once_with("device-code")
print_message.assert_called_once_with(
"Please visit https://github.com/login/device and enter code ABCD-EFGH to authenticate.",
flush=True,
)

55
tests/unit/llms/github_copilot/test_github_copilot_transformation.py
View file

@ -1,29 +1,17 @@
import asyncio
import json
from datetime import datetime, timedelta
from typing import AsyncGenerator
from unittest.mock import AsyncMock, MagicMock, mock_open, patch
import pytest
from unittest.mock import MagicMock, patch
import httpx
from respx import MockRouter
import pytest
import litellm
# Import at the top to make the patch work correctly
import litellm.llms.github_copilot.chat.transformation
from litellm import Choices, Message, ModelResponse, Usage, acompletion, completion
from litellm import ModelResponse, completion
from litellm.exceptions import AuthenticationError
from litellm.llms.github_copilot.authenticator import Authenticator
from litellm.llms.github_copilot.chat.transformation import GithubCopilotConfig
from litellm.llms.github_copilot.common_utils import (
APIKeyExpiredError,
GetAccessTokenError,
GetAPIKeyError,
GetDeviceCodeError,
RefreshAPIKeyError,
)
@ -37,9 +25,7 @@ def test_github_copilot_config_get_openai_compatible_provider_info():
config.authenticator = MagicMock()
config.authenticator.get_api_key.return_value = mock_api_key
# Test with dynamic endpoint
config.authenticator.get_api_base.return_value = (
"https://api.enterprise.githubcopilot.com"
)
config.authenticator.get_api_base.return_value = "https://api.enterprise.githubcopilot.com"
# Test with default values
model = "github_copilot/gpt-4"
@ -73,7 +59,7 @@ def test_github_copilot_config_get_openai_compatible_provider_info():
assert api_base == "https://api.githubcopilot.com"
# Test with authentication failure
config.authenticator.get_api_key.side_effect = GetAPIKeyError(
config.authenticator.get_api_key.side_effect = GetAccessTokenError(
message="Failed to get API key",
status_code=401,
)
@ -145,7 +131,6 @@ def test_completion_github_copilot_mock_response(
def test_transform_messages_disable_copilot_system_to_assistant(monkeypatch):
"""Test that system messages are converted to assistant unless disable_copilot_system_to_assistant is True."""
import litellm
from litellm.llms.github_copilot.chat.transformation import GithubCopilotConfig
# Save original value
@ -365,10 +350,6 @@ def test_x_initiator_header_system_only_messages():
assert headers["X-Initiator"] == "user"
def test_copilot_vision_request_header_with_image():
"""Test that Copilot-Vision-Request header is added when messages contain images"""
config = GithubCopilotConfig()
@ -693,13 +674,8 @@ class TestGithubCopilotTransformResponse:
assert result.choices[0].message.tool_calls is not None
assert len(result.choices[0].message.tool_calls) == 1
assert result.choices[0].message.tool_calls[0]["id"] == "toolu_01ABC"
assert (
result.choices[0].message.tool_calls[0]["function"]["name"] == "get_weather"
)
assert (
'"Boston, MA"'
in result.choices[0].message.tool_calls[0]["function"]["arguments"]
)
assert result.choices[0].message.tool_calls[0]["function"]["name"] == "get_weather"
assert '"Boston, MA"' in result.choices[0].message.tool_calls[0]["function"]["arguments"]
def test_transform_response_anthropic_native_multiple_text_blocks(self):
"""All text blocks must be concatenated, not only the first."""
@ -867,12 +843,8 @@ class TestGithubCopilotTransformParsedResponseDict:
@patch("litellm.llms.openai.openai.OpenAIChatCompletion._get_openai_client")
@patch(
"litellm.llms.openai.openai.OpenAIChatCompletion.make_sync_openai_chat_completion_request"
)
def test_openai_handler_repairs_github_copilot_empty_choices(
mock_request, mock_get_client
):
@patch("litellm.llms.openai.openai.OpenAIChatCompletion.make_sync_openai_chat_completion_request")
def test_openai_handler_repairs_github_copilot_empty_choices(mock_request, mock_get_client):
"""
The OpenAI SDK handler calls convert_to_model_response_object directly on the
SDK's parsed output, bypassing transform_response. convert raises APIError on
@ -929,7 +901,10 @@ def test_validate_environment_uses_per_user_session_and_skips_authenticator():
session = GithubCopilotUserSession(token="user-copilot-token", api_base="https://api.githubcopilot.com")
headers = config.validate_environment(
headers={}, model="github_copilot/gpt-4o", messages=[], optional_params={},
headers={},
model="github_copilot/gpt-4o",
messages=[],
optional_params={},
litellm_params={"github_copilot_user_session": session},
)
assert headers["Authorization"] == "Bearer user-copilot-token"
@ -946,9 +921,7 @@ def test_transform_response_carries_upstream_usage():
"object": "chat.completion",
"created": 1,
"model": "github_copilot/gpt-4o",
"choices": [
{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}
],
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18},
},
)

5
tests/unit/llms/github_copilot/test_github_host.py
View file

@ -32,28 +32,25 @@ def test_normalized_github_host_accepts_only_bare_supported_hosts(raw: object, e
def test_github_com_endpoints_read_url_overrides_at_call_time(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("GITHUB_COPILOT_DEVICE_CODE_URL", "https://device.example/code")
monkeypatch.setenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", "https://device.example/token")
monkeypatch.setenv("GITHUB_COPILOT_API_KEY_URL", "https://api.example/copilot")
monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", "custom-client")
endpoints: Final = github_host_endpoints("github.com")
assert endpoints.device_code_url == "https://device.example/code"
assert endpoints.access_token_url == "https://device.example/token"
assert endpoints.api_key_url == "https://api.example/copilot"
assert endpoints.copilot_user_url == "https://api.github.com/copilot_internal/user"
assert endpoints.client_id == "custom-client"
def test_ghe_endpoints_ignore_github_com_url_overrides(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("GITHUB_COPILOT_DEVICE_CODE_URL", "https://device.example/code")
monkeypatch.setenv("GITHUB_COPILOT_ACCESS_TOKEN_URL", "https://device.example/token")
monkeypatch.setenv("GITHUB_COPILOT_API_KEY_URL", "https://api.example/copilot")
monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", "custom-client")
endpoints: Final = github_host_endpoints("octo.ghe.com")
assert endpoints.device_code_url == "https://octo.ghe.com/login/device/code"
assert endpoints.access_token_url == "https://octo.ghe.com/login/oauth/access_token"
assert endpoints.api_key_url == "https://api.octo.ghe.com/copilot_internal/v2/token"
assert endpoints.copilot_user_url == "https://api.octo.ghe.com/copilot_internal/user"
assert endpoints.user_url == "https://api.octo.ghe.com/user"
assert endpoints.default_copilot_api_base == "https://copilot-api.octo.ghe.com"

410
tests/unit/llms/github_copilot/test_per_user_auth.py
View file

@ -12,8 +12,9 @@ import httpx
import pytest
import litellm
from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST
from litellm.constants import GITHUB_COPILOT_DEFAULT_HOST, GITHUB_COPILOT_USER_SESSION_TTL_SECONDS
from litellm.exceptions import (
APIConnectionError,
AuthenticationError,
BadRequestError,
CallerCredentialAuthenticationError,
@ -21,27 +22,27 @@ from litellm.exceptions import (
ServiceUnavailableError,
)
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.github_copilot.github_host import github_host_endpoints
from litellm.llms.github_copilot.per_user_auth import (
_SESSION_CACHE,
GITHUB_COPILOT_USER_SESSION_KWARG_KEY,
GithubCopilotUserSession,
_session_cache_key,
aattach_github_copilot_user_session,
aexchange_github_token,
acheck_copilot_seat,
apoll_device_flow,
aresolve_copilot_user_session,
astart_device_flow,
attach_github_copilot_user_session,
evict_copilot_user_session,
exchange_github_token,
github_copilot_host_endpoints,
github_copilot_auth_mode,
github_copilot_user_session_from,
resolve_copilot_user_session,
validated_copilot_api_base,
)
from litellm.llms.github_copilot.github_host import github_host_endpoints
from litellm.types.utils import CredentialItem, ModelResponse
EXCHANGE_URL = "https://api.github.com/copilot_internal/v2/token"
COPILOT_USER_URL = "https://api.github.com/copilot_internal/user"
DEVICE_CODE_URL = "https://github.com/login/device/code"
ACCESS_TOKEN_URL = "https://github.com/login/oauth/access_token"
_GITHUB_ENDPOINTS: Final = github_host_endpoints("github.com")
@ -62,10 +63,8 @@ def _credential(
)
def _exchange_payload(token="copilot-token", expires_at=None, api_base="https://api.githubcopilot.com"):
def _user_payload(api_base="https://api.githubcopilot.com"):
return {
"token": token,
"expires_at": expires_at if expires_at is not None else int(time.time()) + 3600,
"endpoints": {"api": api_base},
}
@ -92,10 +91,10 @@ def github_http(respond: Callable[[httpx.Request], httpx.Response]):
yield requests
def _exchange_responder(payload=None, status=200):
def _user_responder(payload=None, status=200):
def respond(request: httpx.Request) -> httpx.Response:
if str(request.url) == EXCHANGE_URL:
return httpx.Response(status, json=payload or _exchange_payload(), request=request)
if str(request.url) == COPILOT_USER_URL:
return httpx.Response(status, json=payload or _user_payload(), request=request)
pytest.fail(f"unexpected request to {request.url}")
return respond
@ -170,17 +169,17 @@ def test_validated_copilot_api_base(raw, expected):
"https://copilot-api.octo.ghe.com:8443",
],
)
def test_ghe_exchange_response_must_stay_on_its_tenant_host(raw: object) -> None:
def test_ghe_user_response_must_stay_on_its_tenant_host(raw: object) -> None:
endpoints: Final = github_host_endpoints("octo.ghe.com")
assert validated_copilot_api_base(raw, endpoints) == "https://copilot-api.octo.ghe.com"
def test_github_com_exchange_rejects_a_ghe_copilot_host() -> None:
def test_github_com_user_response_rejects_a_ghe_copilot_host() -> None:
endpoints: Final = github_host_endpoints("github.com")
assert validated_copilot_api_base("https://copilot-api.octo.ghe.com", endpoints) == "https://api.githubcopilot.com"
def test_ghe_exchange_uses_stored_host_not_request_kwargs():
def test_ghe_user_lookup_uses_stored_host_not_request_kwargs():
credential: Final = CredentialItem(
credential_name="copilot-cred",
credential_values={"github_copilot_auth_type": "per_user_oauth", "github_copilot_host": "octo.ghe.com"},
@ -188,8 +187,8 @@ def test_ghe_exchange_uses_stored_host_not_request_kwargs():
)
def respond(request: httpx.Request) -> httpx.Response:
assert str(request.url) == "https://api.octo.ghe.com/copilot_internal/v2/token"
return httpx.Response(200, json=_exchange_payload(api_base="https://copilot-api.octo.ghe.com"), request=request)
assert str(request.url) == "https://api.octo.ghe.com/copilot_internal/user"
return httpx.Response(200, json=_user_payload(api_base="https://copilot-api.octo.ghe.com"), request=request)
with patch.object(litellm, "credential_list", [credential]), github_http(respond) as requests:
kwargs: Final = _kwargs_with_connection(github_host="octo.ghe.com")
@ -201,7 +200,7 @@ def test_ghe_exchange_uses_stored_host_not_request_kwargs():
assert len(requests) == 1
def test_invalid_stored_host_fails_before_exchange():
def test_invalid_stored_host_fails_before_user_lookup():
credential: Final = CredentialItem(
credential_name="copilot-cred",
credential_values={"github_copilot_auth_type": "per_user_oauth", "github_copilot_host": "https://octo.ghe.com"},
@ -225,35 +224,36 @@ def test_session_repr_and_str_hide_the_token(caplog):
assert "secret-copilot-token" not in caplog.text
def test_exchange_returns_session_and_caches_per_user():
with github_http(_exchange_responder()) as requests:
s1 = exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
def test_user_lookup_returns_github_token_and_caches_per_user():
api_base: Final = "https://tenant.githubcopilot.com"
with github_http(_user_responder(payload=_user_payload(api_base=api_base))) as requests:
s1 = resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
assert len(requests) == 1
assert str(requests[0].url) == COPILOT_USER_URL
assert requests[0].headers["Authorization"] == "token gh-token"
s2 = exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
s2 = resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
assert len(requests) == 1
assert s2.token == s1.token == "copilot-token"
# different user -> separate exchange
exchange_github_token("user-b", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
assert s2.token == s1.token == "gh-token"
assert s1.api_base == api_base
resolve_copilot_user_session("user-b", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
assert len(requests) == 2
# different github token -> separate exchange
exchange_github_token("user-a", "gh-token-2", "copilot-cred", _GITHUB_ENDPOINTS)
resolve_copilot_user_session("user-a", "gh-token-2", "copilot-cred", _GITHUB_ENDPOINTS)
assert len(requests) == 3
def test_cached_session_is_not_reused_after_github_token_changes():
with github_http(_exchange_responder()) as requests:
exchange_github_token("user-a", "gh-token-a", "copilot-cred", _GITHUB_ENDPOINTS)
exchange_github_token("user-a", "gh-token-b", "copilot-cred", _GITHUB_ENDPOINTS)
with github_http(_user_responder()) as requests:
resolve_copilot_user_session("user-a", "gh-token-a", "copilot-cred", _GITHUB_ENDPOINTS)
resolve_copilot_user_session("user-a", "gh-token-b", "copilot-cred", _GITHUB_ENDPOINTS)
assert [request.headers["Authorization"] for request in requests] == ["token gh-token-a", "token gh-token-b"]
@pytest.mark.parametrize("status", [401, 403, 404])
def test_exchange_auth_failures_raise_caller_auth_error_and_evict(status):
with github_http(_exchange_responder(status=status)):
def test_user_lookup_auth_failures_raise_caller_auth_error_and_evict(status):
with github_http(_user_responder(status=status)):
with pytest.raises(CallerCredentialAuthenticationError) as exc:
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
assert "Reconnect GitHub Copilot for credential 'copilot-cred' in the LiteLLM UI (LLM Credentials)" in str(
exc.value
)
@ -262,52 +262,135 @@ def test_exchange_auth_failures_raise_caller_auth_error_and_evict(status):
)
def test_exchange_429_raises_caller_rate_limit():
with github_http(_exchange_responder(status=429)):
def test_user_lookup_429_raises_caller_rate_limit():
with github_http(_user_responder(status=429)):
with pytest.raises(CallerCredentialRateLimitError):
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
def test_exchange_other_failure_raises_service_unavailable_without_body():
with github_http(_exchange_responder(status=500, payload={"secret": "body"})):
def test_user_lookup_other_failure_raises_service_unavailable_without_body():
with github_http(_user_responder(status=500, payload={"secret": "body"})):
with pytest.raises(ServiceUnavailableError) as exc:
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
assert "body" not in str(exc.value)
def test_exchange_expiry_margin_refetches_near_expiry():
with github_http(_exchange_responder(_exchange_payload(expires_at=int(time.time()) + 30))) as requests:
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
# ttl under the 60s safety margin -> not cached, second call refetches
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
def test_session_cache_uses_fixed_ttl():
with (
patch("litellm.llms.github_copilot.per_user_auth.time.time", return_value=1000) as clock,
github_http(_user_responder()) as requests,
):
resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
clock.return_value = 1000 + GITHUB_COPILOT_USER_SESSION_TTL_SECONDS - 1
resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
assert len(requests) == 1
clock.return_value = 1000 + GITHUB_COPILOT_USER_SESSION_TTL_SECONDS + 1
resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
assert len(requests) == 2
@pytest.mark.asyncio
async def test_aexchange_single_flight_one_http_call():
async def test_acheck_copilot_seat_rejects_models_403(per_user_credential):
def respond(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, json=_exchange_payload(), request=request)
if str(request.url) == COPILOT_USER_URL:
return httpx.Response(200, json=_user_payload(), request=request)
if str(request.url) == "https://api.githubcopilot.com/models":
return httpx.Response(403, request=request)
pytest.fail(f"unexpected request to {request.url}")
with github_http(respond) as requests:
s1, s2 = await asyncio.gather(
aexchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS),
aexchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS),
)
assert len(requests) == 1
assert s1.token == s2.token == "copilot-token"
with pytest.raises(CallerCredentialAuthenticationError):
await acheck_copilot_seat("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
assert [str(request.url) for request in requests] == [
COPILOT_USER_URL,
"https://api.githubcopilot.com/models",
]
assert requests[1].headers["Authorization"] == "Bearer gh-token"
@pytest.mark.asyncio
async def test_aexchange_401_raises_and_evicts():
with github_http(_exchange_responder(status=401)):
@pytest.mark.parametrize(
("status", "error_type"),
[(429, CallerCredentialRateLimitError), (500, ServiceUnavailableError)],
)
async def test_acheck_copilot_seat_maps_models_failures(
per_user_credential,
status: int,
error_type: type[Exception],
) -> None:
def respond(request: httpx.Request) -> httpx.Response:
if str(request.url) == COPILOT_USER_URL:
return httpx.Response(200, json=_user_payload(), request=request)
if str(request.url) == "https://api.githubcopilot.com/models":
return httpx.Response(status, request=request)
pytest.fail(f"unexpected request to {request.url}")
with github_http(respond):
with pytest.raises(error_type):
await acheck_copilot_seat("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
@pytest.mark.asyncio
async def test_acheck_copilot_seat_maps_models_request_errors(per_user_credential) -> None:
def respond(request: httpx.Request) -> httpx.Response:
if str(request.url) == COPILOT_USER_URL:
return httpx.Response(200, json=_user_payload(), request=request)
if str(request.url) == "https://api.githubcopilot.com/models":
raise httpx.ConnectError("connection failed", request=request)
pytest.fail(f"unexpected request to {request.url}")
with github_http(respond):
with pytest.raises(APIConnectionError, match="models check request failed"):
await acheck_copilot_seat("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
@pytest.mark.asyncio
async def test_acheck_copilot_seat_passes_when_models_are_usable(per_user_credential):
def respond(request: httpx.Request) -> httpx.Response:
if str(request.url) == COPILOT_USER_URL:
return httpx.Response(200, json=_user_payload(), request=request)
if str(request.url) == "https://api.githubcopilot.com/models":
return httpx.Response(200, json={"data": []}, request=request)
pytest.fail(f"unexpected request to {request.url}")
with github_http(respond) as requests:
session = await acheck_copilot_seat("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
assert session.token == "gh-token"
assert session.api_base == "https://api.githubcopilot.com"
assert [str(request.url) for request in requests] == [
COPILOT_USER_URL,
"https://api.githubcopilot.com/models",
]
@pytest.mark.asyncio
async def test_aresolve_single_flight_one_http_call():
def respond(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, json=_user_payload(), request=request)
with github_http(respond) as requests:
s1, s2 = await asyncio.gather(
aresolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS),
aresolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS),
)
assert len(requests) == 1
assert s1.token == s2.token == "gh-token"
@pytest.mark.asyncio
async def test_aresolve_401_raises_and_evicts():
with github_http(_user_responder(status=401)):
with pytest.raises(CallerCredentialAuthenticationError):
await aexchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
await aresolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
assert (
_SESSION_CACHE.get_cache(_session_cache_key("user-a", "copilot-cred", GITHUB_COPILOT_DEFAULT_HOST)) is None
)
ACCOUNT_URL = "https://api.github.com/copilot_internal/user"
EXCHANGE_URL: Final = "https://api.github.com/copilot_internal/v2/token"
_CUSTOM_CLIENT_ID: Final = "Ov23liCustomOrgApp"
_ACCOUNT_API_BASE: Final = "https://api.business.githubcopilot.com"
@ -327,17 +410,20 @@ def _account_responder(payload=None, status=200, model_payload=None):
return respond
@pytest.fixture
def custom_client_id(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", _CUSTOM_CLIENT_ID)
@pytest.fixture(params=[None, _CUSTOM_CLIENT_ID], ids=["default-client-id", "custom-client-id"])
def copilot_client_id(request, monkeypatch: pytest.MonkeyPatch):
if request.param is None:
monkeypatch.delenv("GITHUB_COPILOT_CLIENT_ID", raising=False)
return
monkeypatch.setenv("GITHUB_COPILOT_CLIENT_ID", request.param)
@pytest.mark.asyncio
async def test_custom_client_id_uses_the_github_token_without_the_copilot_token_swap(custom_client_id):
async def test_user_lookup_uses_the_github_token_for_any_client_id(copilot_client_id):
endpoints: Final = github_host_endpoints("github.com")
with github_http(_account_responder()) as requests:
first = await aexchange_github_token("user-a", "gho_org_app", "copilot-cred", endpoints)
second = await aexchange_github_token("user-a", "gho_org_app", "copilot-cred", endpoints)
first = await aresolve_copilot_user_session("user-a", "gho_org_app", "copilot-cred", endpoints)
second = await aresolve_copilot_user_session("user-a", "gho_org_app", "copilot-cred", endpoints)
assert [str(r.url) for r in requests] == [ACCOUNT_URL]
assert requests[0].headers["Authorization"] == "token gho_org_app"
@ -353,22 +439,24 @@ async def test_custom_client_id_uses_the_github_token_without_the_copilot_token_
_account_payload(access_type_sku="no_access"),
],
)
def test_custom_client_id_without_a_copilot_seat_raises_reconnect(custom_client_id, payload):
def test_user_lookup_without_a_copilot_seat_raises_reconnect(copilot_client_id, payload):
with github_http(_account_responder(payload)):
with pytest.raises(CallerCredentialAuthenticationError):
exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com"))
resolve_copilot_user_session("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com"))
@pytest.mark.parametrize("status", [401, 403, 404])
def test_custom_client_id_rejected_account_lookup_raises_reconnect(custom_client_id, status):
def test_rejected_user_lookup_raises_reconnect(copilot_client_id, status):
with github_http(_account_responder(status=status)):
with pytest.raises(CallerCredentialAuthenticationError):
exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com"))
resolve_copilot_user_session("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com"))
def test_custom_client_id_account_api_base_must_stay_on_copilot_hosts(custom_client_id):
def test_user_lookup_api_base_must_stay_on_copilot_hosts(copilot_client_id):
with github_http(_account_responder(_account_payload(api_base="https://evil.example"))):
session = exchange_github_token("user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com"))
session = resolve_copilot_user_session(
"user-a", "gho_org_app", "copilot-cred", github_host_endpoints("github.com")
)
assert session.api_base == "https://api.githubcopilot.com"
@ -392,16 +480,16 @@ def _kwargs_with_connection(
@pytest.mark.asyncio
async def test_aattach_injects_session_into_kwargs(per_user_credential):
with github_http(_exchange_responder()):
with github_http(_user_responder()):
kwargs = _kwargs_with_connection()
await aattach_github_copilot_user_session(kwargs)
session = github_copilot_user_session_from(kwargs)
assert session is not None
assert session.token == "copilot-token"
assert session.token == "gho_user_token"
def test_attach_sync_injects_session(per_user_credential):
with github_http(_exchange_responder()):
with github_http(_user_responder()):
kwargs = _kwargs_with_connection()
attach_github_copilot_user_session(kwargs)
assert isinstance(kwargs[GITHUB_COPILOT_USER_SESSION_KWARG_KEY], GithubCopilotUserSession)
@ -424,7 +512,7 @@ def test_attach_rejects_connection_from_a_different_github_host(
):
with patch.object(litellm, "credential_list", [_credential(github_host=credential_host)]):
kwargs: Final = _kwargs_with_connection(github_host=connection_host)
with github_http(lambda request: pytest.fail(f"unexpected token exchange at {request.url}")) as requests:
with github_http(lambda request: pytest.fail(f"unexpected user lookup at {request.url}")) as requests:
with pytest.raises(CallerCredentialAuthenticationError) as exc:
attach_github_copilot_user_session(kwargs)
assert "Connect GitHub Copilot for credential 'copilot-cred'" in str(exc.value)
@ -432,11 +520,11 @@ def test_attach_rejects_connection_from_a_different_github_host(
def test_legacy_connection_without_host_uses_github_com(per_user_credential):
with github_http(_exchange_responder()) as requests:
with github_http(_user_responder()) as requests:
kwargs: Final = _kwargs_with_connection(github_host=None)
attach_github_copilot_user_session(kwargs)
assert len(requests) == 1
assert str(requests[0].url) == EXCHANGE_URL
assert str(requests[0].url) == COPILOT_USER_URL
def test_attach_shared_mode_does_nothing_without_connection():
@ -447,10 +535,10 @@ def test_attach_shared_mode_does_nothing_without_connection():
def test_evict_drops_cached_session():
with github_http(_exchange_responder()) as requests:
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
with github_http(_user_responder()) as requests:
resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
evict_copilot_user_session("user-a", "copilot-cred", _GITHUB_ENDPOINTS.host)
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
assert len(requests) == 2
@ -559,7 +647,7 @@ def test_two_per_user_calls_with_identical_bodies_both_reach_upstream( # test-q
monkeypatch.setattr(litellm, "cache", Cache())
with (
github_http(_exchange_responder()),
github_http(_user_responder()),
patch(
"litellm.main.openai_chat_completions.completion",
return_value=_chat_completion_response(),
@ -584,7 +672,7 @@ async def test_two_per_user_async_calls_with_identical_bodies_both_reach_upstrea
monkeypatch.setattr(litellm, "cache", Cache())
with (
github_http(_exchange_responder()),
github_http(_user_responder()),
patch(
"litellm.main.openai_chat_completions.acompletion",
new=AsyncMock(return_value=_chat_completion_response()),
@ -605,27 +693,19 @@ def test_shared_mode_second_identical_call_hits_response_cache( # test-quality-
):
"""Control: shared device-login github_copilot keeps caching exactly as before, so a second
identical call does not touch the upstream."""
import json
from litellm.caching import Cache
# a mounted shared login: api-key.json with a far-future token, no device flow needed
(tmp_path / "api-key.json").write_text(
json.dumps(
{
"token": "shared-copilot-token",
"expires_at": 4102444800,
"endpoints": {"api": "https://api.githubcopilot.com"},
}
)
)
(tmp_path / "access-token").write_text("shared-copilot-token")
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
monkeypatch.setattr(litellm, "cache", Cache())
with patch(
"litellm.main.openai_chat_completions.completion",
return_value=_chat_completion_response(),
) as mock_completion:
with (
github_http(_user_responder()),
patch(
"litellm.main.openai_chat_completions.completion",
return_value=_chat_completion_response(),
) as mock_completion,
):
call = {
"model": "github_copilot/gpt-4o",
"messages": [{"role": "user", "content": "hi"}],
@ -679,8 +759,8 @@ def _anthropic_messages_payload():
def _copilot_surface_responder(model_payload):
def respond(request: httpx.Request) -> httpx.Response:
if str(request.url) == EXCHANGE_URL:
return httpx.Response(200, json=_exchange_payload(), request=request)
if str(request.url) == COPILOT_USER_URL:
return httpx.Response(200, json=_user_payload(), request=request)
if request.url.host.endswith("githubcopilot.com"):
return httpx.Response(200, json=model_payload, request=request)
pytest.fail(f"unexpected request to {request.url}")
@ -696,19 +776,19 @@ def _no_authenticator():
)
def _assert_exchange_then_model_call(requests, github_token, model_host="https://api.githubcopilot.com"):
exchange = [r for r in requests if str(r.url) == EXCHANGE_URL]
def _assert_user_lookup_then_model_call(requests, github_token, model_host="https://api.githubcopilot.com"):
lookup = [r for r in requests if str(r.url) == COPILOT_USER_URL]
upstream = [r for r in requests if r.url.host.endswith("githubcopilot.com")]
assert exchange, "expected the GitHub token exchange request"
assert lookup, "expected the GitHub user lookup request"
assert upstream, "expected the model request to the Copilot host"
assert exchange[0].headers["Authorization"] == f"token {github_token}"
assert upstream[0].headers["Authorization"].startswith("Bearer ")
assert lookup[0].headers["Authorization"] == f"token {github_token}"
assert upstream[0].headers["Authorization"] == f"Bearer {github_token}"
assert upstream[0].url.host == httpx.URL(model_host).host
def test_e2e_completion_per_user_credential(per_user_credential):
with (
github_http(_exchange_responder()) as requests,
github_http(_user_responder()) as requests,
_no_authenticator(),
patch(
"litellm.main.openai_chat_completions.completion",
@ -721,13 +801,13 @@ def test_e2e_completion_per_user_credential(per_user_credential):
litellm_credential_name="copilot-cred",
secret_fields=_secret_fields("user-a", "gho_a"),
)
assert len(requests) == 1 and str(requests[0].url) == EXCHANGE_URL
assert len(requests) == 1 and str(requests[0].url) == COPILOT_USER_URL
assert requests[0].headers["Authorization"] == "token gho_a"
call_kwargs = mock_completion.call_args.kwargs
assert call_kwargs["api_key"] == "copilot-token"
assert call_kwargs["api_key"] == "gho_a"
assert call_kwargs["api_base"] == "https://api.githubcopilot.com"
extra_headers = call_kwargs["optional_params"]["extra_headers"]
assert extra_headers["Authorization"] == "Bearer copilot-token"
assert extra_headers["Authorization"] == "Bearer gho_a"
assert response.usage.total_tokens > 0
@ -735,7 +815,7 @@ def test_e2e_completion_per_user_session_token_wins_over_caller_authorization(pe
"""The session token reaches the wire even when the caller passes their own
Authorization in extra_headers (chat dispatches through _complete_custom_openai)."""
with (
github_http(_exchange_responder()),
github_http(_user_responder()),
_no_authenticator(),
patch(
"litellm.main.openai_chat_completions.completion",
@ -750,17 +830,17 @@ def test_e2e_completion_per_user_session_token_wins_over_caller_authorization(pe
extra_headers={"Authorization": "Bearer caller-token", "x-custom": "keep"},
)
extra_headers = mock_completion.call_args.kwargs["optional_params"]["extra_headers"]
assert extra_headers["Authorization"] == "Bearer copilot-token"
assert extra_headers["Authorization"] == "Bearer gho_a"
assert extra_headers["x-custom"] == "keep"
_EVIL_AUTH_HEADERS = {"authorization": "Bearer caller-evil", "x-keep": "1"}
def _assert_session_token_on_the_wire(requests):
def _assert_session_token_on_the_wire(requests, github_token="gho_a"):
upstream = [r for r in requests if r.url.host.endswith("githubcopilot.com")]
assert upstream, "expected the model request to the Copilot host"
assert upstream[0].headers["authorization"] == "Bearer copilot-token"
assert upstream[0].headers["authorization"] == f"Bearer {github_token}"
def test_e2e_completion_per_user_caller_authorization_cannot_override_session(per_user_credential):
@ -769,7 +849,7 @@ def test_e2e_completion_per_user_caller_authorization_cannot_override_session(pe
handed to openai_chat_completions.completion."""
caller_headers: dict[str, str] = dict(_EVIL_AUTH_HEADERS)
with (
github_http(_exchange_responder()),
github_http(_user_responder()),
_no_authenticator(),
patch(
"litellm.main.openai_chat_completions.completion",
@ -784,7 +864,7 @@ def test_e2e_completion_per_user_caller_authorization_cannot_override_session(pe
extra_headers=caller_headers,
)
extra_headers = mock_completion.call_args.kwargs["optional_params"]["extra_headers"]
assert extra_headers["Authorization"] == "Bearer copilot-token"
assert extra_headers["Authorization"] == "Bearer gho_a"
assert "authorization" not in extra_headers
assert extra_headers["x-keep"] == "1"
assert caller_headers == dict(_EVIL_AUTH_HEADERS), "the caller's dict must not be mutated"
@ -805,9 +885,7 @@ def test_e2e_embedding_per_user_caller_authorization_cannot_override_session(per
_assert_session_token_on_the_wire(requests)
def test_e2e_custom_client_id_sends_the_github_token_to_the_accounts_copilot_host(
per_user_credential, custom_client_id
):
def test_e2e_per_user_sends_the_github_token_to_the_accounts_copilot_host(per_user_credential, copilot_client_id):
with (
github_http(_account_responder(model_payload=_embedding_response_payload())) as requests,
_no_authenticator(),
@ -869,7 +947,7 @@ def test_e2e_embedding_per_user_credential(per_user_credential):
litellm_credential_name="copilot-cred",
secret_fields=_secret_fields("user-a", "gho_a"),
)
_assert_exchange_then_model_call(requests, "gho_a")
_assert_user_lookup_then_model_call(requests, "gho_a")
assert requests[-1].url.path == "/embeddings"
assert response.usage.prompt_tokens == 4
@ -885,7 +963,7 @@ def test_e2e_responses_per_user_credential(per_user_credential):
litellm_credential_name="copilot-cred",
secret_fields=_secret_fields("user-a", "gho_a"),
)
_assert_exchange_then_model_call(requests, "gho_a")
_assert_user_lookup_then_model_call(requests, "gho_a")
assert requests[-1].url.path == "/responses"
assert response.usage.total_tokens == 6
@ -903,7 +981,7 @@ async def test_e2e_anthropic_messages_per_user_credential(per_user_credential):
litellm_credential_name="copilot-cred",
secret_fields=_secret_fields("user-a", "gho_a"),
)
_assert_exchange_then_model_call(requests, "gho_a")
_assert_user_lookup_then_model_call(requests, "gho_a")
assert requests[-1].url.path == "/v1/messages"
assert response["usage"]["input_tokens"] == 4
@ -954,26 +1032,19 @@ async def test_e2e_anthropic_messages_not_connected_raises_connect_401(per_user_
def _shared_login(tmp_path, monkeypatch):
import json
(tmp_path / "api-key.json").write_text(
json.dumps(
{
"token": "shared-copilot-token",
"expires_at": 4102444800,
"endpoints": {"api": "https://api.githubcopilot.com"},
}
)
)
(tmp_path / "access-token").write_text("shared-copilot-token")
monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path))
def test_e2e_completion_shared_mode_still_uses_authenticator(tmp_path, monkeypatch):
_shared_login(tmp_path, monkeypatch)
with patch(
"litellm.main.openai_chat_completions.completion",
return_value=_chat_completion_response(),
) as mock_completion:
with (
github_http(_user_responder()),
patch(
"litellm.main.openai_chat_completions.completion",
return_value=_chat_completion_response(),
) as mock_completion,
):
response = litellm.completion(
model="github_copilot/gpt-4o",
messages=[{"role": "user", "content": "hi"}],
@ -1007,45 +1078,37 @@ def test_e2e_responses_shared_mode_still_uses_authenticator(tmp_path, monkeypatc
@pytest.mark.asyncio
async def test_e2e_anthropic_messages_shared_mode_still_uses_authenticator(tmp_path, monkeypatch):
# The shared Authenticator's access-token refresh refuses to run inside an event
# loop, so on this async surface the mounted-login path can't reach it; patch the
# shared-login edge to emulate a mounted api-key.json.
async def test_e2e_anthropic_messages_shared_mode_still_uses_authenticator():
from litellm.llms.github_copilot.authenticator import Authenticator
with (
github_http(_copilot_surface_responder(_anthropic_messages_payload())) as requests,
patch.object( # test-quality-ok: patching the shared-login edge is the fixture under test
Authenticator, "get_api_key", return_value="shared-copilot-token"
) as mock_api_key,
patch.object( # test-quality-ok: same edge; api_key read is mocked so base must be too
Authenticator, "get_api_base", return_value="https://api.githubcopilot.com"
),
patch.object(Authenticator, "get_api_key", return_value="shared-copilot-token"),
patch.object(Authenticator, "get_api_base", return_value="https://api.githubcopilot.com"),
):
response = await litellm.anthropic.messages.acreate(
model="github_copilot/claude-haiku-4.5",
max_tokens=16,
messages=[{"role": "user", "content": "hi"}],
)
mock_api_key.assert_called()
upstream = [r for r in requests if r.url.host.endswith("githubcopilot.com")]
assert upstream and upstream[0].headers["Authorization"] == "Bearer shared-copilot-token"
assert response["usage"]["input_tokens"] == 4
def test_sync_exchange_single_flight_one_http_call():
def test_sync_resolve_single_flight_one_http_call():
import threading
def slow_respond(request: httpx.Request) -> httpx.Response:
time.sleep(0.05)
return httpx.Response(200, json=_exchange_payload(), request=request)
return httpx.Response(200, json=_user_payload(), request=request)
sessions: list[GithubCopilotUserSession] = []
with github_http(slow_respond) as requests:
threads = [
threading.Thread(
target=lambda: sessions.append(
exchange_github_token("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
resolve_copilot_user_session("user-a", "gh-token", "copilot-cred", _GITHUB_ENDPOINTS)
)
)
for _ in range(6)
@ -1055,13 +1118,13 @@ def test_sync_exchange_single_flight_one_http_call():
for t in threads:
t.join()
assert len(requests) == 1
assert len(sessions) == 6 and all(s.token == "copilot-token" for s in sessions)
assert len(sessions) == 6 and all(s.token == "gh-token" for s in sessions)
def test_two_users_same_github_token_do_not_share_session():
with github_http(_exchange_responder()):
s1: Final = exchange_github_token("user-a", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS)
s2: Final = exchange_github_token("user-b", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS)
with github_http(_user_responder()):
s1: Final = resolve_copilot_user_session("user-a", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS)
s2: Final = resolve_copilot_user_session("user-b", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS)
user_a_key: Final = _session_cache_key("user-a", "copilot-cred", _GITHUB_ENDPOINTS.host)
user_b_key: Final = _session_cache_key("user-b", "copilot-cred", _GITHUB_ENDPOINTS.host)
other_credential_key: Final = _session_cache_key("user-a", "another-credential", _GITHUB_ENDPOINTS.host)
@ -1069,7 +1132,7 @@ def test_two_users_same_github_token_do_not_share_session():
assert user_a_key != user_b_key
assert user_a_key != other_credential_key
assert user_a_key != other_host_key
assert s1.token == s2.token == "copilot-token"
assert s1.token == s2.token == "gh-same"
def test_same_user_and_token_do_not_share_sessions_across_hosts():
@ -1078,11 +1141,11 @@ def test_same_user_and_token_do_not_share_sessions_across_hosts():
ghe_endpoints: Final = github_host_endpoints("octo.ghe.com")
def respond(request: httpx.Request) -> httpx.Response:
return httpx.Response(200, json=_exchange_payload(), request=request)
return httpx.Response(200, json=_user_payload(), request=request)
with github_http(respond) as requests:
github_session: Final = exchange_github_token("user-a", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS)
ghe_session: Final = exchange_github_token("user-a", "gh-same", "copilot-cred", ghe_endpoints)
github_session: Final = resolve_copilot_user_session("user-a", "gh-same", "copilot-cred", _GITHUB_ENDPOINTS)
ghe_session: Final = resolve_copilot_user_session("user-a", "gh-same", "copilot-cred", ghe_endpoints)
assert len(requests) == 2
assert github_session.api_base == "https://api.githubcopilot.com"
assert ghe_session.api_base == "https://copilot-api.octo.ghe.com"
@ -1141,33 +1204,37 @@ def test_provider_info_shared_github_copilot_still_uses_authenticator():
@pytest.mark.asyncio
async def test_aexchange_cancelled_waiter_does_not_cancel_the_shared_exchange(per_user_credential):
"""Two waiters share one in-flight exchange; cancelling one must not cancel the
async def test_aresolve_cancelled_waiter_does_not_cancel_the_shared_session(per_user_credential):
"""Two waiters share one in-flight resolution; cancelling one must not cancel the
underlying future the other is waiting on."""
gate = asyncio.Event()
class _GatedClient:
async def get(self, url, headers=None):
await gate.wait()
return httpx.Response(200, json=_exchange_payload(), request=httpx.Request("GET", url))
return httpx.Response(200, json=_user_payload(), request=httpx.Request("GET", url))
with patch( # test-quality-ok: the gate holds the HTTP edge open so two waiters overlap
"litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=_GatedClient()
):
first = asyncio.ensure_future(aexchange_github_token("user-a", "gh-shared", "copilot-cred", _GITHUB_ENDPOINTS))
first = asyncio.ensure_future(
aresolve_copilot_user_session("user-a", "gh-shared", "copilot-cred", _GITHUB_ENDPOINTS)
)
await asyncio.sleep(0)
second = asyncio.ensure_future(aexchange_github_token("user-a", "gh-shared", "copilot-cred", _GITHUB_ENDPOINTS))
second = asyncio.ensure_future(
aresolve_copilot_user_session("user-a", "gh-shared", "copilot-cred", _GITHUB_ENDPOINTS)
)
await asyncio.sleep(0)
second.cancel()
gate.set()
session = await first
assert session.token == "copilot-token" and session.api_base == "https://api.githubcopilot.com"
assert session.token == "gh-shared" and session.api_base == "https://api.githubcopilot.com"
with pytest.raises(asyncio.CancelledError):
await second
@pytest.mark.asyncio
async def test_aexchange_does_not_share_an_in_flight_session_after_token_changes(per_user_credential):
async def test_aresolve_does_not_share_an_in_flight_session_after_token_changes(per_user_credential):
gate = asyncio.Event()
auth_headers: asyncio.Queue[str] = asyncio.Queue()
@ -1175,26 +1242,29 @@ async def test_aexchange_does_not_share_an_in_flight_session_after_token_changes
async def get(self, url, headers=None):
await auth_headers.put(headers["authorization"])
await gate.wait()
token = "copilot-token-a" if headers["authorization"] == "token gh-a" else "copilot-token-b"
return httpx.Response(
200,
json=_exchange_payload(token=token),
json=_user_payload(),
request=httpx.Request("GET", url),
)
with patch("litellm.llms.custom_httpx.http_handler.get_async_httpx_client", return_value=_GatedClient()):
first = asyncio.ensure_future(aexchange_github_token("user-a", "gh-a", "copilot-cred", _GITHUB_ENDPOINTS))
first = asyncio.ensure_future(
aresolve_copilot_user_session("user-a", "gh-a", "copilot-cred", _GITHUB_ENDPOINTS)
)
assert await asyncio.wait_for(auth_headers.get(), timeout=1) == "token gh-a"
second = asyncio.ensure_future(aexchange_github_token("user-a", "gh-b", "copilot-cred", _GITHUB_ENDPOINTS))
second = asyncio.ensure_future(
aresolve_copilot_user_session("user-a", "gh-b", "copilot-cred", _GITHUB_ENDPOINTS)
)
assert await asyncio.wait_for(auth_headers.get(), timeout=1) == "token gh-b"
gate.set()
sessions = await asyncio.gather(first, second)
assert [session.token for session in sessions] == ["copilot-token-a", "copilot-token-b"]
assert [session.token for session in sessions] == ["gh-a", "gh-b"]
@pytest.mark.asyncio
async def test_aexchange_failure_with_a_single_caller_emits_no_unretrieved_warning(per_user_credential):
async def test_aresolve_failure_with_a_single_caller_emits_no_unretrieved_warning(per_user_credential):
"""The leader removes the failed future before anyone awaits it; the stored
exception must be consumed so the loop does not log 'never retrieved'."""
import litellm.llms.github_copilot.per_user_auth as pua
@ -1205,10 +1275,10 @@ async def test_aexchange_failure_with_a_single_caller_emits_no_unretrieved_warni
loop.set_exception_handler(lambda _loop, context: captured.append((_loop, context)))
try:
with (
github_http(_exchange_responder(status=401)),
github_http(_user_responder(status=401)),
pytest.raises(CallerCredentialAuthenticationError),
):
await aexchange_github_token("user-a", "gh-fail", "copilot-cred", _GITHUB_ENDPOINTS)
await aresolve_copilot_user_session("user-a", "gh-fail", "copilot-cred", _GITHUB_ENDPOINTS)
_IN_FLIGHT_KEYS = list(pua._IN_FLIGHT)
import gc
@ -1240,7 +1310,7 @@ async def test_per_user_streaming_call_writes_nothing_to_the_response_cache(per_
"model": "github_copilot/gpt-4o",
"messages": [{"role": "user", "content": "hi"}],
GITHUB_COPILOT_USER_SESSION_KWARG_KEY: GithubCopilotUserSession(
token="copilot-token", api_base="https://api.githubcopilot.com"
token="gh-token", api_base="https://api.githubcopilot.com"
),
}
await handler.async_get_cache(

193
tests/unit/proxy/anthropic_endpoints/test_claude_code_marketplace.py
View file

@ -80,7 +80,7 @@ def _make_mock_prisma():
return record
async def _update(where, data):
record = store[where["name"]]
record = store[where["name"]] if "name" in where else next(r for r in store.values() if r.id == where["id"])
for k, v in data.items():
setattr(record, k, v)
return record
@ -91,9 +91,20 @@ def _make_mock_prisma():
mock_table.update = AsyncMock(side_effect=_update)
mock_client.db.litellm_claudecodeplugintable = mock_table
mock_client.writer_db = mock_client.db
mock_client.db.query_raw = AsyncMock(return_value=[])
mock_client.db.tx.return_value.__aenter__.return_value = mock_client.db
mock_client.db.litellm_claudecodepluginversion.count = AsyncMock(return_value=0)
mock_client.db.litellm_claudecodepluginversion.create = AsyncMock()
mock_client.db.litellm_claudecodepluginupload.find_unique = AsyncMock(return_value=None)
mock_client.db.litellm_claudecodepluginupload.delete_many = AsyncMock(return_value=0)
mock_client.db.litellm_claudecodepluginupload.upsert = AsyncMock()
return mock_client
def _uploads() -> PluginUploadRepository:
return PluginUploadRepository(litellm.proxy.proxy_server.prisma_client)
_USER = UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key=MASTER_KEY,
@ -131,7 +142,7 @@ async def test_register_plugin_git_subdir_success():
"""git-subdir with both url and path fields registers successfully."""
request = RegisterPluginRequest(name="my-monorepo-plugin", source=_GIT_SUBDIR_SOURCE)
response = await register_plugin(request=request, user_api_key_dict=_USER)
response = await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER)
assert response.status == "success"
assert response.action == "created"
@ -149,7 +160,7 @@ async def _read_stored_manifest(name: str) -> dict:
async def test_register_plugin_duplicate_name_conflicts():
"""A second POST with an existing name returns 409 and leaves the stored plugin untouched."""
name = "my-monorepo-plugin"
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"),
user_api_key_dict=_USER,
)
@ -167,7 +178,7 @@ async def test_register_plugin_duplicate_name_conflicts():
version="2.0.0",
)
with pytest.raises(HTTPException) as exc_info:
await register_plugin(request=conflicting, user_api_key_dict=_USER)
await register_plugin(uploads=_uploads(), request=conflicting, user_api_key_dict=_USER)
assert exc_info.value.status_code == 409
assert "already exists" in exc_info.value.detail["error"]
@ -182,13 +193,13 @@ async def test_register_plugin_duplicate_name_conflicts():
async def test_update_plugin_replaces_existing_source():
"""PUT updates an existing plugin: action=updated and the stored source is replaced."""
name = "my-monorepo-plugin"
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"),
user_api_key_dict=_USER,
)
new_source = {"source": "github", "repo": "org/replacement"}
response = await update_plugin(
response = await update_plugin(uploads=_uploads(),
plugin_name=name,
request=UpdatePluginRequest(source=new_source, version="2.0.0", description="updated"),
user_api_key_dict=_USER,
@ -208,13 +219,13 @@ async def test_update_plugin_replaces_existing_source():
async def test_metadata_edits_preserve_uploaded_source_and_reject_a_changed_http_destination() -> None:
table: Final = litellm.proxy.proxy_server.prisma_client.writer_db.litellm_claudecodeplugintable
await table.create(data={"name": _UPLOAD_NAME, "manifest_json": json.dumps({"source": _UPLOAD_SOURCE})})
request: Final = UpdatePluginRequest(source=_UPLOAD_SOURCE, description="Updated description")
response: Final = await update_plugin(plugin_name=_UPLOAD_NAME, request=request, user_api_key_dict=_USER)
expected: Final = {"name": _UPLOAD_NAME, "source": _UPLOAD_SOURCE, "description": "Updated description"}
request: Final = UpdatePluginRequest(source=_UPLOAD_SOURCE, description="Updated description", version="1.0.0")
response: Final = await update_plugin(uploads=_uploads(), plugin_name=_UPLOAD_NAME, request=request, user_api_key_dict=_USER)
expected: Final = {"name": _UPLOAD_NAME, "source": _UPLOAD_SOURCE, "description": "Updated description", "version": "1.0.0"}
assert response.plugin.description == "Updated description"
assert await _read_stored_manifest(_UPLOAD_NAME) == expected
with pytest.raises(HTTPException) as error:
await update_plugin(
await update_plugin(uploads=_uploads(),
plugin_name=_UPLOAD_NAME, user_api_key_dict=_USER,
request=UpdatePluginRequest(source={**_UPLOAD_SOURCE, "url": "http://other.example/skill.zip"}),
)
@ -227,7 +238,7 @@ async def test_metadata_edits_preserve_uploaded_source_and_reject_a_changed_http
async def test_update_plugin_not_found():
"""PUT on a name that does not exist raises HTTP 404."""
with pytest.raises(HTTPException) as exc_info:
await update_plugin(
await update_plugin(uploads=_uploads(),
plugin_name="does-not-exist",
request=UpdatePluginRequest(source=_GIT_SUBDIR_SOURCE),
user_api_key_dict=_USER,
@ -246,7 +257,7 @@ async def test_register_plugin_create_race_maps_unique_violation_to_409():
table.create = AsyncMock(side_effect=UniqueViolationError({}, message="duplicate name"))
with pytest.raises(HTTPException) as exc_info:
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name="racy-plugin", source=_GIT_SUBDIR_SOURCE),
user_api_key_dict=_USER,
)
@ -262,7 +273,7 @@ async def test_update_plugin_db_error_maps_to_structured_500():
from prisma.errors import PrismaError
name = "my-monorepo-plugin"
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"),
user_api_key_dict=_USER,
)
@ -271,7 +282,7 @@ async def test_update_plugin_db_error_maps_to_structured_500():
table.update = AsyncMock(side_effect=PrismaError("connection lost"))
with pytest.raises(HTTPException) as exc_info:
await update_plugin(
await update_plugin(uploads=_uploads(),
plugin_name=name,
request=UpdatePluginRequest(source={"source": "github", "repo": "org/replacement"}),
user_api_key_dict=_USER,
@ -286,16 +297,17 @@ async def test_update_plugin_deleted_mid_update_returns_404():
"""A concurrent delete between the find_unique pre-check and the update makes prisma's
update return None; that must surface the same 404 as a plain miss, not an AttributeError."""
name = "my-monorepo-plugin"
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"),
user_api_key_dict=_USER,
)
table = litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable
table.update = AsyncMock(return_value=None)
existing: Final = await table.find_unique(where={"name": name})
table.find_unique = AsyncMock(side_effect=[existing, None])
with pytest.raises(HTTPException) as exc_info:
await update_plugin(
await update_plugin(uploads=_uploads(),
plugin_name=name,
request=UpdatePluginRequest(source={"source": "github", "repo": "org/replacement"}),
user_api_key_dict=_USER,
@ -307,7 +319,7 @@ async def test_update_plugin_deleted_mid_update_returns_404():
@pytest.mark.asyncio
async def test_get_marketplace_skips_plugin_with_null_manifest():
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name="good-plugin", source=_GIT_SUBDIR_SOURCE, version="1.0.0"),
user_api_key_dict=_USER,
)
@ -333,7 +345,7 @@ def _granted_user(skills: list[str]) -> UserAPIKeyAuth:
async def _register_public_and_private_plugins() -> None:
for name, enabled in (("public-skill", True), ("private-skill", False)):
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"),
user_api_key_dict=_USER,
)
@ -401,7 +413,7 @@ async def test_get_marketplace_key_query_param_adds_granted_disabled_plugins(mon
@pytest.mark.asyncio
@pytest.mark.parametrize("preference", ["available", "auto_install", "required"])
async def test_get_marketplace_emits_installation_preference(preference):
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(
name="s3-skill",
source=_ARCHIVE_SOURCE,
@ -434,7 +446,7 @@ def test_plugin_request_rejects_unknown_installation_preference(build_request):
@pytest.mark.asyncio
async def test_get_marketplace_omits_installation_preference_when_unset():
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name="manual-install-plugin", source=_GIT_SUBDIR_SOURCE),
user_api_key_dict=_USER,
)
@ -447,12 +459,12 @@ async def test_get_marketplace_omits_installation_preference_when_unset():
@pytest.mark.asyncio
async def test_update_plugin_propagates_installation_preference_to_get_and_list():
name = "updated-install-plugin"
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name=name, source=_ARCHIVE_SOURCE),
user_api_key_dict=_USER,
)
response = await update_plugin(
response = await update_plugin(uploads=_uploads(),
plugin_name=name,
request=UpdatePluginRequest(
source=_ARCHIVE_SOURCE,
@ -472,12 +484,12 @@ async def test_update_plugin_propagates_installation_preference_to_get_and_list(
@pytest.mark.asyncio
async def test_update_plugin_without_installation_preference_clears_it():
name = "cleared-install-plugin"
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name=name, source=_ARCHIVE_SOURCE, installation_preference="auto_install"),
user_api_key_dict=_USER,
)
await update_plugin(
await update_plugin(uploads=_uploads(),
plugin_name=name,
request=UpdatePluginRequest(source=_ARCHIVE_SOURCE),
user_api_key_dict=_USER,
@ -518,7 +530,7 @@ async def test_register_accepts_the_camel_case_installation_preference_key():
{"name": "s3-skill", "source": _ARCHIVE_SOURCE, "installationPreference": "auto_install"}
)
response = await register_plugin(request=request, user_api_key_dict=_USER)
response = await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER)
marketplace = json.loads((await get_marketplace(request=MagicMock())).body)
assert response.plugin.installation_preference == "auto_install"
@ -536,13 +548,13 @@ _UNPINNED_ARCHIVE_SOURCE = {"source": "archive", "url": "https://artifacts.examp
async def _assert_preference_refused(source, preference):
with pytest.raises(HTTPException) as register_error:
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name="no-op-plugin", source=source, installation_preference=preference),
user_api_key_dict=_USER,
)
await register_plugin(request=RegisterPluginRequest(name="no-op-plugin", source=source), user_api_key_dict=_USER)
await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name="no-op-plugin", source=source), user_api_key_dict=_USER)
with pytest.raises(HTTPException) as update_error:
await update_plugin(
await update_plugin(uploads=_uploads(),
plugin_name="no-op-plugin",
request=UpdatePluginRequest(source=source, installation_preference=preference),
user_api_key_dict=_USER,
@ -568,7 +580,7 @@ async def test_self_installing_preference_needs_the_archive_sha256(preference):
@pytest.mark.asyncio
async def test_available_preference_is_accepted_on_an_unpinned_archive():
response = await register_plugin(
response = await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(
name="listed-plugin", source=_UNPINNED_ARCHIVE_SOURCE, installation_preference="available"
),
@ -626,7 +638,7 @@ async def test_register_plugin_git_subdir_missing_url():
)
with pytest.raises(HTTPException) as exc_info:
await register_plugin(request=request, user_api_key_dict=_USER)
await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER)
assert exc_info.value.status_code == 400
assert "url" in exc_info.value.detail["error"]
@ -641,7 +653,7 @@ async def test_register_plugin_git_subdir_empty_url():
)
with pytest.raises(HTTPException) as exc_info:
await register_plugin(request=request, user_api_key_dict=_USER)
await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER)
assert exc_info.value.status_code == 400
assert "url" in exc_info.value.detail["error"]
@ -656,7 +668,7 @@ async def test_register_plugin_git_subdir_missing_path():
)
with pytest.raises(HTTPException) as exc_info:
await register_plugin(request=request, user_api_key_dict=_USER)
await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER)
assert exc_info.value.status_code == 400
assert "path" in exc_info.value.detail["error"]
@ -675,7 +687,7 @@ async def test_register_plugin_git_subdir_empty_path():
)
with pytest.raises(HTTPException) as exc_info:
await register_plugin(request=request, user_api_key_dict=_USER)
await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER)
assert exc_info.value.status_code == 400
assert "path" in exc_info.value.detail["error"]
@ -703,7 +715,7 @@ async def test_register_plugin_git_subdir_path_traversal():
)
with pytest.raises(HTTPException) as exc_info:
await register_plugin(request=request, user_api_key_dict=_USER)
await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER)
assert exc_info.value.status_code == 400
assert "relative" in exc_info.value.detail["error"]
@ -718,7 +730,7 @@ async def test_register_plugin_unknown_source_type():
)
with pytest.raises(HTTPException) as exc_info:
await register_plugin(request=request, user_api_key_dict=_USER)
await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_USER)
assert exc_info.value.status_code == 400
assert "git-subdir" in exc_info.value.detail["error"]
@ -727,7 +739,7 @@ async def test_register_plugin_unknown_source_type():
@pytest.mark.asyncio
async def test_archive_source_registers_and_is_served_verbatim_in_marketplace():
response = await register_plugin(
response = await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name="s3-skill", source=_ARCHIVE_SOURCE),
user_api_key_dict=_USER,
)
@ -743,7 +755,7 @@ async def test_archive_source_registers_and_is_served_verbatim_in_marketplace():
async def test_archive_source_without_sha256_is_accepted():
source = {"source": "archive", "url": "https://artifacts.example.com/plugin.zip"}
response = await register_plugin(
response = await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name="unpinned-skill", source=source),
user_api_key_dict=_USER,
)
@ -754,12 +766,12 @@ async def test_archive_source_without_sha256_is_accepted():
@pytest.mark.asyncio
async def test_update_plugin_to_archive_source():
name = "my-monorepo-plugin"
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE),
user_api_key_dict=_USER,
)
response = await update_plugin(
response = await update_plugin(uploads=_uploads(),
plugin_name=name,
request=UpdatePluginRequest(source=_ARCHIVE_SOURCE),
user_api_key_dict=_USER,
@ -787,7 +799,7 @@ async def test_update_plugin_to_archive_source():
)
async def test_register_plugin_archive_rejects_malformed_source(source, expected_fragment):
with pytest.raises(HTTPException) as exc_info:
await register_plugin(request=RegisterPluginRequest(name="bad-plugin", source=source), user_api_key_dict=_USER)
await register_plugin(uploads=_uploads(), request=RegisterPluginRequest(name="bad-plugin", source=source), user_api_key_dict=_USER)
assert exc_info.value.status_code == 400
assert expected_fragment in exc_info.value.detail["error"]
@ -799,7 +811,7 @@ async def test_register_plugin_rejects_non_admin():
request = RegisterPluginRequest(name="attacker-plugin", source=_GIT_SUBDIR_SOURCE)
with pytest.raises(HTTPException) as exc_info:
await register_plugin(request=request, user_api_key_dict=_NON_ADMIN_USER)
await register_plugin(uploads=_uploads(), request=request, user_api_key_dict=_NON_ADMIN_USER)
assert exc_info.value.status_code == 403
@ -811,14 +823,14 @@ async def test_register_plugin_rejects_non_admin():
async def test_update_plugin_rejects_non_admin_overwrite():
"""A non-admin key cannot overwrite an existing plugin's source."""
name = "trusted-plugin"
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"),
user_api_key_dict=_USER,
)
malicious_source = {"source": "github", "repo": "attacker/malicious-repo"}
with pytest.raises(HTTPException) as exc_info:
await update_plugin(
await update_plugin(uploads=_uploads(),
plugin_name=name,
request=UpdatePluginRequest(source=malicious_source),
user_api_key_dict=_NON_ADMIN_USER,
@ -834,7 +846,7 @@ async def test_update_plugin_rejects_non_admin_overwrite():
async def test_enable_disable_delete_plugin_reject_non_admin():
"""Non-admin keys cannot enable, disable, or delete catalog entries."""
name = "trusted-plugin-2"
await register_plugin(
await register_plugin(uploads=_uploads(),
request=RegisterPluginRequest(name=name, source=_GIT_SUBDIR_SOURCE, version="1.0.0"),
user_api_key_dict=_USER,
)
@ -843,6 +855,7 @@ async def test_enable_disable_delete_plugin_reject_non_admin():
enable_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER),
disable_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER),
delete_plugin(plugin_name=name, user_api_key_dict=_NON_ADMIN_USER),
claude_code_marketplace.get_plugin_downloads(plugin_name=name, uploads=None, user_api_key_dict=_NON_ADMIN_USER),
):
with pytest.raises(HTTPException) as exc_info:
await coro
@ -869,7 +882,7 @@ async def test_register_plugin():
user_id="test-user",
)
response: Final = await register_plugin(
response: Final = await register_plugin(uploads=_uploads(),
request=request,
user_api_key_dict=user_api_key_dict,
)
@ -904,7 +917,7 @@ async def test_get_marketplace():
user_id="test-user",
)
await register_plugin(
await register_plugin(uploads=_uploads(),
request=request,
user_api_key_dict=user_api_key_dict,
)
@ -946,7 +959,7 @@ async def test_register_plugin_git_subdir():
user_id="test-user",
)
response: Final = await register_plugin(
response: Final = await register_plugin(uploads=_uploads(),
request=request,
user_api_key_dict=user_api_key_dict,
)
@ -1211,14 +1224,20 @@ _UPLOAD_SOURCE: Final = {"source": "archive", "url": f"http://test{_UPLOAD_PATH}
async def upload_storage(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[_UploadClient]:
prisma: Final = litellm.proxy.proxy_server.prisma_client
prisma.writer_db = MagicMock()
prisma.writer_db.execute_raw = AsyncMock(return_value=1)
plugins: Final = prisma.writer_db.litellm_claudecodeplugintable
archives: Final = prisma.writer_db.litellm_claudecodepluginupload
prisma.writer_db.litellm_claudecodeplugindownload.create = AsyncMock(return_value=None)
plugins.create = AsyncMock(return_value=prisma_models.LiteLLM_ClaudeCodePluginTable(
id="uploaded-id", name=_UPLOAD_NAME, enabled=True,
id="uploaded-id", name=_UPLOAD_NAME, enabled=True, anonymous_download_count=0,
))
plugins.find_unique = AsyncMock(return_value=prisma_models.LiteLLM_ClaudeCodePluginTable(
id="uploaded-id", name=_UPLOAD_NAME, enabled=True, manifest_json=json.dumps({"source": _UPLOAD_SOURCE}),
id="uploaded-id", name=_UPLOAD_NAME, enabled=True, anonymous_download_count=0, manifest_json=json.dumps({"source": _UPLOAD_SOURCE}),
))
prisma.writer_db.query_raw = AsyncMock(return_value=[])
prisma.writer_db.tx.return_value.__aenter__.return_value = prisma.writer_db
prisma.writer_db.litellm_claudecodepluginversion.create = AsyncMock()
archives.create = AsyncMock()
archives.find_first = AsyncMock(return_value=prisma_models.LiteLLM_ClaudeCodePluginUpload(
plugin_name=_UPLOAD_NAME, content=Base64.encode(b"stored archive"),
))
@ -1252,7 +1271,8 @@ async def _post_skill_upload(
async def test_skill_upload_creates_metadata_and_bytes_atomically_and_rejects_duplicates(
upload_storage: _UploadClient,
) -> None:
client, plugins, _ = upload_storage
client, plugins, archives = upload_storage
plugins.find_unique = AsyncMock(return_value=None)
buffer: Final = io.BytesIO()
with zipfile.ZipFile(buffer, "w") as archive:
archive.writestr("SKILL.md", "Read the personal instructions")
@ -1261,7 +1281,7 @@ async def test_skill_upload_creates_metadata_and_bytes_atomically_and_rejects_du
assert response.status_code == 200, response.text
plugins.create.assert_awaited_once()
stored: Final = plugins.create.await_args.kwargs["data"]
content: Final = stored["uploaded_archive"]["create"]["content"].decode()
content: Final = archives.create.await_args.kwargs["data"]["content"].decode()
digest: Final = hashlib.sha256(content).hexdigest()
source: Final = {
"source": "archive",
@ -1272,7 +1292,7 @@ async def test_skill_upload_creates_metadata_and_bytes_atomically_and_rejects_du
assert stored == {
"name": _UPLOAD_NAME, "description": "My notes", "version": "2.0.0",
"manifest_json": stored["manifest_json"], "created_by": litellm.proxy.proxy_server.litellm_proxy_admin_name,
"uploaded_archive": {"create": {"content": Base64.encode(content)}},
"files_json": "{}", "updated_at": stored["updated_at"],
}
assert response.json() == {"status": "success", "action": "created", "plugin": {
"id": "uploaded-id", "name": _UPLOAD_NAME, "version": "2.0.0", "description": "My notes",
@ -1318,7 +1338,7 @@ async def test_archive_reads_only_authorized_current_bytes(
client, plugins, archives = upload_storage
manifest: Final = json.dumps({"source": source})
plugins.find_unique.return_value = None if source is None else prisma_models.LiteLLM_ClaudeCodePluginTable(
id="uploaded-id", name=_UPLOAD_NAME, enabled=enabled, manifest_json=manifest,
id="uploaded-id", name=_UPLOAD_NAME, enabled=enabled, anonymous_download_count=0, manifest_json=manifest,
)
if not blob:
archives.find_first.return_value = None
@ -1327,9 +1347,16 @@ async def test_archive_reads_only_authorized_current_bytes(
assert (response.content == b"stored archive") is (status == 200)
plugins.find_unique.assert_awaited_once_with(where={"name": _UPLOAD_NAME})
assert archives.find_first.await_args_list == ([call(where={
"plugin_name": _UPLOAD_NAME, "plugin": {"is": {"manifest_json": manifest, "enabled": enabled}},
"plugin_name": _UPLOAD_NAME, "plugin": {"is": {
"id": "uploaded-id", "manifest_json": manifest, "enabled": enabled,
}},
})] if reads else [])
litellm.proxy.proxy_server.prisma_client.db.litellm_claudecodeplugintable.find_unique.assert_not_awaited()
events: Final = litellm.proxy.proxy_server.prisma_client.writer_db.litellm_claudecodeplugindownload
assert events.create.await_count == int(status == 200 and key is not None)
assert litellm.proxy.proxy_server.prisma_client.writer_db.execute_raw.await_count == int(
status == 200 and key is None
)
@pytest.mark.asyncio
@ -1359,3 +1386,59 @@ async def test_archive_requires_header_credentials_for_private_content(
assert response.status_code == status, response.text
assert (response.content == b"stored archive") is (status == 200)
assert archives.find_first.await_count == int(status == 200)
events: Final = litellm.proxy.proxy_server.prisma_client.writer_db.litellm_claudecodeplugindownload
assert events.create.await_count == int(status == 200)
@pytest.mark.asyncio
@pytest.mark.parametrize("key, query, exists, status", [
(_LIMITED_KEY, {}, True, 403),
(None, {}, True, 401),
(MASTER_KEY, {"cursor": "9223372036854775808"}, True, 422),
(MASTER_KEY, {"cursor": "0"}, True, 422),
(MASTER_KEY, {"limit": 101}, True, 422),
(MASTER_KEY, {}, False, 404),
])
async def test_download_history_rejects_unauthorized_invalid_and_unknown_requests(
upload_storage: _UploadClient, key: str | None, query: Mapping[str, str | int], exists: bool, status: int,
) -> None:
client, plugins, _ = upload_storage
if not exists:
plugins.find_unique.return_value = None
response: Final = await client.get(
f"/claude-code/plugins/{_UPLOAD_NAME}/downloads",
headers={"Authorization": f"Bearer {key}"} if key else {}, params=query,
)
assert response.status_code == status, response.text
events: Final = litellm.proxy.proxy_server.prisma_client.writer_db.litellm_claudecodeplugindownload
events.find_many.assert_not_called()
@pytest.mark.asyncio
async def test_published_updates_auto_bump_and_reject_stale_editors() -> None:
from datetime import datetime, timezone
source: Final = {"source": "github", "repo": "example/skill"}
created: Final = await register_plugin(
uploads=_uploads(), request=RegisterPluginRequest(name="versioned", source=source, version=None),
user_api_key_dict=_USER,
)
updated: Final = await update_plugin(
uploads=_uploads(), plugin_name="versioned", request=UpdatePluginRequest(source=source),
user_api_key_dict=_USER,
)
assert (created.plugin.version, updated.plugin.version) == ("1.0.0", "1.0.1")
before: Final = await _read_stored_manifest("versioned")
with pytest.raises(HTTPException) as stale:
await update_plugin(
uploads=_uploads(), plugin_name="versioned", user_api_key_dict=_USER,
request=UpdatePluginRequest(source=source, expected_updated_at=datetime(2000, 1, 1, tzinfo=timezone.utc)),
)
assert stale.value.status_code == 409
with pytest.raises(HTTPException) as error:
await update_plugin(
uploads=_uploads(), plugin_name="versioned", user_api_key_dict=_USER,
request=UpdatePluginRequest(source=source, version="1.0.0", description="rejected edit"),
)
assert error.value.status_code == 400
assert await _read_stored_manifest("versioned") == before

34
tests/unit/proxy/auth/test_route_checks.py
View file

@ -363,18 +363,24 @@ def test_proxy_admin_viewer_allowed_management_reads(allowed_read_route):
)
def test_virtual_key_allowed_routes_with_litellm_routes_member_name_allowed():
@pytest.mark.parametrize("group, route", (
("openai_routes", "/chat/completions"),
("skill_hub_routes", "/skill-hub/skills/release-notes/downloads"),
("management_routes", "/skill-hub/skills/release-notes/downloads"),
("anthropic_routes", "/claude-code/plugins/release-notes/downloads"),
))
def test_virtual_key_allowed_routes_with_litellm_routes_member_name_allowed(group: str, route: str) -> None:
"""Test that virtual key is allowed to call routes when allowed_routes contains LiteLLMRoutes member name"""
# Create a UserAPIKeyAuth with allowed_routes containing a LiteLLMRoutes member name
valid_token = UserAPIKeyAuth(
user_id="test_user",
allowed_routes=["openai_routes"], # This is a member name in LiteLLMRoutes enum
allowed_routes=[group],
)
# Test that a route from the openai_routes group is allowed
result = RouteChecks.is_virtual_key_allowed_to_call_route(
route="/chat/completions", # This is in LiteLLMRoutes.openai_routes.value
route=route,
valid_token=valid_token,
)
@ -5102,3 +5108,25 @@ def test_every_dashboard_role_can_send_telemetry_ui_events(
)
except HTTPException as exc:
pytest.fail(f"{user_role.value} was denied {method} {route}: {exc.detail}")
@pytest.mark.parametrize("route", [
"/skill-hub/skills/{name}/versions",
"/skill-hub/skills/{name}/versions/{version_id}",
"/skill-hub/skills/{name}/versions/{version_id}/archive",
])
def test_member_can_reach_skill_history_visibility_checks(route: str) -> None:
request: Final = MagicMock(spec=Request)
request.method = "GET"
request.query_params = {}
try:
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=LiteLLM_UserTable(user_id="reader", user_role=LitellmUserRoles.INTERNAL_USER.value),
_user_role=LitellmUserRoles.INTERNAL_USER,
route=route,
request=request,
valid_token=UserAPIKeyAuth(user_id="reader", user_role=LitellmUserRoles.INTERNAL_USER),
request_data={},
)
except HTTPException as exc:
pytest.fail(f"History visibility could not run: {exc.detail}")

6
tests/unit/proxy/common_utils/test_check_batch_cost.py
View file

@ -3324,9 +3324,13 @@ class TestMultiPodBatchCostClaim:
@pytest.mark.parametrize("failure_phase", ("before_write", "after_write", "external"))
@pytest.mark.parametrize("collector_enabled", (False, True))
async def test_builtin_accounting_retries_only_before_a_possible_charge(
self, failure_phase: str, collector_enabled: bool
self, failure_phase: str, collector_enabled: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
import base64
from litellm.proxy import proxy_server
monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False)
from datetime import datetime
import httpx

42
tests/unit/proxy/credential_endpoints/test_endpoints.py
View file

@ -2069,7 +2069,10 @@ from litellm.proxy.credential_endpoints import user_provider_credentials as upc
_DEVICE_CODE_URL = "https://github.com/login/device/code"
_ACCESS_TOKEN_URL = "https://github.com/login/oauth/access_token"
_COPILOT_TOKEN_URL = "https://api.github.com/copilot_internal/v2/token"
_COPILOT_USER_URL = "https://api.github.com/copilot_internal/user"
_COPILOT_API_BASE = "https://api.githubcopilot.com"
_COPILOT_MODELS_URL = f"{_COPILOT_API_BASE}/models"
_COPILOT_USER_RESPONSE = {"endpoints": {"api": _COPILOT_API_BASE}}
_GITHUB_USER_URL = "https://api.github.com/user"
@ -2245,7 +2248,8 @@ def test_user_connection_poll_connected_persists_and_reports_login(monkeypatch):
{
_DEVICE_CODE_URL: _DEVICE_FLOW_START,
_ACCESS_TOKEN_URL: {"access_token": "gho_1"},
_COPILOT_TOKEN_URL: {"token": "copilot-tok", "expires_at": 4102444800},
_COPILOT_USER_URL: _COPILOT_USER_RESPONSE,
_COPILOT_MODELS_URL: {},
_GITHUB_USER_URL: {"login": "octo"},
}
):
@ -2276,19 +2280,18 @@ def test_poll_host_change_race_requires_reconnect_on_the_next_request(monkeypatc
table: Final = _patch_user_connection_env(monkeypatch, [_per_user_credential(host="octo.ghe.com")])
def change_host_after_token_exchange(request: httpx.Request) -> dict[str, str]:
def change_host_after_access_token(request: httpx.Request) -> dict[str, str]:
monkeypatch.setattr(litellm, "credential_list", [_per_user_credential()])
return {"access_token": "gho_1"}
with _github_http(
{
"https://octo.ghe.com/login/device/code": _DEVICE_FLOW_START,
"https://octo.ghe.com/login/oauth/access_token": change_host_after_token_exchange,
"https://api.octo.ghe.com/copilot_internal/v2/token": {
"token": "copilot-tok",
"expires_at": 4102444800,
"endpoints": {"api": "https://copilot-api.octo.ghe.com"},
"https://octo.ghe.com/login/oauth/access_token": change_host_after_access_token,
"https://api.octo.ghe.com/copilot_internal/user": {
"endpoints": {"api": "https://copilot-api.octo.ghe.com"}
},
"https://copilot-api.octo.ghe.com/models": {},
"https://api.octo.ghe.com/user": {"login": "octo"},
}
):
@ -2318,7 +2321,7 @@ def test_poll_host_change_race_requires_reconnect_on_the_next_request(monkeypatc
},
}
module_level_client: Final = MagicMock()
module_level_client.get.side_effect = AssertionError("unexpected token exchange")
module_level_client.get.side_effect = AssertionError("unexpected GitHub user lookup")
with _github_http({}), patch.object(litellm, "module_level_client", module_level_client):
with pytest.raises(litellm.CallerCredentialAuthenticationError) as exc:
attach_github_copilot_user_session(kwargs)
@ -2326,17 +2329,16 @@ def test_poll_host_change_race_requires_reconnect_on_the_next_request(monkeypatc
assert stored_connection.github_host == "octo.ghe.com"
def test_ghe_user_connection_uses_tenant_urls_for_device_flow_exchange_and_login(monkeypatch):
def test_ghe_user_connection_uses_tenant_urls_for_device_flow_and_user_lookup(monkeypatch):
table: Final = _patch_user_connection_env(monkeypatch, [_per_user_credential(host="octo.ghe.com")])
with _github_http(
{
"https://octo.ghe.com/login/device/code": _DEVICE_FLOW_START,
"https://octo.ghe.com/login/oauth/access_token": {"access_token": "gho_1"},
"https://api.octo.ghe.com/copilot_internal/v2/token": {
"token": "copilot-tok",
"expires_at": 4102444800,
"endpoints": {"api": "https://copilot-api.octo.ghe.com"},
"https://api.octo.ghe.com/copilot_internal/user": {
"endpoints": {"api": "https://copilot-api.octo.ghe.com"}
},
"https://copilot-api.octo.ghe.com/models": {},
"https://api.octo.ghe.com/user": {"login": "octo"},
}
):
@ -2471,7 +2473,9 @@ def test_user_connection_poll_reports_no_copilot_seat(monkeypatch, seat_status):
return httpx.Response(200, json=_DEVICE_FLOW_START, request=request)
if url == _ACCESS_TOKEN_URL:
return httpx.Response(200, json={"access_token": "gho_seat"}, request=request)
if url == _COPILOT_TOKEN_URL:
if url == _COPILOT_USER_URL:
return httpx.Response(200, json=_COPILOT_USER_RESPONSE, request=request)
if url == _COPILOT_MODELS_URL:
return httpx.Response(seat_status, json={}, request=request)
pytest.fail(f"unexpected request to {url}")
@ -2499,7 +2503,7 @@ def test_user_connection_poll_rate_limit_returns_429(monkeypatch):
return httpx.Response(200, json=_DEVICE_FLOW_START, request=request)
if url == _ACCESS_TOKEN_URL:
return httpx.Response(200, json={"access_token": "gho_limited"}, request=request)
if url == _COPILOT_TOKEN_URL:
if url == _COPILOT_USER_URL:
return httpx.Response(429, json={}, request=request)
pytest.fail(f"unexpected request to {url}")
@ -2875,7 +2879,8 @@ def test_user_connection_poll_valid_handle_works_with_a_fresh_cache(monkeypatch)
{
_DEVICE_CODE_URL: _DEVICE_FLOW_START,
_ACCESS_TOKEN_URL: {"access_token": "gho_1"},
_COPILOT_TOKEN_URL: {"token": "copilot-tok", "expires_at": 4102444800},
_COPILOT_USER_URL: _COPILOT_USER_RESPONSE,
_COPILOT_MODELS_URL: {},
_GITHUB_USER_URL: {"login": "octo"},
}
):
@ -2959,7 +2964,8 @@ def test_user_connection_poll_reports_503_when_the_cache_overwrite_cannot_land(m
{
_DEVICE_CODE_URL: _DEVICE_FLOW_START,
_ACCESS_TOKEN_URL: {"access_token": "gho_1"},
_COPILOT_TOKEN_URL: {"token": "copilot-tok", "expires_at": 4102444800},
_COPILOT_USER_URL: _COPILOT_USER_RESPONSE,
_COPILOT_MODELS_URL: {},
_GITHUB_USER_URL: {"login": "octo"},
}
):

34
tests/unit/proxy/db/test_db_spend_update_writer.py
View file

@ -44,6 +44,15 @@ from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpush
from tests.unit.proxy.db.fake_prisma_engine import engine_call
@pytest.fixture(autouse=True)
def legacy_spend_writers(monkeypatch: pytest.MonkeyPatch) -> None:
"""These tests cover the legacy writers, still reachable as the rollback path with use_spend_logs_ledger off."""
from litellm.proxy import proxy_server
monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False)
@pytest.mark.asyncio
async def test_daily_spend_tracking_with_disabled_spend_logs():
"""
@ -3714,11 +3723,12 @@ async def test_failed_per_entity_increment_from_redis_restores_only_what_may_sti
proxy_logging_obj=proxy_logging_obj,
)
mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once()
restored = mock_redis_update_buffer.restore_transactions_to_redis.call_args.kwargs["db_spend_update_transactions"]
assert restored["user_list_transactions"] is None
assert restored["team_list_transactions"] == {"team-1": 1.5}
assert restored["key_list_transactions"] == ({"key-1": 1.5} if safe_to_resend else None)
restores = mock_redis_update_buffer.restore_transactions_to_redis.await_args_list
assert [
{name: rows for name, rows in call.kwargs["db_spend_update_transactions"].items() if rows}
for call in restores
] == ([{"key_list_transactions": {"key-1": 1.5}}] if safe_to_resend else [])
mock_batcher.litellm_teamtable.update_many.assert_called_once()
mock_batcher.litellm_usertable.update_many.assert_called_once()
@ -3892,13 +3902,13 @@ async def test_commit_spend_updates_to_db_reports_each_completed_table():
)
assert on_table_committed.call_args_list == [
call("user_list_transactions"),
call("end_user_list_transactions"),
call("key_list_transactions"),
call("team_list_transactions"),
call("team_member_list_transactions"),
call("org_list_transactions"),
call("org_member_list_transactions"),
call("team_list_transactions"),
call("user_list_transactions"),
call("end_user_list_transactions"),
call("team_member_list_transactions"),
call("key_list_transactions"),
call("project_list_transactions"),
call("tag_list_transactions"),
call("model_access_group_list_transactions"),
@ -5808,7 +5818,9 @@ async def test_completed_spend_survives_error_log_suppression(
prisma.get_request_status.return_value = status
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "disable_spend_logs", False)
monkeypatch.setattr(proxy_server, "general_settings", {"disable_error_logs": disable_error_logs})
monkeypatch.setattr(
proxy_server, "general_settings", {"disable_error_logs": disable_error_logs, "use_spend_logs_ledger": False}
)
now: Final = datetime(2026, 10, 20, 12, tzinfo=timezone.utc)
result: Final = ModelResponse(
model="gpt-4o-mini", usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)

128
tests/unit/proxy/discovery_endpoints/test_agent_skills_archive.py
View file

@ -4,8 +4,12 @@ import stat
import zipfile
from typing import Final
import pytest
from litellm.proxy.discovery_endpoints.agent_skills_archive import (
MAX_ARCHIVE_ENTRIES,
agent_skill_frontmatter,
build_agent_skill_archive,
build_skill_archive,
)
@ -118,3 +122,127 @@ def test_discovery_keeps_its_permission_and_finder_metadata_policy() -> None:
with zipfile.ZipFile(io.BytesIO(result.content)) as archive:
assert {entry.create_system for entry in archive.infolist()} == {3}
assert {entry.external_attr >> 16 for entry in archive.infolist()} == {0o644}
@pytest.mark.parametrize(
"manifest",
(
pytest.param(b"Instructions without frontmatter", id="missing-frontmatter"),
pytest.param(b"---\n- pdf-summarizer\n- description\n---\n", id="non-mapping"),
pytest.param(b"---\nname: [unterminated\ndescription: invalid\n---\n", id="invalid-yaml"),
pytest.param(b"---\nname: pdf-summarizer\ndescription: !!python/object:unsafe {}\n---\n", id="unsafe-yaml-tag"),
pytest.param(b"---\nname: Bad_Name\ndescription: invalid\n---\n", id="invalid-name"),
pytest.param(b"---\nname: " + b"a" * 65 + b"\ndescription: too long\n---\n", id="name-too-long"),
pytest.param(b"---\nname: pdf-summarizer\n---\n", id="missing-description"),
pytest.param(b"---\nname: pdf-summarizer\ndescription:\n - item\n---\n", id="non-string-description"),
pytest.param(b"---\nname: pdf-summarizer\ndescription: ' '\n---\n", id="empty-description"),
pytest.param(b"---\nname: pdf-summarizer\ndescription: " + b"a" * 1025 + b"\n---\n", id="description-too-long"),
pytest.param(b"---\nname: pdf-summarizer\nname: other-skill\ndescription: ok\n---\n", id="duplicate-name"),
pytest.param(b"---\nname: a\ndescription: ok\nlicense: MIT\nlicense: MIT\n---\n", id="duplicate-optional-key"),
),
)
def test_portable_upload_rejects_malformed_or_invalid_frontmatter(manifest: bytes) -> None:
assert build_agent_skill_archive(zip_bytes({"SKILL.md": manifest})) is None
@pytest.mark.parametrize(
"unsafe_path",
("../escape", "scripts/../../escape", "/absolute", r"scripts\escape", "C:/escape", "./scripts/run.sh"),
)
def test_portable_upload_rejects_unsafe_paths(unsafe_path: str) -> None:
assert build_agent_skill_archive(zip_bytes({"SKILL.md": MANIFEST, unsafe_path: b"unsafe"})) is None
@pytest.mark.parametrize("file_type", (stat.S_IFLNK, stat.S_IFIFO, stat.S_IFCHR))
def test_portable_upload_rejects_links_and_special_files(file_type: int) -> None:
entry: Final = zipfile.ZipInfo("scripts/run.sh")
entry.create_system = 3
entry.external_attr = (file_type | 0o777) << 16
assert build_agent_skill_archive(zip_bytes({"SKILL.md": MANIFEST, entry: b"../../private"})) is None
@pytest.mark.parametrize("size, accepted", ((10 * 1024 * 1024, True), (10 * 1024 * 1024 + 1, False)))
def test_file_size_limit_applies_to_selected_skill_members(size: int, accepted: bool) -> None:
content: Final = zip_bytes(
{
"selected/SKILL.md": MANIFEST,
"selected/asset.bin": b"x" * size,
"unrelated.bin": b"x" * (size + 1),
}
)
assert (build_agent_skill_archive(content, skill_path="selected") is not None) == accepted
def test_portable_upload_rejects_entries_that_normalize_to_the_same_path() -> None:
content: Final = zip_bytes({"SKILL.md": MANIFEST, "scripts//run.sh": b"first", "scripts/run.sh": b"second"})
assert build_agent_skill_archive(content) is None
@pytest.mark.parametrize("bom", (b"", b"\xef\xbb\xbf"))
@pytest.mark.parametrize("newline", ("\n", "\r\n"))
@pytest.mark.parametrize(
("name", "description", "expected"),
(
("on", "yes", ("on", "yes")),
("yes", "on", ("yes", "on")),
('"123"', "'true'", ("123", "true")),
("!!str 123", "!!str null", ("123", "null")),
('"null"', '"false"', ("null", "false")),
("date-reader", "2026-11-05", ("date-reader", "2026-11-05")),
("block-reader", "|-\n yes", ("block-reader", "yes")),
),
)
def test_frontmatter_preserves_yaml12_strings_bom_newlines_and_extensions(
bom: bytes, newline: str, name: str, description: str, expected: tuple[str, str]
) -> None:
frontmatter: Final = (
f"---\nname: {name}\ndescription: {description}\n"
"metadata:\n"
" tags: [on, yes, 123]\n"
"x-extension:\n"
" nested: {count: 1e3, enabled: false}\n"
"---\nPreserve this body and its trailing spaces. \n"
)
manifest: Final = bom + frontmatter.replace("\n", newline).encode()
result: Final = build_agent_skill_archive(zip_bytes({"SKILL.md": manifest}))
assert result is not None
assert (result.declared_name, result.declared_description) == expected
assert result.manifest_content == manifest
@pytest.mark.parametrize("field", ("name", "description"))
@pytest.mark.parametrize("scalar", ("123", "1e3", "012", "0o12", "0x12", "true", "FALSE", "null", "~"))
def test_frontmatter_rejects_non_string_required_scalars(field: str, scalar: str) -> None:
manifest: Final = f"---\n{field}: {scalar}\n".encode() + (
b"description: Instructions\n---\n" if field == "name" else b"name: pdf-summarizer\n---\n"
)
assert build_agent_skill_archive(zip_bytes({"SKILL.md": manifest})) is None
@pytest.mark.parametrize(
"extension",
(
pytest.param("metadata:\n owner: first\n owner: second", id="duplicate-metadata-key"),
pytest.param("metadata:\n owners:\n - {owner: first, owner: second}", id="duplicate-key-in-sequence"),
pytest.param("metadata:\n owner: !!python/object:unsafe {}", id="unsafe-nested-tag"),
pytest.param("metadata:\n count: !!int invalid", id="invalid-explicit-scalar"),
pytest.param("metadata:\n ? [complex, key]\n : value", id="complex-key"),
pytest.param("metadata: &loop\n self: *loop", id="recursive-alias"),
pytest.param("metadata: " + "[" * 40 + "nested" + "]" * 40, id="excessive-depth"),
pytest.param("metadata: [unterminated", id="malformed-extension"),
),
)
def test_frontmatter_rejects_invalid_nested_metadata(extension: str) -> None:
manifest: Final = f"---\nname: pdf-summarizer\ndescription: Instructions\n{extension}\n---\n".encode()
assert build_agent_skill_archive(zip_bytes({"SKILL.md": manifest})) is None
@pytest.mark.parametrize("name", ("résumé", "мой-навык", "技能", "ｆｏｏ"))
def test_hub_registry_retains_its_ascii_name_policy(name: str) -> None:
manifest: Final = f"---\nname: {name}\ndescription: Instructions\n---\n".encode()
assert agent_skill_frontmatter(manifest) == {"name": name, "description": "Instructions"}
assert build_agent_skill_archive(zip_bytes({"SKILL.md": manifest})) is None

127
tests/unit/proxy/discovery_endpoints/test_agent_skills_endpoints.py
View file

@ -1,21 +1,36 @@
import hashlib
import io
import json
import zipfile
from collections.abc import Mapping
from datetime import datetime, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from prisma import Base64
from prisma import models as prisma_models
import litellm
from litellm.models.skills import LiteLLM_SkillsTable
from litellm.proxy.anthropic_endpoints.claude_code_endpoints.skill_upload_archive import build_uploaded_skill_archive
from litellm.proxy.deps.repositories import get_plugin_upload_repository
from litellm.proxy.discovery_endpoints.agent_skills_archive import build_agent_skill_archive
from litellm.proxy.discovery_endpoints.agent_skills_endpoints import (
MAX_CACHED_ARCHIVE_BYTES,
MAX_INDEXED_SKILLS,
_ARCHIVE_CACHE,
router,
stored_skill,
stored_skills,
)
from litellm.proxy.discovery_endpoints.skill_hub_endpoints import router as hub_router
from litellm.repositories.plugin_upload_repository import PluginUploadRepository
from litellm.types.proxy.discovery_endpoints.agent_skills_endpoints import (
AGENT_SKILLS_DISCOVERY_SCHEMA_URL,
MAX_SKILL_NAME_LENGTH,
)
WELL_KNOWN_PATHS = ("/.well-known/agent-skills/index.json", "/.well-known/skills/index.json")
@ -54,9 +69,10 @@ def skill(
)
def client_for(*skills: LiteLLM_SkillsTable) -> TestClient:
def client_for(*skills: LiteLLM_SkillsTable, uploads: PluginUploadRepository | None = None) -> TestClient:
app = FastAPI()
app.include_router(router)
app.include_router(hub_router)
def _skills() -> tuple[LiteLLM_SkillsTable, ...]:
return skills
@ -66,9 +82,51 @@ def client_for(*skills: LiteLLM_SkillsTable) -> TestClient:
app.dependency_overrides[stored_skills] = _skills
app.dependency_overrides[stored_skill] = _skill
app.dependency_overrides[get_plugin_upload_repository] = lambda: uploads
return TestClient(app)
def hub_record(
name: str,
) -> tuple[prisma_models.LiteLLM_ClaudeCodePluginTable, prisma_models.LiteLLM_ClaudeCodePluginUpload]:
manifest: Final = MANIFEST.replace(b"pdf-summarizer", name.encode())
portable: Final = build_agent_skill_archive(zip_bytes({"SKILL.md": manifest}))
assert portable is not None
wrapper: Final = build_uploaded_skill_archive(portable.content, name)
assert wrapper is not None
digest: Final = hashlib.sha256(wrapper).hexdigest()
return (
prisma_models.LiteLLM_ClaudeCodePluginTable(
id=f"hub-{name}", name=name, enabled=True, anonymous_download_count=0, manifest_json=json.dumps({
"source": {
"source": "archive", "sha256": digest,
"url": f"http://testserver/claude-code/plugins/{name}/archive/{digest}.zip",
},
"agent_skill": {"name": name, "description": portable.declared_description, "digest": portable.digest},
}),
),
prisma_models.LiteLLM_ClaudeCodePluginUpload(plugin_name=name, content=Base64.encode(wrapper)),
)
def hub_repository(*names: str) -> PluginUploadRepository:
records: Final = tuple(hub_record(name) for name in names)
prisma: Final = MagicMock()
prisma.writer_db.execute_raw = AsyncMock(return_value=1)
plugins: Final = prisma.writer_db.litellm_claudecodeplugintable
plugins.find_many = AsyncMock(return_value=[record for record, _ in records])
async def metadata(where: Mapping[str, str]) -> prisma_models.LiteLLM_ClaudeCodePluginTable | None:
return next((record for record, _ in records if record.name == where["name"]), None)
async def content(where: Mapping[str, object]) -> prisma_models.LiteLLM_ClaudeCodePluginUpload | None:
return next((archive for _, archive in records if archive.plugin_name == where["plugin_name"]), None)
plugins.find_unique = AsyncMock(side_effect=metadata)
prisma.writer_db.litellm_claudecodepluginupload.find_first = AsyncMock(side_effect=content)
return PluginUploadRepository(prisma)
@pytest.fixture
def index_enabled(monkeypatch):
monkeypatch.setattr(litellm, "public_skills_index", True)
@ -149,6 +207,55 @@ def test_skills_sharing_a_title_get_distinct_install_names(index_enabled):
assert names == ["pdf-summarizer", "pdf-summarizer-2"]
@pytest.mark.parametrize("hosted_names, hub_names, aliases", (
(("foo",), ("foo", "foo-2"), ("foo-3", "foo-2")),
(("a" * 64,), ("a" * 64, "a" * 62 + "-2"), ("a" * 62 + "-3", "a" * 62 + "-2")),
(("foo", "foo"), ("foo", "foo-2"), ("foo-3", "foo-2-2")),
(("foo", "foo", "foo-2"), ("foo", "foo-2", "foo-3"), ("foo-4", "foo-2-2", "foo-3-2")),
(("bar", "bar", "foo", "foo"), ("bar", "bar-2", "foo", "foo-2"), ("bar-3", "bar-2-2", "foo-3", "foo-2-2")),
))
def test_hub_collisions_keep_every_hosted_entry_and_only_alias_hub_install_names(
index_enabled: None, hosted_names: tuple[str, ...], hub_names: tuple[str, ...], aliases: tuple[str, ...],
) -> None:
hosted: Final = tuple(
skill(f"hosted-{position}", files={"s/SKILL.md": MANIFEST.replace(b"pdf-summarizer", name.encode())})
for position, name in enumerate(hosted_names)
)
before: Final = client_for(*reversed(hosted)).get(WELL_KNOWN_PATHS[0]).json()["skills"]
client: Final = client_for(*reversed(hosted), uploads=hub_repository(*hub_names))
response: Final = client.get(WELL_KNOWN_PATHS[0])
assert response.status_code == 200, response.text
entries: Final = response.json()["skills"]
assert entries[:len(hosted)] == before
assert len({entry["name"] for entry in entries}) == len(hosted) + len(hub_names)
assert all(len(entry["name"]) <= MAX_SKILL_NAME_LENGTH for entry in entries)
assert tuple(entry["name"] for entry in entries[len(hosted):]) == aliases
assert client.get(WELL_KNOWN_PATHS[1]).json() == response.json()
assert client.get(WELL_KNOWN_PATHS[0]).json() == response.json()
for entry, declared_name in zip(entries, hosted_names + hub_names, strict=True):
downloaded: Final = client.get(entry["url"])
assert downloaded.status_code == 200, downloaded.text
assert entry["digest"] == f"sha256:{hashlib.sha256(downloaded.content).hexdigest()}"
with zipfile.ZipFile(io.BytesIO(downloaded.content)) as archive:
assert archive.read("SKILL.md") == MANIFEST.replace(b"pdf-summarizer", declared_name.encode())
@pytest.mark.parametrize("hosted_count", (MAX_INDEXED_SKILLS - 1, MAX_INDEXED_SKILLS + 1))
def test_combined_index_cap_keeps_hosted_entries_before_hub_entries(index_enabled: None, hosted_count: int) -> None:
hosted: Final = tuple(skill(f"hosted-{position}") for position in range(hosted_count))
client: Final = client_for(*hosted, uploads=hub_repository("first", "second"))
response: Final = client.get(WELL_KNOWN_PATHS[0])
assert response.status_code == 200, response.text
entries: Final = response.json()["skills"]
retained: Final = min(hosted_count, MAX_INDEXED_SKILLS)
assert len(entries) == MAX_INDEXED_SKILLS
assert len({entry["name"] for entry in entries}) == MAX_INDEXED_SKILLS
assert [entry["url"] for entry in entries[:retained]] == [
f"http://testserver/v1/skills/{item.skill_id}/archive" for item in reversed(hosted[:MAX_INDEXED_SKILLS])
]
assert [entry["name"] for entry in entries[retained:]] == (["first"] if retained < MAX_INDEXED_SKILLS else [])
def test_uploads_without_a_root_manifest_are_left_out_of_the_index(index_enabled):
client = client_for(
skill("litellm_skill_1"),
@ -182,6 +289,24 @@ def test_a_stored_skill_is_repacked_once_per_version(index_enabled):
)
def test_retained_manifest_bytes_count_toward_the_archive_cache_limit(
index_enabled: None,
) -> None:
stamp: Final = datetime(2026, 9, 6, 9, 0, tzinfo=timezone.utc)
name: Final = "litellm_skill_large_manifest"
manifest: Final = MANIFEST + b"x" * MAX_CACHED_ARCHIVE_BYTES
client: Final = client_for(skill(name, files={"s/SKILL.md": manifest}, updated_at=stamp))
response: Final = client.get(f"/v1/skills/{name}/archive")
assert response.status_code == 200
assert len(response.content) < MAX_CACHED_ARCHIVE_BYTES
with zipfile.ZipFile(io.BytesIO(response.content)) as archive:
stored: Final = archive.read("SKILL.md")
actual: Final = (len(stored), hashlib.sha256(stored).hexdigest())
expected: Final = (len(manifest), hashlib.sha256(manifest).hexdigest())
assert actual == expected
assert _ARCHIVE_CACHE.get_cache(f"{name}:{stamp.isoformat()}") is None
def test_a_skill_edited_since_the_last_read_is_republished(index_enabled):
stamp = datetime(2026, 9, 6, 9, 0, tzinfo=timezone.utc)
before = client_for(skill("litellm_skill_edited", files={"s/SKILL.md": MANIFEST}, updated_at=stamp))

808
tests/unit/proxy/discovery_endpoints/test_skill_hub_endpoints.py Normal file
View file

@ -0,0 +1,808 @@
from __future__ import annotations
import asyncio
import hashlib
import io
import json
import stat
import threading
import zipfile
import zlib
from collections.abc import AsyncIterator, Mapping
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
from prisma import models as prisma_models
from prisma.errors import ForeignKeyViolationError, UniqueViolationError
import litellm
from litellm.caching.caching import DualCache
from litellm.proxy import proxy_server
from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth
from litellm.proxy.deps.repositories import get_plugin_upload_repository, get_skill_submission_repository
from litellm.proxy.discovery_endpoints.agent_skills_archive import (
MAX_ARCHIVE_ENTRIES,
MAX_ARCHIVE_UNPACKED_BYTES,
)
from litellm.proxy.discovery_endpoints.skill_hub_endpoints import import_http_client
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
from litellm.proxy.utils import InternalUsageCache
from litellm.repositories.plugin_upload_repository import PluginUploadRepository
from litellm.repositories.skill_submission_repository import SkillSubmissionRepository
from tests._master_key import MASTER_KEY
from tests.unit.proxy.skill_submission_rows import SubmissionRows
NAME: Final = "portable-notes"
MANIFEST: Final = (
b"---\nname: portable-notes\ndescription: Summarize synthetic notes.\n---\n\nRead references/guide.md.\n"
)
FILES: Final = {
"SKILL.md": MANIFEST,
"references/guide.md": b"Use only the supplied notes.\n",
"assets/example.bin": b"\x00\x01\xff",
"scripts/check.sh": b"#!/bin/sh\nprintf 'synthetic notes\\n'\n",
}
HUB: Final = "/skill-hub/skills"
INDEX: Final = "/.well-known/agent-skills/index.json"
GRANTED_KEY: Final = "sk-hub-granted"
PLAIN_KEY: Final = "sk-hub-plain"
MEMBER_KEY: Final = "sk-hub-member"
OTHER_KEY: Final = "sk-hub-other"
VIEWER_KEY: Final = "sk-hub-viewer"
def _zip(files: Mapping[str, bytes], prefix: str = "") -> bytes:
buffer: Final = io.BytesIO()
with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive:
for path, content in files.items():
entry: Final = zipfile.ZipInfo(f"{prefix}{path}", (2024, 1, 1, 0, 0, 0))
entry.create_system = 3
entry.external_attr = (stat.S_IFREG | (0o755 if path.endswith(".sh") else 0o644)) << 16
archive.writestr(entry, content, compress_type=zipfile.ZIP_DEFLATED)
return buffer.getvalue()
def _headers(key: str | None) -> dict[str, str]:
return {"Authorization": f"Bearer {key}"} if key is not None else {}
@dataclass(frozen=True, slots=True)
class HubClient:
client: httpx.AsyncClient
plugins: MagicMock
archives: MagicMock
downloads: MagicMock
anonymous_counter: AsyncMock
upstream: MagicMock
clock: MagicMock
database: MagicMock
@pytest.fixture
async def hub(monkeypatch: pytest.MonkeyPatch) -> AsyncIterator[HubClient]:
prisma: Final = MagicMock()
prisma.writer_db = prisma.db
prisma.db.execute_raw = AsyncMock(return_value=1)
prisma.get_data = AsyncMock(return_value=None)
prisma.db.litellm_skillstable.find_many = AsyncMock(return_value=[])
prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
plugins: Final = prisma.db.litellm_claudecodeplugintable
archives: Final = prisma.db.litellm_claudecodepluginupload
async def create_row(data: Mapping[str, object]) -> prisma_models.LiteLLM_ClaudeCodePluginTable:
return prisma_models.LiteLLM_ClaudeCodePluginTable.model_validate({"id": "hub-id", "enabled": True, "anonymous_download_count": 0, **data})
plugins.create = AsyncMock(side_effect=create_row)
downloads: Final = prisma.db.litellm_claudecodeplugindownload
downloads.create = AsyncMock()
plugins.find_unique = AsyncMock(return_value=None)
plugins.find_many = AsyncMock(return_value=[])
async def update_row(where: Mapping[str, str], data: Mapping[str, object]) -> prisma_models.LiteLLM_ClaudeCodePluginTable:
current: Final = plugins.find_unique.return_value.model_copy(update=dict(data))
plugins.find_unique.return_value = current
plugins.find_many.return_value = [current]
return current
plugins.update = AsyncMock(side_effect=update_row)
plugins.delete = AsyncMock()
archives.find_first = AsyncMock(return_value=None)
archives.find_unique = AsyncMock(return_value=None)
archives.create = AsyncMock()
archives.upsert = AsyncMock()
archives.delete_many = AsyncMock()
prisma.db.tx.return_value.__aenter__.return_value = prisma.db
submissions: Final = SubmissionRows()
prisma.db.litellm_skillsubmission = submissions
prisma.db.query_raw = AsyncMock(return_value=[])
prisma.db.litellm_claudecodepluginversion.create = AsyncMock()
prisma.db.litellm_claudecodepluginversion.count = AsyncMock(return_value=0)
prisma.db.litellm_claudecodepluginversion.find_first = AsyncMock(return_value=None)
upstream: Final = MagicMock(return_value=httpx.Response(404))
clock: Final = MagicMock(return_value=datetime.now())
limiter: Final = PROXY_MaxParallelRequestsHandler_v3(InternalUsageCache(DualCache()), time_provider=clock)
monkeypatch.setattr(proxy_server.proxy_logging_obj, "proxy_hook_mapping", {"parallel_request_limiter": limiter})
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "master_key", MASTER_KEY)
monkeypatch.setattr(proxy_server, "general_settings", {})
monkeypatch.setattr(proxy_server, "user_api_key_cache", DualCache())
monkeypatch.setattr(litellm, "public_skills_index", True)
monkeypatch.setenv("PROXY_BASE_URL", "http://test")
monkeypatch.delenv("SERVER_ROOT_PATH", raising=False)
for key, skills in ((GRANTED_KEY, [NAME]), (PLAIN_KEY, [])):
proxy_server.user_api_key_cache.set_cache(
key=proxy_server.hash_token(key),
value=UserAPIKeyAuth(
token=proxy_server.hash_token(key),
user_role=proxy_server.LitellmUserRoles.INTERNAL_USER,
object_permission=LiteLLM_ObjectPermissionTable(object_permission_id=key, skills=skills),
),
)
for key, role in (
(MEMBER_KEY, proxy_server.LitellmUserRoles.INTERNAL_USER),
(OTHER_KEY, proxy_server.LitellmUserRoles.INTERNAL_USER),
(VIEWER_KEY, proxy_server.LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY),
):
proxy_server.user_api_key_cache.set_cache(
key=proxy_server.hash_token(key),
value=UserAPIKeyAuth(token=proxy_server.hash_token(key), user_role=role, user_id=key),
)
async with (
httpx.AsyncClient(transport=httpx.MockTransport(upstream), trust_env=False) as source_client,
httpx.AsyncClient(transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test") as client,
):
monkeypatch.setattr(proxy_server.app, "dependency_overrides", {
get_plugin_upload_repository: lambda: PluginUploadRepository(prisma),
get_skill_submission_repository: lambda: SkillSubmissionRepository(prisma),
import_http_client: lambda: source_client,
})
yield HubClient(client, plugins, archives, downloads, prisma.db.execute_raw, upstream, clock, prisma.db)
async def _upload(
hub: HubClient, content: bytes, key: str | None = MASTER_KEY, version: str | None = "2.0.0",
) -> httpx.Response:
return await hub.client.post(
f"{HUB}/upload",
headers=_headers(key),
data={"version": version} if version is not None else {},
files={"file": ("notes.zip", content, "application/zip")},
)
def _stored_upload(hub: HubClient, enabled: bool = True) -> prisma_models.LiteLLM_ClaudeCodePluginTable:
data: Final = hub.plugins.create.await_args.kwargs["data"]
record: Final = prisma_models.LiteLLM_ClaudeCodePluginTable(
id="hub-id", name=data["name"], version=data["version"], description=data["description"], anonymous_download_count=0,
manifest_json=data["manifest_json"], enabled=enabled, updated_at=data["updated_at"],
)
hub.plugins.find_unique.return_value = record
hub.plugins.find_many.return_value = [record]
hub.archives.find_first.return_value = prisma_models.LiteLLM_ClaudeCodePluginUpload(
plugin_name=record.name, content=hub.archives.create.await_args.kwargs["data"]["content"],
)
hub.archives.find_unique.return_value = hub.archives.find_first.return_value
return record
def _assert_portable(response: httpx.Response) -> None:
assert response.status_code == 200, response.text
assert response.headers["content-type"] == "application/zip"
with zipfile.ZipFile(io.BytesIO(response.content)) as archive:
assert {name: archive.read(name) for name in archive.namelist()} == FILES
assert stat.S_IMODE(archive.getinfo("scripts/check.sh").external_attr >> 16) == 0o755
@pytest.mark.parametrize("key, expected", ((MASTER_KEY, 7), (PLAIN_KEY, None)))
async def test_catalog_counts_are_admin_only_and_batched_for_the_page(
hub: HubClient, key: str, expected: int | None,
) -> None:
assert (await _upload(hub, _zip(FILES))).status_code == 200
record: Final = _stored_upload(hub)
hub.plugins.find_many.return_value = [record, record.model_copy(update={"id": "other", "name": "zzz-other"})]
hub.database.query_raw.reset_mock()
hub.database.query_raw.return_value = [{"id": record.id, "download_count": 7}]
response: Final = await hub.client.get(f"{HUB}?limit=1", headers=_headers(key))
assert response.status_code == 200, response.text
assert [(item["id"], item["download_count"]) for item in response.json()["skills"]] == [(record.id, expected)]
if expected is None:
hub.database.query_raw.assert_not_awaited()
else:
assert hub.database.query_raw.await_count == 1
assert hub.database.query_raw.await_args.args[1] == (record.id,)
assert (await hub.client.get(f"{HUB}?offset=2", headers=_headers(key))).json()["skills"] == []
assert hub.database.query_raw.await_count == 1
@pytest.mark.parametrize("operation", ("upload", "import"))
@pytest.mark.parametrize("public_base, root_path", (
("http://test", ""), ("https://public.example", ""), ("https://public.example", "/gateway"),
))
async def test_publication_round_trips_both_formats_without_inflating_on_the_request_thread(
hub: HubClient, monkeypatch: pytest.MonkeyPatch, operation: str, public_base: str, root_path: str,
) -> None:
monkeypatch.delenv("PROXY_BASE_URL")
hub.upstream.return_value = httpx.Response(200, content=_zip(FILES, "repo-revision/"))
response: Final = (
await _upload(hub, _zip(FILES)) if operation == "upload" else await hub.client.post(
f"{HUB}/import", headers=_headers(MASTER_KEY),
json={"url": f"https://skills.sh/synthetic-author/notes-library/{NAME}"},
)
)
assert response.status_code == 200, response.text
assert response.json()["version"] == ("2.0.0" if operation == "upload" else "1.0.0")
record: Final = _stored_upload(hub)
manifest: Final = json.loads(record.manifest_json)
monkeypatch.setattr(proxy_server.app, "root_path", root_path)
monkeypatch.setenv("SERVER_ROOT_PATH", root_path)
serving_base: Final = public_base + root_path
hub.client.base_url = httpx.URL(serving_base)
details: Final = await hub.client.get(f"{HUB}/{NAME}", headers=_headers(PLAIN_KEY))
listed: Final = await hub.client.get(HUB, headers=_headers(PLAIN_KEY))
assert listed.json()["skills"] == [details.json()]
assert "has_agent_skill_metadata" not in listed.text + details.text
assert details.json()["download_url"] == f"{serving_base}{HUB}/{NAME}/archive"
assert details.json()["content_url"] == f"{serving_base}{HUB}/{NAME}/content"
catalog: Final = await hub.client.get("/claude-code/plugins", headers=_headers(PLAIN_KEY))
assert catalog.status_code == 200 and "has_agent_skill_metadata" not in catalog.text
hub.downloads.create.assert_not_awaited()
hub.anonymous_counter.assert_not_awaited()
request_thread: Final = threading.get_ident()
decompress: Final = zlib.decompressobj
def off_thread_decompressor(wbits: int) -> zlib.Decompress:
assert threading.get_ident() != request_thread
return decompress(wbits)
monkeypatch.setattr(zlib, "decompressobj", off_thread_decompressor)
portable: Final = await hub.client.get(details.json()["download_url"])
content: Final = await hub.client.get(details.json()["content_url"])
assert content.status_code == 200 and content.content == MANIFEST
digest: Final = "sha256:" + hashlib.sha256(portable.content).hexdigest()
assert manifest["agent_skill"]["digest"] == digest
entry: Final = (await hub.client.get(INDEX)).json()["skills"][0]
assert entry["digest"] == digest
assert entry["url"].startswith(f"{serving_base}{HUB}/{NAME}/archive?")
assert (await hub.client.get(entry["url"])).content == portable.content
assert (await hub.client.get(f"{HUB}/{NAME}/archive?digest=sha256:stale")).status_code == 404
legacy_path: Final = httpx.URL(manifest["source"]["url"]).path
legacy: Final = await hub.client.get(legacy_path, headers=_headers(GRANTED_KEY))
assert legacy.status_code == 200
assert hashlib.sha256(legacy.content).hexdigest() == manifest["source"]["sha256"]
assert (await hub.client.get(legacy_path.replace(manifest["source"]["sha256"], "stale"))).status_code == 404
assert hub.anonymous_counter.await_count == 2
assert hub.anonymous_counter.await_args.args[1:] == (record.id,)
hub.downloads.create.assert_awaited_once()
monkeypatch.setattr(zlib, "decompressobj", decompress)
_assert_portable(portable)
with zipfile.ZipFile(io.BytesIO(legacy.content)) as archive:
assert {path: archive.read(path) for path in archive.namelist()} == {
".claude-plugin/plugin.json": json.dumps({"name": NAME}).encode(),
**{f"skills/{NAME}/{path}": content for path, content in FILES.items()},
}
@pytest.mark.parametrize("operation", ("upload", "import"))
@pytest.mark.parametrize("boundary", ("entries", "bytes", "file"))
@pytest.mark.parametrize("fits", (True, False))
async def test_publication_reserves_wrapper_capacity_before_persisting(
hub: HubClient, operation: str, boundary: str, fits: bool,
) -> None:
file_limit: Final = 10 * 1024 * 1024
padding_size: Final = (
MAX_ARCHIVE_UNPACKED_BYTES - len(MANIFEST) - len(json.dumps({"name": NAME}).encode()) * int(fits)
if boundary == "bytes" else file_limit + int(not fits)
)
files: Final = (
{"SKILL.md": MANIFEST, **{
f"references/{index}.txt": b"x" for index in range(MAX_ARCHIVE_ENTRIES - 1 - int(fits))
}} if boundary == "entries" else {
"SKILL.md": MANIFEST,
**({
f"assets/{index}.bin": b"x" * min(file_limit, padding_size - index * file_limit)
for index in range(5)
} if boundary == "bytes" else {"assets/padding.bin": b"x" * padding_size}),
}
)
content: Final = _zip(files, prefix="repo-revision/" if operation == "import" else "")
hub.upstream.return_value = httpx.Response(200, content=content)
response: Final = (
await _upload(hub, content) if operation == "upload" else await hub.client.post(
f"{HUB}/import", headers=_headers(MASTER_KEY),
json={"url": "https://github.com/synthetic-author/notes-library"},
)
)
assert response.status_code == (200 if fits else 422), response.text
if not fits:
hub.plugins.create.assert_not_awaited()
return
_stored_upload(hub)
downloaded: Final = await hub.client.get(response.json()["download_url"])
assert downloaded.status_code == 200, downloaded.text
with zipfile.ZipFile(io.BytesIO(downloaded.content)) as archive:
actual: Final = {
path: (archive.getinfo(path).file_size, hashlib.sha256(archive.read(path)).hexdigest())
for path in archive.namelist()
}
expected: Final = {path: (len(content), hashlib.sha256(content).hexdigest()) for path, content in files.items()}
assert actual == expected
@pytest.mark.parametrize("enabled, public_index, auth_required, key, status", (
(True, True, False, None, 200), (False, True, False, None, 404),
(True, False, False, None, 401), (True, True, True, None, 401),
(False, True, False, GRANTED_KEY, 200), (False, True, False, PLAIN_KEY, 404),
(True, True, False, "", 401), (True, True, False, "sk-invalid", 401),
))
async def test_download_policy_requires_publication_or_a_valid_grant(
hub: HubClient, monkeypatch: pytest.MonkeyPatch, enabled: bool,
public_index: bool, auth_required: bool, key: str | None, status: int,
) -> None:
assert (await _upload(hub, _zip(FILES))).status_code == 200
_stored_upload(hub, enabled=enabled)
monkeypatch.setattr(litellm, "public_skills_index", public_index)
monkeypatch.setattr(proxy_server, "general_settings", {"claude_code_marketplace_auth_required": auth_required})
for suffix in ("archive", "content"):
response: Final = await hub.client.get(f"{HUB}/{NAME}/{suffix}", headers=_headers(key))
assert response.status_code == status, response.text
assert hub.downloads.create.await_count == int(status == 200 and key is not None)
assert hub.anonymous_counter.await_count == int(status == 200 and key is None)
if status != 200:
hub.archives.find_first.assert_not_awaited()
index: Final = await hub.client.get(INDEX)
assert index.status_code == (200 if public_index else 404)
if public_index:
assert [entry["name"] for entry in index.json()["skills"]] == ([NAME] if enabled and not auth_required else [])
@pytest.mark.parametrize("key", (None, GRANTED_KEY))
async def test_deleted_skill_is_not_served_when_download_cannot_be_recorded(hub: HubClient, key: str | None) -> None:
assert (await _upload(hub, _zip(FILES))).status_code == 200
record: Final = _stored_upload(hub)
legacy_url: Final = json.loads(record.manifest_json)["source"]["url"]
hub.anonymous_counter.return_value = 0
hub.downloads.create.side_effect = ForeignKeyViolationError(data={})
for path in (f"{HUB}/{NAME}/archive", legacy_url):
response: Final = await hub.client.get(path, headers=_headers(key))
assert response.status_code == 404, response.text
assert "Skill archive not found" in response.text
assert hub.anonymous_counter.await_count == (2 if key is None else 0)
assert hub.downloads.create.await_count == (0 if key is None else 2)
async def test_search_paginates_visible_matches(hub: HubClient) -> None:
assert (await _upload(hub, _zip(FILES))).status_code == 200
record: Final = _stored_upload(hub)
hub.plugins.find_many.return_value = [
record.model_copy(update={"name": "private-notes", "enabled": False}), record,
record.model_copy(update={"name": "other-notes"}),
record.model_copy(update={"name": "unmatched", "description": "Different topic"}),
]
for offset, name in enumerate(("other-notes", NAME)):
response: Final = await hub.client.get(
HUB, headers=_headers(PLAIN_KEY), params={"query": "SYNTHETIC", "limit": 1, "offset": offset},
)
assert response.status_code == 200
assert [item["name"] for item in response.json()["skills"]] == [name]
assert response.json()["count"] == 2
assert response.json()["next_offset"] == (1 if offset == 0 else None)
async def test_metadata_publication_and_deletion_preserve_files_and_grants(hub: HubClient) -> None:
assert (await _upload(hub, _zip(FILES))).status_code == 200
original: Final = _stored_upload(hub)
manifest: Final = json.loads(original.manifest_json)
for route, metadata in (
(f"/claude-code/plugins/{NAME}", {"source": manifest["source"], "description": "Legacy edit"}),
(f"{HUB}/{NAME}", {"description": "Neutral edit", "version": "3.0.0", "keywords": ["notes"]}),
(f"{HUB}/{NAME}", {}),
):
response: Final = await hub.client.put(route, headers=_headers(MASTER_KEY), json=metadata)
assert response.status_code == 200, response.text
assert json.loads(hub.plugins.find_unique.return_value.manifest_json)["agent_skill"] == manifest["agent_skill"]
details: Final = await hub.client.get(f"{HUB}/{NAME}", headers=_headers(PLAIN_KEY))
assert details.json()["description"] == metadata.get("description")
_assert_portable(await hub.client.get(details.json()["download_url"]))
for action, public in (("unpublish", False), ("publish", True)):
changed: Final = await hub.client.post(f"{HUB}/{NAME}/{action}", headers=_headers(MASTER_KEY))
assert changed.status_code == 200
assert changed.json()["visibility"] == ("public" if public else "private")
assert (await hub.client.get(f"{HUB}/{NAME}/archive")).status_code == (200 if public else 404)
_assert_portable(await hub.client.get(f"{HUB}/{NAME}/archive", headers=_headers(GRANTED_KEY)))
assert (await hub.client.get(f"{HUB}/{NAME}", headers=_headers(PLAIN_KEY))).status_code == (200 if public else 404)
assert bool((await hub.client.get(INDEX)).json()["skills"]) == public
assert (await hub.client.delete(f"{HUB}/{NAME}", headers=_headers(MASTER_KEY))).status_code == 204
hub.plugins.delete.assert_awaited_once_with(where={"name": NAME})
@pytest.mark.parametrize("method, suffix", (("PUT", ""), ("POST", "/publish"), ("POST", "/unpublish"), ("DELETE", "")))
async def test_granted_readers_cannot_manage_skills(hub: HubClient, method: str, suffix: str) -> None:
response: Final = await hub.client.request(
method, f"{HUB}/{NAME}{suffix}", headers=_headers(GRANTED_KEY), json={} if method == "PUT" else None,
)
assert response.status_code == 403
hub.plugins.update.assert_not_awaited()
hub.plugins.delete.assert_not_awaited()
@pytest.mark.parametrize("content, key, status", (
(b"not a zip", MASTER_KEY, 422),
(_zip(FILES), PLAIN_KEY, 403), (_zip(FILES), None, 401),
), ids=("invalid-zip", "forbidden", "unauthenticated"))
async def test_rejected_uploads_do_not_write(hub: HubClient, content: bytes, key: str | None, status: int) -> None:
assert (await _upload(hub, content, key)).status_code == status
hub.plugins.create.assert_not_awaited()
async def test_stored_content_survives_duplicate_upload_and_is_hidden_after_source_replacement(hub: HubClient) -> None:
assert (await _upload(hub, _zip(FILES))).status_code == 200
original: Final = _stored_upload(hub)
hub.plugins.create.side_effect = UniqueViolationError({"message": "duplicate", "error_code": "P2002"})
response: Final = await _upload(hub, _zip({**FILES, "references/guide.md": b"replacement"}))
assert response.status_code == 409
_assert_portable(await hub.client.get(f"{HUB}/{NAME}/archive"))
hub.plugins.update.assert_not_awaited()
hub.archives.find_first.reset_mock()
manifest: Final = json.loads(original.manifest_json)
external: Final = {
**manifest["source"],
"url": str(httpx.URL(manifest["source"]["url"]).copy_with(scheme="https", host="external.example")),
}
replaced: Final = await hub.client.put(
f"/claude-code/plugins/{NAME}", headers=_headers(MASTER_KEY),
json={"source": external, "agent_skill": manifest["agent_skill"]},
)
assert replaced.status_code == 200, replaced.text
assert "agent_skill" not in json.loads(hub.plugins.find_unique.return_value.manifest_json)
details: Final = await hub.client.get(f"{HUB}/{NAME}", headers=_headers(MASTER_KEY))
assert details.json()["source"] == external and details.json()["download_url"] is None
for suffix in ("archive", "content"):
assert (await hub.client.get(f"{HUB}/{NAME}/{suffix}")).status_code == 404
assert (await hub.client.get(httpx.URL(manifest["source"]["url"]).path)).status_code == 404
hub.archives.find_first.assert_not_awaited()
assert (await hub.client.get(INDEX)).json()["skills"] == []
@pytest.mark.parametrize("user_id, role, admitted", (
("publisher", proxy_server.LitellmUserRoles.PROXY_ADMIN, 10),
(None, proxy_server.LitellmUserRoles.PROXY_ADMIN, 20),
("publisher", proxy_server.LitellmUserRoles.INTERNAL_USER, 10),
), ids=("admin-user", "admin-key", "member-user"))
async def test_publication_limit_is_shared_by_user_and_route_and_resets(
hub: HubClient, monkeypatch: pytest.MonkeyPatch, user_id: str | None, role: str, admitted: int,
) -> None:
for key in (PLAIN_KEY, GRANTED_KEY):
proxy_server.user_api_key_cache.set_cache(
key=proxy_server.hash_token(key),
value=UserAPIKeyAuth(
token=proxy_server.hash_token(key), user_id=user_id,
user_role=role,
),
)
async def attempt(index: int) -> httpx.Response:
key: Final = PLAIN_KEY if index % 2 else GRANTED_KEY
if index % 3 == 0:
return await _upload(hub, b"invalid", key)
if index % 3 == 1:
return await hub.client.post(
f"{HUB}/import", headers=_headers(key), json={"url": "https://example.com/skill"},
)
return await hub.client.put(
f"{HUB}/{NAME}/upload", headers=_headers(key), files={"file": ("invalid.zip", b"invalid")},
)
responses: Final = await asyncio.gather(*(attempt(index) for index in range(20)))
assert sum(response.status_code in (400, 404, 422) for response in responses) == admitted
assert sum(response.status_code == 429 for response in responses) == 20 - admitted
assert all(response.headers["retry-after"] == "60" for response in responses if response.status_code == 429)
hub.plugins.create.assert_not_awaited()
hub.upstream.assert_not_called()
assert (await _upload(hub, _zip(FILES))).status_code == 200
_stored_upload(hub)
for _ in range(11):
assert (await hub.client.get(INDEX)).status_code == 200
assert (await hub.client.get(f"{HUB}/{NAME}/archive", headers=_headers(PLAIN_KEY))).status_code == 200
hub.clock.return_value += timedelta(seconds=60)
assert (await _upload(hub, b"invalid", GRANTED_KEY)).status_code == 422
monkeypatch.setattr(proxy_server.proxy_logging_obj, "proxy_hook_mapping", {})
assert (await _upload(hub, b"invalid")).status_code == 503
@pytest.mark.asyncio
async def test_neutral_metadata_publication_pairs_version_and_preserved_skill_snapshot(
hub: HubClient, monkeypatch: pytest.MonkeyPatch,
) -> None:
created: Final = await _upload(hub, _zip(FILES), version=None)
assert created.status_code == 200, created.text
assert created.json()["version"] == "1.0.0"
original: Final = _stored_upload(hub)
hub.plugins.find_unique.return_value = original.model_copy(update={
"manifest_json": json.dumps({**json.loads(original.manifest_json), "author": "legacy author"}),
})
assert (await hub.client.get(f"{HUB}/{NAME}", headers=_headers(MASTER_KEY))).status_code == 200
monkeypatch.delenv("PROXY_BASE_URL")
hub.client.base_url = httpx.URL("https://public.example")
first: Final = hub.database.litellm_claudecodepluginversion.create.await_args.kwargs["data"]
hub.database.litellm_claudecodepluginversion.count.return_value = 1
async def replace_row(
where: Mapping[str, str], data: Mapping[str, object],
) -> prisma_models.LiteLLM_ClaudeCodePluginTable:
current: Final = original.model_copy(update=dict(data))
hub.plugins.find_unique.return_value = current
return current
hub.plugins.update.side_effect = replace_row
updated: Final = await hub.client.put(
f"{HUB}/{NAME}", headers=_headers(MASTER_KEY),
json={"description": "Registry description", "domain": "notes", "namespace": "team", "keywords": ["notes"]},
)
assert updated.status_code == 200, updated.text
assert updated.json()["version"] == "1.0.1"
assert updated.json()["download_url"] == f"https://public.example{HUB}/{NAME}/archive"
latest: Final = hub.database.litellm_claudecodepluginversion.create.await_args.kwargs["data"]
manifest: Final = json.loads(latest["manifest_json"])
assert latest["version"] == manifest["version"] == updated.json()["version"]
assert latest["manifest_json"] == hub.plugins.find_unique.return_value.manifest_json
assert manifest["agent_skill"] == json.loads(original.manifest_json)["agent_skill"]
assert (manifest["domain"], manifest["namespace"], manifest["keywords"]) == ("notes", "team", ["notes"])
assert first["version"] == "1.0.0"
assert latest["content"].decode() == first["content"].decode()
_assert_portable(await hub.client.get(updated.json()["download_url"]))
@pytest.mark.parametrize("key,status", ((MASTER_KEY, 200), (GRANTED_KEY, 200), (PLAIN_KEY, 404)))
@pytest.mark.asyncio
async def test_neutral_history_list_detail_and_portable_download_share_current_visibility(
hub: HubClient, key: str, status: int,
) -> None:
assert (await _upload(hub, _zip(FILES), version=None)).status_code == 200
original: Final = _stored_upload(hub, enabled=False)
data: Final = hub.database.litellm_claudecodepluginversion.create.await_args.kwargs["data"]
version: Final = prisma_models.LiteLLM_ClaudeCodePluginVersion.model_validate({"id": "retained", **data})
hub.plugins.find_unique.return_value = original.model_copy(update={"version": "3.0.0"})
hub.database.litellm_claudecodepluginversion.find_first.return_value = version
hub.database.query_raw.return_value = [{
"id": version.id, "version": version.version, "description": version.description,
"manifest": json.loads(version.manifest_json), "published_at": version.published_at,
"published_by": version.published_by, "archive_available": version.archive_available,
}]
listed: Final = await hub.client.get(f"{HUB}/{NAME}/versions", headers=_headers(key))
detail: Final = await hub.client.get(f"{HUB}/{NAME}/versions/retained", headers=_headers(key))
downloaded: Final = await hub.client.get(f"{HUB}/{NAME}/versions/retained/archive", headers=_headers(key))
assert listed.status_code == detail.status_code == downloaded.status_code == status
assert hub.downloads.create.await_count == int(status == 200)
if status == 200:
event: Final = hub.downloads.create.await_args.kwargs["data"]
assert event["version"] == version.version != "3.0.0"
assert event["digest"] == hashlib.sha256(downloaded.content).hexdigest()
assert listed.json() == {"versions": [detail.json()], "next_offset": None}
assert detail.json()["manifest"]["agent_skill"] == json.loads(original.manifest_json)["agent_skill"]
assert detail.json()["download_url"] == f"http://test{HUB}/{NAME}/versions/retained/archive"
_assert_portable(downloaded)
assert downloaded.headers["etag"] == f'"sha256:{hashlib.sha256(downloaded.content).hexdigest()}"'
hub.database.query_raw.return_value = []
assert (await hub.client.get(f"{HUB}/{NAME}/versions/missing", headers=_headers(key))).status_code == 404
hub.database.litellm_claudecodepluginversion.find_first.return_value = None
assert (await hub.client.get(
f"{HUB}/{NAME}/versions/missing/archive", headers=_headers(key),
)).status_code == 404
@pytest.mark.asyncio
async def test_neutral_history_labels_external_sources_and_bounds_pagination(hub: HubClient) -> None:
assert (await _upload(hub, _zip(FILES))).status_code == 200
_stored_upload(hub)
row: Final = {
"id": "reference", "version": "3.0.0", "description": None,
"manifest": {"source": {"source": "github", "repo": "example/skill"}},
"published_at": None, "published_by": None, "archive_available": False,
}
hub.database.query_raw.return_value = [row, {**row, "id": "older"}]
listed: Final = await hub.client.get(f"{HUB}/{NAME}/versions?limit=1", headers=_headers(MASTER_KEY))
assert listed.status_code == 200, listed.text
assert listed.json()["next_offset"] == 1
assert len(listed.json()["versions"]) == 1
assert listed.json()["versions"][0]["content_kind"] == "source_reference"
assert listed.json()["versions"][0]["download_url"] is None
assert (await hub.client.get(
f"{HUB}/{NAME}/versions?limit=101", headers=_headers(MASTER_KEY),
)).status_code == 422
@pytest.fixture
async def pending_hub(hub: HubClient) -> HubClient:
hub.database.query_raw.side_effect = hub.database.litellm_skillsubmission.query_raw
return hub
@pytest.mark.parametrize("operation,approve", (("upload", True), ("import", True), ("upload", False)),
ids=("approve-upload", "approve-import", "reject-upload"))
async def test_member_archive_stays_private_until_admin_review(
pending_hub: HubClient, operation: str, approve: bool,
) -> None:
hub: Final = pending_hub
hub.upstream.return_value = httpx.Response(200, content=_zip(FILES, "repository/"))
pending: Final = (
await _upload(hub, _zip(FILES), MEMBER_KEY, version=None) if operation == "upload"
else await hub.client.post(
f"{HUB}/import", headers=_headers(MEMBER_KEY),
json={"url": f"https://skills.sh/synthetic-author/notes-library/{NAME}"},
)
)
assert pending.status_code == 200, pending.text
receipt: Final = pending.json()
assert receipt["status"] == "pending_review" and receipt["proposal"]["version"] is None
assert receipt["proposal"]["agent_skill"]["name"] == NAME
hub.plugins.create.assert_not_awaited()
hub.database.litellm_claudecodepluginversion.create.assert_not_awaited()
path: Final = f"/skill-hub/submissions/{receipt['submission_id']}"
assert receipt["proposal"]["source"]["url"] == f"http://test{path}/archive"
own: Final = await hub.client.get(path, headers=_headers(MEMBER_KEY))
assert own.json()["current_skill"] is None
assert own.json()["submission"]["proposal"] == receipt["proposal"]
for suffix in ("", "/archive"):
assert (await hub.client.get(path + suffix, headers=_headers(OTHER_KEY))).status_code == 404
_assert_portable(await hub.client.get(path + "/archive", headers=_headers(MEMBER_KEY)))
assert (await hub.client.get(path, headers=_headers(VIEWER_KEY))).status_code == 200
denied: Final = await hub.client.put(path + "/approve", headers=_headers(VIEWER_KEY))
assert denied.status_code == 403, denied.text
assert (await hub.client.get(f"{HUB}/{NAME}/archive")).status_code == 404
if not approve:
rejected: Final = await hub.client.put(
path + "/reject", headers=_headers(MASTER_KEY), json={"reason": "Needs revision"},
)
assert rejected.status_code == 200, rejected.text
assert rejected.json()["skill"] is None and rejected.json()["action"] == "rejected"
assert rejected.json()["submission"]["rejection_reason"] == "Needs revision"
assert rejected.json()["submission"]["proposal"] == receipt["proposal"]
hub.plugins.create.assert_not_awaited()
return
approved: Final = await hub.client.put(path + "/approve", headers=_headers(MASTER_KEY))
assert approved.status_code == 200, approved.text
assert approved.json()["action"] == "approved"
assert approved.json()["skill"]["version"] == "1.0.0"
assert approved.json()["submission"]["proposal"] == receipt["proposal"]
_stored_upload(hub)
_assert_portable(await hub.client.get(approved.json()["skill"]["download_url"]))
assert (await hub.client.put(path + "/reject", headers=_headers(MASTER_KEY))).status_code == 409
listed: Final = await hub.client.get("/skill-hub/submissions?status=approved", headers=_headers(MEMBER_KEY))
assert (listed.json()["count"], listed.json()["has_more"]) == (1, False)
@pytest.mark.parametrize("replacement", (False, True), ids=("metadata", "archive"))
async def test_member_updates_preserve_manifest_and_original_editor_base(pending_hub: HubClient, replacement: bool) -> None:
hub: Final = pending_hub
assert (await _upload(hub, _zip(FILES))).status_code == 200
original: Final = _stored_upload(hub)
manifest: Final = {
**json.loads(original.manifest_json), "domain": "Tools", "namespace": "research",
"keywords": ["notes"], "extension": {"retained": True}, "author": "legacy author",
}
hub.plugins.find_unique.return_value = original.model_copy(update={"manifest_json": json.dumps(manifest)})
before: Final = hub.database.litellm_claudecodepluginversion.create.await_count
updated_files: Final = {**FILES, "references/guide.md": b"Updated supporting instructions.\n"}
expected: Final = original.updated_at.isoformat()
response: Final = (
await hub.client.put(
f"{HUB}/{NAME}/upload", headers=_headers(MEMBER_KEY),
data={"expected_updated_at": expected}, files={"file": ("new.zip", _zip(updated_files), "application/zip")},
) if replacement else await hub.client.put(
f"{HUB}/{NAME}", headers=_headers(MEMBER_KEY),
json={"description": "Reviewed description", "expected_updated_at": expected,
"domain": "Tools", "namespace": "research", "keywords": ["notes"]},
)
)
assert response.status_code == 200, response.text
path: Final = f"/skill-hub/submissions/{response.json()['submission_id']}"
detail: Final = (await hub.client.get(path, headers=_headers(MEMBER_KEY))).json()
assert detail["submission"]["base_updated_at"] == original.updated_at.isoformat().replace("+00:00", "Z")
assert detail["submission"]["proposal"]["domain"] == "Tools"
assert detail["submission"]["proposal"]["keywords"] == ["notes"]
assert detail["current_skill"]["namespace"] == "research"
assert hub.database.litellm_claudecodepluginversion.create.await_count == before
hub.plugins.update.assert_not_awaited()
approved: Final = await hub.client.put(path + "/approve", headers=_headers(MASTER_KEY))
assert approved.status_code == 200, approved.text
stored: Final = json.loads(hub.plugins.find_unique.return_value.manifest_json)
assert (stored["extension"], stored["author"]) == (manifest["extension"], manifest["author"])
assert stored["agent_skill"]["digest"] == response.json()["proposal"]["agent_skill"]["digest"]
assert approved.json()["skill"]["version"] == "2.0.1"
assert approved.json()["submission"]["proposal"]["version"] is None
wrapper: Final = hub.archives.upsert.await_args.kwargs["data"]["update"]["content"].decode()
with zipfile.ZipFile(io.BytesIO(wrapper)) as archive:
assert archive.read(f"skills/{NAME}/references/guide.md") == (
updated_files["references/guide.md"] if replacement else FILES["references/guide.md"]
)
if not replacement:
assert stored["agent_skill"] == manifest["agent_skill"]
@pytest.mark.parametrize("failure,status", (("stale", 409), ("wrong-name", 422), ("viewer", 403)),
ids=("stale", "wrong-name", "viewer"))
async def test_replacement_rejections_do_not_stage_or_change_files(pending_hub: HubClient, failure: str, status: int) -> None:
hub: Final = pending_hub
assert (await _upload(hub, _zip(FILES))).status_code == 200
_stored_upload(hub)
response: Final = await hub.client.put(
f"{HUB}/{NAME}/upload", headers=_headers(VIEWER_KEY if failure == "viewer" else MEMBER_KEY),
data={"expected_updated_at": "2020-01-01T00:00:00Z"} if failure == "stale" else {},
files={"file": ("new.zip", _zip({"SKILL.md": MANIFEST.replace(NAME.encode(), b"different-name")}), "application/zip")},
)
assert response.status_code == status, response.text
assert hub.database.litellm_skillsubmission.rows == ()
hub.plugins.update.assert_not_awaited()
hub.archives.upsert.assert_not_awaited()
@pytest.mark.parametrize("key", (MEMBER_KEY, OTHER_KEY, VIEWER_KEY))
@pytest.mark.parametrize("method", ("GET", "POST"))
async def test_skill_scan_requires_full_admin(hub: HubClient, key: str, method: str) -> None:
response = await hub.client.request(method, "/skill-hub/submissions/absent/scan", headers=_headers(key))
assert response.status_code == 403, response.text
async def test_unconfigured_scanning_preserves_manual_approval(pending_hub: HubClient) -> None:
hub = pending_hub
created = await _upload(hub, _zip(FILES), MEMBER_KEY, version=None)
path = f"/skill-hub/submissions/{created.json()['submission_id']}"
state = await hub.client.get(path + "/scan", headers=_headers(MASTER_KEY))
assert state.status_code == 200, state.text
assert state.json() == {"configured": False, "scan": None}
unavailable = await hub.client.post(path + "/scan", headers=_headers(MASTER_KEY))
assert unavailable.status_code == 503, unavailable.text
approved = await hub.client.put(path + "/approve", headers=_headers(MASTER_KEY))
assert approved.status_code == 200, approved.text
assert approved.json()["submission"]["status"] == "approved"
@pytest.mark.parametrize("scanner_status", (200, 503), ids=("report", "error"))
async def test_configured_scan_persists_result_without_changing_manual_approval(
pending_hub: HubClient, monkeypatch: pytest.MonkeyPatch, scanner_status: int,
) -> None:
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy.discovery_endpoints import skill_scanning
hub: Final = pending_hub
report: Final = {
"scanner_version": "test-version", "complete": True, "scan_mode": "static-only",
"risk_score": 0, "findings": [], "limitations": [],
}
def respond(request: httpx.Request) -> httpx.Response:
assert request.headers["authorization"] == "Bearer scanner-only-token"
return httpx.Response(scanner_status, json=report)
handler: Final = AsyncHTTPHandler(transport=httpx.MockTransport(respond))
def scanner_client(provider: str, params: Mapping[str, object]) -> AsyncHTTPHandler:
return handler
monkeypatch.setattr(skill_scanning, "get_async_httpx_client", scanner_client)
monkeypatch.setattr(proxy_server, "general_settings", {"skill_scanner": {
"endpoint": "https://scanner.example/scan", "api_key_env": "SCANNER_TEST_KEY",
}})
monkeypatch.setenv("SCANNER_TEST_KEY", "scanner-only-token")
created: Final = await _upload(hub, _zip(FILES), MEMBER_KEY, version=None)
path: Final = f"/skill-hub/submissions/{created.json()['submission_id']}"
before: Final = (await hub.client.get(path, headers=_headers(MASTER_KEY))).json()
async with handler.client:
started: Final = await hub.client.post(path + "/scan", headers=_headers(MASTER_KEY))
assert started.status_code == 200, started.text
completed: Final = (await hub.client.get(path + "/scan", headers=_headers(MASTER_KEY))).json()["scan"]
assert completed["id"] == started.json()["scan"]["id"]
expected: Final = ("completed", report) if scanner_status == 200 else ("failed", None)
assert (completed["status"], completed["report"]) == expected
assert completed["error"] is None if scanner_status == 200 else "HTTP 503" in completed["error"]
monkeypatch.setattr(proxy_server, "general_settings", {"skill_scanner": {"endpoint": "invalid"}})
retained: Final = await hub.client.get(path + "/scan", headers=_headers(MASTER_KEY))
assert retained.status_code == 200 and retained.json() == {"configured": False, "scan": completed}
unavailable: Final = await hub.client.post(path + "/scan", headers=_headers(MASTER_KEY))
assert unavailable.status_code == 503 and "configuration is invalid" in unavailable.text
assert (await hub.client.get(path, headers=_headers(MASTER_KEY))).json() == before
approved: Final = await hub.client.put(path + "/approve", headers=_headers(MASTER_KEY))
assert approved.status_code == 200 and approved.json()["submission"]["status"] == "approved"

189
tests/unit/proxy/discovery_endpoints/test_skill_import.py Normal file
View file

@ -0,0 +1,189 @@
import hashlib
import io
import stat
import struct
import zipfile
import zlib
from collections.abc import AsyncIterator
from dataclasses import dataclass
from typing import Final, NoReturn
import httpx
import pytest
from litellm.proxy.discovery_endpoints.agent_skills_archive import (
MAX_ARCHIVE_ENTRIES,
MAX_ARCHIVE_UNPACKED_BYTES,
MAX_SKILL_MANIFEST_BYTES,
)
from litellm.proxy.discovery_endpoints.skill_import import (
MAX_IMPORT_BYTES,
GitHubSkillSource,
SkillImportError,
fetch_skill_source,
parse_skill_source,
select_imported_skill,
)
MANIFEST: Final = b"---\nname: pdf-summarizer\ndescription: Summarize PDFs.\n---\n\nRead the PDF.\n"
SOURCE: Final = GitHubSkillSource("author", "repo", "HEAD", "pdf-summarizer", None)
NESTED: Final = "skills/pdf-summarizer"
def repository_zip(entries: tuple[tuple[str | zipfile.ZipInfo, bytes], ...]) -> bytes:
buffer: Final = io.BytesIO()
with zipfile.ZipFile(buffer, "w") as archive:
for name, content in entries:
entry: Final = zipfile.ZipInfo(name) if isinstance(name, str) else name
archive.writestr(entry, content, compress_type=zipfile.ZIP_DEFLATED)
return buffer.getvalue()
@pytest.mark.parametrize(("url", "skill", "ref", "path", "expected"), (
("https://skills.sh/author/repo/pdf-summarizer", None, "HEAD", None, SOURCE),
("https://github.com/author/repo.git/", "pdf-summarizer", "release/v2", NESTED,
GitHubSkillSource("author", "repo", "release/v2", "pdf-summarizer", NESTED)),
(f"https://github.com/author/repo/tree/release%2Fv2/{NESTED}", None, "HEAD", None,
GitHubSkillSource("author", "repo", "release/v2", None, NESTED)),
("http://github.com/author/repo", None, "HEAD", None, 400),
("https://github.com.example/author/repo", None, "HEAD", None, 400),
("https://user:secret@github.com/author/repo", None, "HEAD", None, 400),
("https://github.com/author/repo?token=secret", None, "HEAD", None, 400),
("https://github.com/author/repo#fragment", None, "HEAD", None, 400),
("https://github.com/author%2Fother/repo", None, "HEAD", None, 400),
("https://[invalid/author/repo", None, "HEAD", None, 400),
("https://skills.sh/author/repo", None, "HEAD", None, 400),
("https://skills.sh/author/repo/pdf-summarizer", "other", "HEAD", None, 400),
(f"https://github.com/author/repo/tree/main/{NESTED}", None, "other", None, 400),
(f"https://github.com/author/repo/tree/main/{NESTED}", None, "HEAD", "other", 400),
))
def test_source_urls(
url: str, skill: str | None, ref: str, path: str | None, expected: GitHubSkillSource | int
) -> None:
result: Final = parse_skill_source(url, skill=skill, ref=ref, path=path)
assert (result.status if isinstance(result, SkillImportError) else result) == expected
@dataclass(frozen=True, slots=True)
class DownloadStream(httpx.AsyncByteStream):
overflow: bool
async def __aiter__(self) -> AsyncIterator[bytes]:
yield b"a" * MAX_IMPORT_BYTES
if self.overflow:
yield b"!"
pytest.fail("The download kept reading after exceeding the byte limit")
@pytest.mark.asyncio
@pytest.mark.parametrize(("mode", "status"), (
("bounded", None), ("overflow", 413), ("redirect", 422), ("http-error", 422), ("timeout", 502),
))
async def test_public_download_transport(mode: str, status: int | None) -> None:
credentials: Final = {
"Authorization": "Bearer unit-secret",
"Proxy-Authorization": "Basic unit-secret",
"X-API-Key": "unit-secret",
"Cookie": "session=unit-secret",
}
def respond(request: httpx.Request) -> httpx.Response:
assert request.method == "GET"
assert str(request.url) == "https://codeload.github.com/author/repo/zip/release%2Fv2"
assert not frozenset(key.lower() for key in credentials).intersection(request.headers)
if mode == "timeout":
raise httpx.ReadTimeout("simulated timeout", request=request)
return httpx.Response(
307 if mode == "redirect" else 404 if mode == "http-error" else 200,
headers={"Location": "https://untrusted.example/private", "Content-Length": "1"},
stream=DownloadStream(mode == "overflow"),
)
async with httpx.AsyncClient(
transport=httpx.MockTransport(respond), headers=credentials, auth=httpx.BasicAuth("unit-user", "unit-secret"),
follow_redirects=True, trust_env=False,
) as client:
result: Final = await fetch_skill_source(GitHubSkillSource("author", "repo", "release/v2", None, None), client)
if status is None:
assert isinstance(result, bytes)
actual: Final = (len(result), hashlib.sha256(result).hexdigest())
expected: Final = (MAX_IMPORT_BYTES, hashlib.sha256(b"a" * MAX_IMPORT_BYTES).hexdigest())
assert actual == expected
else:
assert isinstance(result, SkillImportError) and result.status == status
@pytest.mark.parametrize(("skill", "path", "root_name", "expected"), (
(None, "", b"root", "root"),
(None, "/", b"root", "root"),
(None, "SKILL.md", b"root", 422),
(None, NESTED, b"root", "pdf-summarizer"),
("pdf-summarizer", None, b"root", "pdf-summarizer"),
("pdf-summarizer", f"/{NESTED}/", b"pdf-summarizer", "pdf-summarizer"),
(None, None, b"root", 422),
("missing", None, b"root", 422),
("pdf-summarizer", "", b"root", 422),
("pdf-summarizer", None, b"pdf-summarizer", 422),
(None, "", None, 422),
))
def test_repository_selection(skill: str | None, path: str | None, root_name: bytes | None, expected: str | int) -> None:
root: Final = (
(("repo/SKILL.md", MANIFEST.replace(b"pdf-summarizer", root_name)),)
if root_name is not None else ()
)
result: Final = select_imported_skill(
repository_zip(root + ((f"repo/{NESTED}/SKILL.md", MANIFEST),)),
GitHubSkillSource("author", "repo", "HEAD", skill, path),
)
assert (result.status if isinstance(result, SkillImportError) else result.declared_name) == expected
def test_corrupt_manifest_compression_returns_an_import_error() -> None:
content: Final = repository_zip(((f"repo/{NESTED}/SKILL.md", MANIFEST),))
with zipfile.ZipFile(io.BytesIO(content)) as archive:
entry: Final = archive.infolist()[0]
filename_size, extra_size = struct.unpack_from("<HH", content, entry.header_offset + 26)
offset: Final = entry.header_offset + 30 + filename_size + extra_size
invalid: Final = content[:offset] + b"\xff" * entry.compress_size + content[offset + entry.compress_size:]
result: Final = select_imported_skill(invalid, SOURCE)
assert isinstance(result, SkillImportError) and result.status == 422
@pytest.mark.parametrize(("violation", "expected"), (
("path", 422), ("inside-link", 422), ("outside-link", "pdf-summarizer"), ("entries", 413), ("unpacked-bytes", 413),
))
def test_repository_validation_scope(violation: str, expected: str | int) -> None:
paths: Final = {"path": "repo/../escape", "inside-link": f"repo/{NESTED}/run.sh"}
entry: Final = zipfile.ZipInfo(paths.get(violation, "repo/CLAUDE.md"))
entry.create_system = 3
entry.external_attr = ((stat.S_IFLNK if violation.endswith("-link") else stat.S_IFREG) | 0o644) << 16
payload: Final = b"x" * (MAX_ARCHIVE_UNPACKED_BYTES if violation == "unpacked-bytes" else 1)
extra: Final = (
tuple((f"repo/unrelated/{index}", b"x") for index in range(MAX_ARCHIVE_ENTRIES))
if violation == "entries" else ((entry, payload),)
)
result: Final = select_imported_skill(
repository_zip(((f"repo/{NESTED}/SKILL.md", MANIFEST),) + extra), SOURCE
)
assert (result.status if isinstance(result, SkillImportError) else result.declared_name) == expected
@pytest.mark.parametrize(("violation", "status"), (("path", 422), ("symlink", 422), ("manifest-bytes", 422)))
def test_candidate_metadata_is_validated_before_manifest_decompression(
violation: str, status: int, monkeypatch: pytest.MonkeyPatch
) -> None:
entry: Final = zipfile.ZipInfo(
"repo/../pdf-summarizer/SKILL.md" if violation == "path" else f"repo/{NESTED}/SKILL.md"
)
entry.create_system = 3
entry.external_attr = ((stat.S_IFLNK if violation == "symlink" else stat.S_IFREG) | 0o644) << 16
content: Final = repository_zip(
((entry, b"x" * (MAX_SKILL_MANIFEST_BYTES + 1) if violation == "manifest-bytes" else MANIFEST),)
)
def reject_decompression(*args: object, **kwargs: object) -> NoReturn:
pytest.fail("Invalid candidate metadata reached the decompressor before rejection")
monkeypatch.setattr(zlib, "decompressobj", reject_decompression)
result: Final = select_imported_skill(content, SOURCE)
assert isinstance(result, SkillImportError) and result.status == status

Some files were not shown because too many files have changed in this diff Show more