mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
188 lines
9 KiB
Python
188 lines
9 KiB
Python
import json
|
|
import queue
|
|
import uuid
|
|
from urllib.parse import parse_qs, urlsplit
|
|
from typing import Final, Literal
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from integration._support.client import Gateway, eventually
|
|
from integration._support.database import read_rows
|
|
from integration._support.mcp import McpPeer, call_tool, mcp_peer, register_mcp, tool_names
|
|
from integration._support.process import owned_proxy
|
|
from integration._support.wire import Reply, Request, wire_server
|
|
|
|
|
|
@pytest.mark.covers("other.mcp.oauth.discovery_cannot_erase_configured_authorization_endpoint")
|
|
def test_partial_discovery_and_unrelated_edit_keep_actual_authorization_destination(
|
|
gateway: Gateway, tmp_path: Path
|
|
) -> None:
|
|
def discovery(request: Request) -> Reply:
|
|
if request.target.startswith("/configured-authorize"):
|
|
return Reply(body=b'{"synthetic_authorization_endpoint":true}')
|
|
if request.target == "/mcp":
|
|
return Reply(body=b'{"synthetic_resource":true}')
|
|
if request.target.startswith("/.well-known/oauth-protected-resource"):
|
|
return Reply(
|
|
body=json.dumps(
|
|
{
|
|
"resource": wire.url + "/mcp",
|
|
"authorization_servers": [wire.url],
|
|
"scopes_supported": ["tools.read"],
|
|
}
|
|
).encode()
|
|
)
|
|
if request.method == "GET":
|
|
return Reply(
|
|
body=json.dumps(
|
|
{
|
|
"issuer": wire.url,
|
|
"token_endpoint": wire.url + "/discovered-token",
|
|
"scopes_supported": ["tools.read"],
|
|
}
|
|
).encode()
|
|
)
|
|
return Reply(status=401, body=b'{"error":"synthetic OAuth requirement"}')
|
|
|
|
with (
|
|
wire_server(discovery) as wire,
|
|
owned_proxy(gateway, tmp_path, {"LITELLM_MCP_OAUTH_DISCOVERY_ON_STARTUP": "true"}) as candidate,
|
|
candidate.scenario() as scenario,
|
|
):
|
|
gateway = candidate
|
|
alias: Final = "integration" + uuid.uuid4().hex
|
|
endpoint: Final = wire.url + "/configured-authorize"
|
|
identity: Final = register_mcp(
|
|
scenario,
|
|
McpPeer(wire.url + "/mcp", queue.Queue()),
|
|
alias,
|
|
auth_type="oauth2",
|
|
authorization_url=endpoint,
|
|
token_url=wire.url + "/configured-token",
|
|
oauth2_flow="authorization_code",
|
|
credentials={"client_id": "synthetic-oauth-client"},
|
|
)
|
|
discovered = []
|
|
|
|
def observed() -> tuple[Request, ...]:
|
|
discovered.extend(wire.drain())
|
|
return tuple(item for item in discovered if item.method == "GET" and ".well-known/" in item.target)
|
|
|
|
assert eventually(observed, bool, seconds=10)
|
|
for generation in range(2):
|
|
rows: Final = read_rows(
|
|
'SELECT authorization_url FROM "LiteLLM_MCPServerTable" WHERE server_id=%s', (identity,)
|
|
)
|
|
assert rows == [{"authorization_url": endpoint}]
|
|
response: Final = gateway.request(
|
|
"GET",
|
|
f"/v1/mcp/server/oauth/{identity}/authorize",
|
|
params={
|
|
"redirect_uri": "http://127.0.0.1:8765/callback",
|
|
"state": "synthetic-state",
|
|
"code_challenge": "A" * 43,
|
|
"code_challenge_method": "S256",
|
|
"response_type": "code",
|
|
},
|
|
)
|
|
assert response.status_code in (302, 307), response.text
|
|
location: Final = urlsplit(response.headers["location"])
|
|
assert location.scheme + "://" + location.netloc + location.path == endpoint
|
|
query: Final = parse_qs(location.query)
|
|
assert query["client_id"] == ["synthetic-oauth-client"]
|
|
assert query["scope"] == ["tools.read"], (
|
|
"Discovery metadata must be applied before checking endpoint preservation"
|
|
)
|
|
assert query["code_challenge"] == ["A" * 43] and query["code_challenge_method"] == ["S256"]
|
|
selected: Final = gateway.client.get(response.headers["location"])
|
|
assert selected.status_code == 200 and selected.json() == {"synthetic_authorization_endpoint": True}
|
|
if generation == 0:
|
|
updated: Final = gateway.request(
|
|
"PUT", "/v1/mcp/server", {"server_id": identity, "server_name": alias + "renamed"}
|
|
)
|
|
assert updated.status_code == 202, updated.text
|
|
|
|
|
|
@pytest.mark.covers("other.mcp.oauth.same_url_credentials_are_isolated_by_user_and_server")
|
|
@pytest.mark.parametrize("transition", ("revoke", "expire"))
|
|
def test_same_url_oauth_credentials_and_revocation_are_isolated_by_user_and_server(
|
|
gateway: Gateway,
|
|
transition: Literal["revoke", "expire"],
|
|
) -> None:
|
|
with mcp_peer() as peer, gateway.scenario() as scenario:
|
|
servers: Final = tuple(
|
|
register_mcp(
|
|
scenario,
|
|
peer,
|
|
"oauth" + uuid.uuid4().hex,
|
|
auth_type="oauth2",
|
|
oauth2_flow="authorization_code",
|
|
authorization_url=peer.url + "/authorize",
|
|
token_url=peer.url + "/token",
|
|
credentials={"client_id": "synthetic-oauth-client"},
|
|
)
|
|
for _ in range(2)
|
|
)
|
|
users: Final = tuple(scenario.user(user_role="internal_user") for _ in range(2))
|
|
keys: Final = tuple(
|
|
scenario.key(user_id=user, object_permission={"mcp_servers": list(servers)}) for user in users
|
|
)
|
|
for user_index, key in enumerate(keys):
|
|
for server_index, server_id in enumerate(servers):
|
|
stored: Final = gateway.request(
|
|
"POST",
|
|
f"/v1/mcp/server/{server_id}/oauth-user-credential",
|
|
{"access_token": f"synthetic-user-{user_index}-server-{server_index}", "expires_in": 3600},
|
|
key=key,
|
|
)
|
|
assert stored.status_code == 200 and stored.json()["has_credential"] is True, stored.text
|
|
scenario.cleanups.callback(
|
|
gateway.request,
|
|
"DELETE",
|
|
f"/v1/mcp/server/{server_id}/oauth-user-credential",
|
|
key=key,
|
|
)
|
|
names: Final = tuple(tool_names(gateway, keys[0], server) for server in servers)
|
|
for generation in range(2):
|
|
for user_index, key in enumerate(keys):
|
|
for server_index, server_id in enumerate(servers):
|
|
peer.drain()
|
|
discovery: Final = gateway.request(
|
|
"GET",
|
|
"/mcp-rest/tools/list",
|
|
key=key,
|
|
params={"server_id": server_id},
|
|
)
|
|
call: Final = call_tool(gateway, key, server_id, names[server_index]["add"], {"a": 3, "b": 5})
|
|
observed: Final = peer.drain()
|
|
if generation == 1 and user_index == 0 and server_index == 0:
|
|
for rejected in (discovery, call):
|
|
assert rejected.status_code == 401, rejected.text
|
|
assert rejected.json() == {"detail": "Unauthorized"}, rejected.text
|
|
assert "resource_metadata=" in rejected.headers["www-authenticate"]
|
|
assert observed == (), "unusable credentials must not fall back to another user or server"
|
|
else:
|
|
assert discovery.status_code == 200, discovery.text
|
|
assert {tool["name"] for tool in discovery.json()["tools"]} == set(names[server_index].values())
|
|
assert call.status_code == 200 and call.json()["isError"] is False, call.text
|
|
assert call.json()["content"][0]["text"] == "8", call.text
|
|
calls: Final = tuple(item for item in observed if item["body"].get("method") == "tools/call")
|
|
assert len(calls) == 1
|
|
expected: Final = f"Bearer synthetic-user-{user_index}-server-{server_index}".encode()
|
|
assert calls[0]["headers"][b"authorization"] == expected
|
|
assert all(item["headers"].get(b"authorization") == expected for item in observed)
|
|
if generation == 0:
|
|
changed: Final = gateway.request(
|
|
"DELETE" if transition == "revoke" else "POST",
|
|
f"/v1/mcp/server/{servers[0]}/oauth-user-credential",
|
|
None
|
|
if transition == "revoke"
|
|
else {
|
|
"access_token": "synthetic-expired-user-0-server-0",
|
|
"expires_in": -60,
|
|
},
|
|
key=keys[0],
|
|
)
|
|
assert changed.status_code == 200, changed.text
|
|
assert changed.json()["has_credential"] is (transition == "expire"), changed.text
|