scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 154

fix(ci): fix remaining CI failures - migration, lint, tests

Browse source 
- Create Prisma migration for BYOK MCP fields, MCPUserCredentials,
  JWTKeyMapping models, and re-add spec_path column
- Fix syntax error in test_byok_oauth_endpoints.py (duplicate line)
- Remove unused get_user_credential import (Ruff F401)
- Extract _resolve_byok_auth helper to fix PLR0915 (too many statements)
- Add BYOK mock fields to MCP server tests (Pydantic validation)
- Fix JWT handler test by initializing litellm_jwtauth on jwt_handler

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Harshit28j 2026-03-06 09:39:14 +05:30
parent c453595e92
commit 5a1578b1de
6 changed files with 115 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
litellm-proxy-extras/litellm_proxy_extras/migrations/20260306000000_add_byok_mcp_jwt_key_mapping/migration.sql Normal file
View file

@ -0,0 +1,49 @@
-- Re-add spec_path (was added in 20260220, dropped in 20260224, re-added in schema)
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "spec_path" TEXT;
-- Add BYOK MCP fields
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "tool_name_to_display_name" JSONB DEFAULT '{}';
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "tool_name_to_description" JSONB DEFAULT '{}';
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "is_byok" BOOLEAN NOT NULL DEFAULT false;
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "byok_description" TEXT[] DEFAULT ARRAY[]::TEXT[];
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "byok_api_key_help_url" TEXT;
-- CreateTable
CREATE TABLE "LiteLLM_MCPUserCredentials" (
"id" TEXT NOT NULL,
"user_id" TEXT NOT NULL,
"server_id" TEXT NOT NULL,
"credential_b64" TEXT NOT NULL,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT "LiteLLM_MCPUserCredentials_pkey" PRIMARY KEY ("id")
);
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_MCPUserCredentials_user_id_server_id_key" ON "LiteLLM_MCPUserCredentials"("user_id", "server_id");
-- CreateTable
CREATE TABLE "LiteLLM_JWTKeyMapping" (
"id" TEXT NOT NULL,
"jwt_claim_name" TEXT NOT NULL,
"jwt_claim_value" TEXT NOT NULL,
"token" TEXT NOT NULL,
"description" TEXT,
"is_active" BOOLEAN NOT NULL DEFAULT true,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"created_by" TEXT,
"updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_by" TEXT,
CONSTRAINT "LiteLLM_JWTKeyMapping_pkey" PRIMARY KEY ("id")
);
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_key" ON "LiteLLM_JWTKeyMapping"("jwt_claim_name", "jwt_claim_value");
-- CreateIndex
CREATE INDEX "LiteLLM_JWTKeyMapping_jwt_claim_name_jwt_claim_value_is_act_idx" ON "LiteLLM_JWTKeyMapping"("jwt_claim_name", "jwt_claim_value", "is_active");
-- AddForeignKey
ALTER TABLE "LiteLLM_JWTKeyMapping" ADD CONSTRAINT "LiteLLM_JWTKeyMapping_token_fkey" FOREIGN KEY ("token") REFERENCES "LiteLLM_VerificationToken"("token") ON DELETE RESTRICT ON UPDATE CASCADE;

59
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -1644,6 +1644,37 @@ if MCP_AVAILABLE:
},
)
async def _resolve_byok_auth(
mcp_server: MCPServer,
user_api_key_auth: Optional[UserAPIKeyAuth],
mcp_auth_header: Optional[str],
) -> Optional[str]:
"""Resolve BYOK credential for a server, returning the auth header to use."""
if not mcp_server.is_byok:
return mcp_auth_header
if not mcp_auth_header:
byok_cred = await _get_byok_credential(mcp_server, user_api_key_auth)
if byok_cred is None:
raise HTTPException(
status_code=401,
detail={
"error": "byok_auth_required",
"server_id": mcp_server.server_id,
"server_name": mcp_server.server_name or mcp_server.name,
"message": (
"No stored credential found for this BYOK server. "
"Complete the OAuth authorization flow to provide your API key."
),
},
headers={
"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'
},
)
return byok_cred
# External auth header supplied; still enforce user-identity check.
await _check_byok_credential(mcp_server, user_api_key_auth)
return mcp_auth_header
async def execute_mcp_tool(
name: str,
arguments: Dict[str, Any],
@ -1733,30 +1764,10 @@ if MCP_AVAILABLE:
"mcp_tool_call_metadata"
] = standard_logging_mcp_tool_call
# BYOK: retrieve the stored per-user credential. A single DB call
# both checks existence and fetches the value, avoiding a double query.
if mcp_server.is_byok and not mcp_auth_header:
byok_cred = await _get_byok_credential(mcp_server, user_api_key_auth)
if byok_cred is None:
raise HTTPException(
status_code=401,
detail={
"error": "byok_auth_required",
"server_id": mcp_server.server_id,
"server_name": mcp_server.server_name or mcp_server.name,
"message": (
"No stored credential found for this BYOK server. "
"Complete the OAuth authorization flow to provide your API key."
),
},
headers={
"WWW-Authenticate": 'Bearer resource_metadata="/.well-known/oauth-protected-resource"'
},
)
mcp_auth_header = byok_cred
elif mcp_server.is_byok:
# External auth header supplied; still enforce user-identity check.
await _check_byok_credential(mcp_server, user_api_key_auth)
# BYOK: retrieve or validate the per-user credential
mcp_auth_header = await _resolve_byok_auth(
mcp_server, user_api_key_auth, mcp_auth_header
)
# Check if tool exists in local registry first (for OpenAPI-based tools)
# These tools are registered with their prefixed names

