litellm/CLAUDE.md
Ishaan Jaff 1bb713bc7b
feat(mcp): BYOK MCP servers with OAuth 2.1 PKCE authorization flow (#22850)
* feat(mcp): BYOK (Bring Your Own Key) for OpenAPI MCP servers with OAuth 2.1 flow

Adds per-user credential storage for BYOK MCP servers so external clients
can authenticate via standard OAuth 2.1 PKCE without needing a full identity
provider.

Backend:
- New DB table LiteLLM_MCPUserCredentials (user_id, server_id, credential_b64)
- is_byok, byok_description, byok_api_key_help_url fields on MCPServerTable
- OAuth 2.1 authorization server endpoints (/.well-known/oauth-authorization-server,
  /.well-known/oauth-protected-resource, /v1/mcp/oauth/authorize, /v1/mcp/oauth/token)
- 401 challenge with WWW-Authenticate header when BYOK server has no credential
- CRUD endpoints: POST/DELETE /v1/mcp/server/{id}/user-credential
- has_user_credential annotated on GET /v1/mcp/server response

UI:
- ByokCredentialModal: 2-step Connect flow (access description + API key entry)
- BYOK toggle + description fields on admin MCP server create form
- Connect/Connected state in MCP server table
- BYOK Demo page (/tools/byok-demo) showing full OAuth 2.1 PKCE flow

* feat(mcp/byok): redesign OAuth authorize page to match 2-step Connect mockup

- Step 1: L→S logos, requested access checklist, How it works box, Continue button
- Step 2: API key input, Save toggle, Duration pills (1h/24h/7d/30d/until_revoked), security note
- Matches screenshots: white modal on dark bg, progress dots, dark CTA buttons
- Authorize handler now fetches byok_description and byok_api_key_help_url from server registry
- CLAUDE.md: replace SQL snippet with proper DB migration troubleshooting guidance

* fix: address greptile review feedback (greploop iteration 1)

- XSS: escape all user-supplied values in _build_authorize_html() with html.escape()
- Open redirect: validate redirect_uri scheme and URL-encode code/state in redirect
- N+1 query: batch BYOK credential lookup into single find_many() call
- Critical path DB: add 60s TTL in-memory cache to _check_byok_credential()
- Encrypt BYOK credentials at rest using encrypt_value_helper/decrypt_value_helper

* fix(byok): update OAuth popup with LiteLLM logo, MCP title suffix, remove emojis

* fix(byok-demo): fix token endpoint URL (/v1/mcp/oauth/token not /v1/mcp/token)

* feat(byok): inject stored BYOK credential as mcp_auth_header on tool execution

* feat(byok): use contextvars to inject per-user credential into OpenAPI tool closures; remove byok-demo from LiteLLM UI

OpenAPI tools have auth headers baked into their closures at registration time. BYOK servers have
no static auth token, so per-user credentials were never reaching the HTTP calls.

Fix: add _request_auth_header ContextVar in openapi_to_mcp_generator.py. create_tool_function now
reads this var at call time and overrides the Authorization header if set. execute_mcp_tool resolves
the MCP server and performs BYOK checks before the local-tool dispatch branch, then sets the
ContextVar around _handle_local_mcp_tool so the credential flows into the HTTP request.

Also remove the /tools/byok-demo page from the LiteLLM UI dashboard — the demo lives at
~/Downloads/litellm-byok-demo/index.html (served separately on port 8080).

* fix: address greptile review feedback (greploop iteration 2)

- Cache invalidation: add _invalidate_byok_cred_cache() and call it after
  store_user_credential() in both token endpoint and management endpoint
- Unbounded cache: add _BYOK_CRED_CACHE_MAX_SIZE=4096 with clear-on-overflow
- Unbounded auth codes: add _AUTH_CODES_MAX_SIZE=1000 with 503 on overflow
- Double DB query: merge _check_byok_credential + _get_byok_credential into
  single _get_byok_credential call; raise 401 inline if None returned
- Sidebar: remove byok-demo entry (page was deleted in prior commit)
- JWT comment: document why byok_session HS256 token can't be used as proxy auth

* fix: address greptile review feedback (greploop iteration 3)

- auth_type: pre-format Authorization header (Bearer/ApiKey/Basic) in server.py
  before setting ContextVar so openapi_to_mcp_generator respects server auth_type
- cache invalidation on delete: call _invalidate_byok_cred_cache after
  delete_user_credential so stale True entries don't persist for 60s
- ContextVar guard: only set _request_auth_header when mcp_auth_header is set,
  avoiding unnecessary ContextVar overhead on non-BYOK tool calls

* fix: address greptile review feedback (greploop iteration 4)

- Unified credential cache: store actual credential value (Optional[str])
  instead of just bool so _get_byok_credential also benefits from caching —
  eliminates the DB hit on every BYOK tool call within the 60s TTL window
- Extracted _write_byok_cred_cache() helper for consistent cache writes
- Replaced has_user_credential with get_user_credential in _check_byok_credential
  so one DB call satisfies both existence check and value retrieval
- Remove false 'encrypted at rest' claim from OAuth HTML and ByokCredentialModal

* Update tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

* Update tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
2026-03-04 21:19:25 -08:00

6.4 KiB

CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

Development Commands

Installation

  • make install-dev - Install core development dependencies
  • make install-proxy-dev - Install proxy development dependencies with full feature set
  • make install-test-deps - Install all test dependencies

Testing

  • make test - Run all tests
  • make test-unit - Run unit tests (tests/test_litellm) with 4 parallel workers
  • make test-integration - Run integration tests (excludes unit tests)
  • pytest tests/ - Direct pytest execution

Code Quality

  • make lint - Run all linting (Ruff, MyPy, Black, circular imports, import safety)
  • make format - Apply Black code formatting
  • make lint-ruff - Run Ruff linting only
  • make lint-mypy - Run MyPy type checking only

Single Test Files

  • poetry run pytest tests/path/to/test_file.py -v - Run specific test file
  • poetry run pytest tests/path/to/test_file.py::test_function -v - Run specific test

Running Scripts

  • poetry run python script.py - Run Python scripts (use for non-test files)

GitHub Issue & PR Templates

When contributing to the project, use the appropriate templates:

Bug Reports (.github/ISSUE_TEMPLATE/bug_report.yml):

  • Describe what happened vs. what you expected
  • Include relevant log output
  • Specify your LiteLLM version

Feature Requests (.github/ISSUE_TEMPLATE/feature_request.yml):

  • Describe the feature clearly
  • Explain the motivation and use case

Pull Requests (.github/pull_request_template.md):

  • Add at least 1 test in tests/litellm/
  • Ensure make test-unit passes

Architecture Overview

LiteLLM is a unified interface for 100+ LLM providers with two main components:

Core Library (litellm/)

  • Main entry point: litellm/main.py - Contains core completion() function
  • Provider implementations: litellm/llms/ - Each provider has its own subdirectory
  • Router system: litellm/router.py + litellm/router_utils/ - Load balancing and fallback logic
  • Type definitions: litellm/types/ - Pydantic models and type hints
  • Integrations: litellm/integrations/ - Third-party observability, caching, logging
  • Caching: litellm/caching/ - Multiple cache backends (Redis, in-memory, S3, etc.)

Proxy Server (litellm/proxy/)

  • Main server: proxy_server.py - FastAPI application
  • Authentication: auth/ - API key management, JWT, OAuth2
  • Database: db/ - Prisma ORM with PostgreSQL/SQLite support
  • Management endpoints: management_endpoints/ - Admin APIs for keys, teams, models
  • Pass-through endpoints: pass_through_endpoints/ - Provider-specific API forwarding
  • Guardrails: guardrails/ - Safety and content filtering hooks
  • UI Dashboard: Served from _experimental/out/ (Next.js build)

Key Patterns

Provider Implementation

  • Providers inherit from base classes in litellm/llms/base.py
  • Each provider has transformation functions for input/output formatting
  • Support both sync and async operations
  • Handle streaming responses and function calling

Error Handling

  • Provider-specific exceptions mapped to OpenAI-compatible errors
  • Fallback logic handled by Router system
  • Comprehensive logging through litellm/_logging.py

Configuration

  • YAML config files for proxy server (see proxy/example_config_yaml/)
  • Environment variables for API keys and settings
  • Database schema managed via Prisma (proxy/schema.prisma)

Development Notes

Code Style

  • Uses Black formatter, Ruff linter, MyPy type checker
  • Pydantic v2 for data validation
  • Async/await patterns throughout
  • Type hints required for all public APIs
  • Avoid imports within methods — place all imports at the top of the file (module-level). Inline imports inside functions/methods make dependencies harder to trace and hurt readability. The only exception is avoiding circular imports where absolutely necessary.

Testing Strategy

  • Unit tests in tests/test_litellm/
  • Integration tests for each provider in tests/llm_translation/
  • Proxy tests in tests/proxy_unit_tests/
  • Load tests in tests/load_tests/
  • Always add tests when adding new entity types or features — if the existing test file covers other entity types, add corresponding tests for the new one

UI / Backend Consistency

  • When wiring a new UI entity type to an existing backend endpoint, verify the backend API contract (single value vs. array, required vs. optional params) and ensure the UI controls match — e.g., use a single-select dropdown when the backend accepts a single value, not a multi-select

Database Migrations

  • Prisma handles schema migrations
  • Migration files auto-generated with prisma migrate dev
  • Always test migrations against both PostgreSQL and SQLite

Proxy database access

  • Do not write raw SQL for proxy DB operations. Use Prisma model methods instead of execute_raw / query_raw.
  • Use the generated client: prisma_client.db.<model> (e.g. litellm_tooltable, litellm_usertable) with .upsert(), .find_many(), .find_unique(), .update(), .update_many() as appropriate. This avoids schema/client drift, keeps code testable with simple mocks, and matches patterns used in spend logs and other proxy code.

Enterprise Features

  • Enterprise-specific code in enterprise/ directory
  • Optional features enabled via environment variables
  • Separate licensing and authentication for enterprise features

Troubleshooting: DB schema out of sync after proxy restart

litellm-proxy-extras runs prisma migrate deploy on startup using its own bundled migration files, which may lag behind schema changes in the current worktree. Symptoms: Unknown column, Invalid prisma invocation, or missing data on new fields.

Diagnose: Run \d "TableName" in psql and compare against schema.prisma — missing columns confirm the issue.

Fix options:

  1. Create a Prisma migration (permanent) — run prisma migrate dev --name <description> in the worktree. The generated file will be picked up by prisma migrate deploy on next startup.
  2. Apply manually for local devpsql -d litellm -c "ALTER TABLE ... ADD COLUMN IF NOT EXISTS ..." after each proxy start. Fine for dev, not for production.
  3. Update litellm-proxy-extras — if the package is installed from PyPI, its migration directory must include the new file. Either update the package or run the migration manually until the next release ships it.