1
litellm/proxy/management_endpoints/mcp_management_endpoints.py
View file

@ -81,7 +81,6 @@ if MCP_AVAILABLE:
delete_user_credential,
get_all_mcp_servers_for_user,
get_mcp_server,
get_user_credential,
store_user_credential,
update_mcp_server,
)

18
tests/mcp_tests/test_mcp_server.py
View file

@ -1458,6 +1458,12 @@ async def test_add_update_server_with_alias():
mock_mcp_server.available_on_public_internet = True
mock_mcp_server.created_at = None
mock_mcp_server.updated_at = None
# BYOK fields
mock_mcp_server.tool_name_to_display_name = None
mock_mcp_server.tool_name_to_description = None
mock_mcp_server.is_byok = False
mock_mcp_server.byok_description = []
mock_mcp_server.byok_api_key_help_url = None
# Add server to manager
await test_manager.add_server(mock_mcp_server)
@ -1504,6 +1510,12 @@ async def test_add_update_server_without_alias():
mock_mcp_server.available_on_public_internet = True
mock_mcp_server.created_at = None
mock_mcp_server.updated_at = None
# BYOK fields
mock_mcp_server.tool_name_to_display_name = None
mock_mcp_server.tool_name_to_description = None
mock_mcp_server.is_byok = False
mock_mcp_server.byok_description = []
mock_mcp_server.byok_api_key_help_url = None
# Add server to manager
await test_manager.add_server(mock_mcp_server)
@ -1550,6 +1562,12 @@ async def test_add_update_server_fallback_to_server_id():
mock_mcp_server.available_on_public_internet = True
mock_mcp_server.created_at = None
mock_mcp_server.updated_at = None
# BYOK fields
mock_mcp_server.tool_name_to_display_name = None
mock_mcp_server.tool_name_to_description = None
mock_mcp_server.is_byok = False
mock_mcp_server.byok_description = []
mock_mcp_server.byok_api_key_help_url = None
# Add server to manager
await test_manager.add_server(mock_mcp_server)

13
tests/proxy_unit_tests/test_user_api_key_auth.py
View file

@ -1044,6 +1044,19 @@ async def test_jwt_non_admin_team_route_access(monkeypatch):
litellm.proxy.proxy_server, "general_settings", {"enable_jwt_auth": True}
)
# Initialize jwt_handler with litellm_jwtauth so user_api_key_auth can
# access jwt_handler.litellm_jwtauth before auth_builder is called
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.caching import DualCache
jwt_handler = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=DualCache(),
litellm_jwtauth=LiteLLM_JWTAuth(),
)
monkeypatch.setattr(litellm.proxy.proxy_server, "jwt_handler", jwt_handler)
# Mock enterprise license check and JWTAuthManager.auth_builder
# License check must be mocked to avoid environment variable pollution
# in parallel test execution

2
tests/test_litellm/proxy/_experimental/mcp_server/test_byok_oauth_endpoints.py
View file

@ -477,7 +477,6 @@ async def test_check_byok_credential_missing_credential():
with patch(
"litellm.proxy._experimental.mcp_server.db.get_user_credential",
new=AsyncMock(return_value=None),
), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
with pytest.raises(HTTPException) as exc_info:
await _check_byok_credential(server, user_auth)
@ -511,7 +510,6 @@ async def test_check_byok_credential_has_credential():
with patch(
"litellm.proxy._experimental.mcp_server.db.get_user_credential",
new=AsyncMock(return_value="some-credential-value"),
), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
), patch("litellm.proxy.proxy_server.prisma_client", mock_prisma):
# Should not raise
await _check_byok_credential(server, user_auth)