mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
fc27dfdc7d
* test(integration): give each owned proxy a private database and Redis by default owned_proxy and owned_proxy_process now boot on a scratch database and a per-proxy redis-server unless the caller opts into the worker's state with state=OwnedState.SHARED_WITH_WORKER, passes its own DATABASE_URL, or strips DATABASE_URL. Under the replica profile the private database gets the writer and reader roles so the reader path stays covered. Gateway carries the database_url its proxy writes to, and the Scenario cleanup checks and delete_key_if_present read from it, so they keep checking something on a private database. The 189 test files with owned proxy sites that depend on the worker's database or Redis opt in with a shares_worker_state module marker, which a module-scoped conftest fixture turns into the shared state. The 7 helper modules with such sites opt in per call, since conftest fixtures and other test files reach them. * test(integration): reach session fixtures with the shared marker and keep the worker retry default The module-scoped fixture that applied shares_worker_state ran after session-scoped fixtures, so proxies those booted went private in marked files. The marker is now applied around each marked item's setup, call and teardown, which covers fixtures of every scope. Owned proxies on the worker database inherited the worker gateway's router_settings.num_retries of 0 from it. A private proxy now gets the same default in its config copy unless the config sets num_retries itself. Five providers files read the owned proxy's spend rows or response cache through the worker's database or Redis and opt into sharing.
567 lines
27 KiB
Python
567 lines
27 KiB
Python
import json
|
|
import queue
|
|
import threading
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Final, Literal
|
|
from urllib.parse import parse_qs, urlsplit
|
|
|
|
import httpx
|
|
import pytest
|
|
from integration._support.client import Gateway, Scenario, 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.oauth_server import oauth_server
|
|
from integration._support.process import owned_proxy
|
|
from integration._support.wire import Reply, Request, Wire, wire_server
|
|
from pydantic import TypeAdapter
|
|
|
|
pytestmark: Final = pytest.mark.shares_worker_state
|
|
|
|
_Upstream = Callable[[Request], Reply]
|
|
CLIENT_REDIRECT: Final = "http://127.0.0.1:9/cb"
|
|
|
|
|
|
def _response_object(response: httpx.Response) -> dict[str, object]:
|
|
return TypeAdapter(dict[str, object]).validate_python(response.json())
|
|
|
|
|
|
def _start_callback_flow(
|
|
gateway: Gateway,
|
|
alias: str,
|
|
key: str,
|
|
client_id: str,
|
|
expected_authorization_url: str | None = None,
|
|
) -> tuple[str, str, str]:
|
|
started: Final = gateway.client.get(
|
|
f"/{alias}/authorize",
|
|
params={
|
|
"client_id": client_id,
|
|
"redirect_uri": CLIENT_REDIRECT,
|
|
"response_type": "code",
|
|
"state": "client-state",
|
|
"code_challenge": "A" * 43,
|
|
"code_challenge_method": "S256",
|
|
},
|
|
headers={"x-litellm-api-key": key},
|
|
)
|
|
assert started.status_code in (302, 307), started.text
|
|
location: Final = urlsplit(started.headers["location"])
|
|
if expected_authorization_url is not None:
|
|
assert location.scheme + "://" + location.netloc + location.path == expected_authorization_url, started.text
|
|
relay_state: Final = parse_qs(location.query)["state"][0]
|
|
cookie_name: Final = f"mcp_oauth_state_{relay_state}"
|
|
assert cookie_name in started.cookies, started.headers
|
|
return relay_state, cookie_name, started.cookies[cookie_name]
|
|
|
|
|
|
def _assert_cleared_oauth_state_cookie(response: httpx.Response, cookie_name: str) -> None:
|
|
cleared: Final = response.headers["set-cookie"]
|
|
assert cleared.startswith(f'{cookie_name}="";'), cleared
|
|
assert "Max-Age=0" in cleared, cleared
|
|
assert "Path=/" in cleared, cleared
|
|
assert "HttpOnly" in cleared, cleared
|
|
assert "SameSite=lax" in cleared, cleared
|
|
|
|
|
|
def test_callback_forwards_the_code_only_for_the_sealed_issuer_and_relays_idp_errors(
|
|
gateway: Gateway,
|
|
) -> None:
|
|
with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario:
|
|
alias: Final = "rt4" + uuid.uuid4().hex[:10]
|
|
identity: Final = register_mcp(
|
|
scenario,
|
|
peer,
|
|
alias,
|
|
issuer=auth.issuer,
|
|
authorization_url=auth.issuer + "/authorize",
|
|
token_url=auth.issuer + "/token",
|
|
registration_url=auth.issuer + "/register",
|
|
auth_type="oauth2",
|
|
oauth2_flow="authorization_code",
|
|
credentials={"client_id": "ac-client", "client_secret": "ac-secret", "scopes": ["tools.call"]},
|
|
)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
registered: Final = gateway.client.post(
|
|
f"/{alias}/register", json={"redirect_uris": [CLIENT_REDIRECT], "client_name": "integration"}
|
|
)
|
|
assert registered.status_code in (200, 201), registered.text
|
|
client_id: Final = str(_response_object(registered)["client_id"])
|
|
|
|
wrong_relay, wrong_cookie, wrong_cookie_value = _start_callback_flow(gateway, alias, key, client_id)
|
|
wrong_code: Final = "callback-code-" + uuid.uuid4().hex
|
|
wrong: Final = gateway.client.get(
|
|
"/callback",
|
|
params={"code": wrong_code, "state": wrong_relay, "iss": "http://127.0.0.1:1/other"},
|
|
cookies={wrong_cookie: wrong_cookie_value},
|
|
)
|
|
assert wrong.status_code == 400, wrong.text
|
|
assert wrong.headers["content-type"] == "text/html; charset=utf-8", wrong.headers
|
|
assert wrong.text == (
|
|
"<html><body><h2>Authentication failed</h2><p><strong>Error:</strong> invalid_issuer</p>"
|
|
"<p>This authorization response came from a different identity provider than the one this "
|
|
"MCP server is configured to use.</p>"
|
|
"<p>You can close this window and try again.</p></body></html>"
|
|
), wrong.text
|
|
assert "location" not in wrong.headers, wrong.headers
|
|
assert wrong_code not in wrong.text, wrong.text
|
|
_assert_cleared_oauth_state_cookie(wrong, wrong_cookie)
|
|
|
|
correct_relay, correct_cookie, correct_cookie_value = _start_callback_flow(gateway, alias, key, client_id)
|
|
correct: Final = gateway.client.get(
|
|
"/callback",
|
|
params={"code": "c", "state": correct_relay, "iss": auth.issuer},
|
|
cookies={correct_cookie: correct_cookie_value},
|
|
)
|
|
assert correct.status_code == 302, correct.text
|
|
correct_location: Final = urlsplit(correct.headers["location"])
|
|
assert correct_location.scheme + "://" + correct_location.netloc + correct_location.path == CLIENT_REDIRECT, (
|
|
correct.headers["location"]
|
|
)
|
|
correct_query: Final = parse_qs(correct_location.query)
|
|
assert set(correct_query) == {"code", "state"}, correct.headers["location"]
|
|
assert correct_query["state"] == ["client-state"], correct.headers["location"]
|
|
assert len(correct_query["code"]) == 1 and correct_query["code"][0].startswith("llm_scope_"), correct.headers[
|
|
"location"
|
|
]
|
|
_assert_cleared_oauth_state_cookie(correct, correct_cookie)
|
|
|
|
error_relay, error_cookie, error_cookie_value = _start_callback_flow(gateway, alias, key, client_id)
|
|
error: Final = gateway.client.get(
|
|
"/callback",
|
|
params={"error": "access_denied", "error_description": "x", "state": error_relay},
|
|
cookies={error_cookie: error_cookie_value},
|
|
)
|
|
assert error.status_code == 302, error.text
|
|
error_location: Final = urlsplit(error.headers["location"])
|
|
assert parse_qs(error_location.query) == {
|
|
"error": ["access_denied"],
|
|
"error_description": ["x"],
|
|
"state": ["client-state"],
|
|
}, error.headers["location"]
|
|
_assert_cleared_oauth_state_cookie(error, error_cookie)
|
|
|
|
error_issuer_relay, error_issuer_cookie, error_issuer_cookie_value = _start_callback_flow(
|
|
gateway, alias, key, client_id
|
|
)
|
|
error_issuer: Final = gateway.client.get(
|
|
"/callback",
|
|
params={
|
|
"error": "access_denied",
|
|
"error_description": "x",
|
|
"state": error_issuer_relay,
|
|
"iss": "http://127.0.0.1:1/other",
|
|
},
|
|
cookies={error_issuer_cookie: error_issuer_cookie_value},
|
|
)
|
|
assert error_issuer.status_code == 400, error_issuer.text
|
|
assert error_issuer.headers["content-type"] == "text/html; charset=utf-8", error_issuer.headers
|
|
assert error_issuer.text == (
|
|
"<html><body><h2>Authentication failed</h2><p><strong>Error:</strong> invalid_issuer</p>"
|
|
"<p>Unexpected authorization issuer</p>"
|
|
"<p>You can close this window and try again.</p></body></html>"
|
|
), error_issuer.text
|
|
assert "location" not in error_issuer.headers, error_issuer.headers
|
|
_assert_cleared_oauth_state_cookie(error_issuer, error_issuer_cookie)
|
|
|
|
missing: Final = gateway.client.get("/callback")
|
|
assert missing.status_code == 400, missing.text
|
|
assert missing.headers["content-type"] == "text/html; charset=utf-8", missing.headers
|
|
assert missing.text == (
|
|
"<html><body><h2>Authentication failed</h2><p><strong>Error:</strong> invalid_request</p>"
|
|
"<p>Missing authorization 'code' and 'state' parameter(s).</p>"
|
|
"<p>You can close this window and try again.</p></body></html>"
|
|
), missing.text
|
|
|
|
|
|
def test_callback_refuses_a_response_without_iss_when_the_authorization_server_advertises_iss_support(
|
|
gateway: Gateway,
|
|
) -> None:
|
|
wire_holder: Final[list[Wire]] = []
|
|
|
|
def respond(request: Request) -> Reply:
|
|
path: Final = urlsplit(request.target).path
|
|
if request.method == "GET" and path in (
|
|
"/.well-known/oauth-authorization-server",
|
|
"/.well-known/openid-configuration",
|
|
):
|
|
issuer: Final = wire_holder[0].url
|
|
return Reply(
|
|
body=json.dumps(
|
|
{
|
|
"issuer": issuer,
|
|
"authorization_endpoint": issuer + "/authorize",
|
|
"token_endpoint": issuer + "/token",
|
|
"registration_endpoint": issuer + "/register",
|
|
"response_types_supported": ["code"],
|
|
"grant_types_supported": ["authorization_code", "refresh_token"],
|
|
"code_challenge_methods_supported": ["S256"],
|
|
"token_endpoint_auth_methods_supported": ["client_secret_post"],
|
|
"authorization_response_iss_parameter_supported": True,
|
|
}
|
|
).encode()
|
|
)
|
|
return Reply(status=404, body=json.dumps({"error": "not_found"}).encode())
|
|
|
|
with wire_server(respond) as wire, mcp_peer() as peer, gateway.scenario() as scenario:
|
|
wire_holder.append(wire)
|
|
alias: Final = "rt5" + uuid.uuid4().hex[:10]
|
|
identity: Final = register_mcp(
|
|
scenario,
|
|
peer,
|
|
alias,
|
|
issuer=wire.url,
|
|
auth_type="oauth2",
|
|
oauth2_flow="authorization_code",
|
|
credentials={"client_id": "ac-client", "client_secret": "ac-secret", "scopes": ["tools.call"]},
|
|
)
|
|
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
|
|
|
relay, cookie, cookie_value = _start_callback_flow(
|
|
gateway, alias, key, "ac-client", expected_authorization_url=wire.url + "/authorize"
|
|
)
|
|
code: Final = "callback-code-" + uuid.uuid4().hex
|
|
missing_issuer: Final = gateway.client.get(
|
|
"/callback",
|
|
params={"code": code, "state": relay},
|
|
cookies={cookie: cookie_value},
|
|
)
|
|
assert missing_issuer.status_code == 400, missing_issuer.text
|
|
assert missing_issuer.headers["content-type"] == "text/html; charset=utf-8", missing_issuer.headers
|
|
assert missing_issuer.text == (
|
|
"<html><body><h2>Authentication failed</h2><p><strong>Error:</strong> invalid_issuer</p>"
|
|
"<p>This authorization response came from a different identity provider than the one this "
|
|
"MCP server is configured to use.</p>"
|
|
"<p>You can close this window and try again.</p></body></html>"
|
|
), missing_issuer.text
|
|
assert "location" not in missing_issuer.headers, missing_issuer.headers
|
|
assert code not in missing_issuer.text, missing_issuer.text
|
|
_assert_cleared_oauth_state_cookie(missing_issuer, cookie)
|
|
|
|
trusted_relay, trusted_cookie, trusted_cookie_value = _start_callback_flow(
|
|
gateway, alias, key, "ac-client", expected_authorization_url=wire.url + "/authorize"
|
|
)
|
|
trusted_code: Final = "callback-code-" + uuid.uuid4().hex
|
|
trusted: Final = gateway.client.get(
|
|
"/callback",
|
|
params={"code": trusted_code, "state": trusted_relay, "iss": wire.url},
|
|
cookies={trusted_cookie: trusted_cookie_value},
|
|
)
|
|
assert trusted.status_code == 302, trusted.text
|
|
trusted_location: Final = urlsplit(trusted.headers["location"])
|
|
assert trusted_location.scheme + "://" + trusted_location.netloc + trusted_location.path == CLIENT_REDIRECT, (
|
|
trusted.headers["location"]
|
|
)
|
|
trusted_query: Final = parse_qs(trusted_location.query)
|
|
assert set(trusted_query) == {"code", "state"}, trusted.headers["location"]
|
|
assert trusted_query["state"] == ["client-state"], trusted.headers["location"]
|
|
assert len(trusted_query["code"]) == 1 and trusted_query["code"][0].startswith("llm_scope_"), trusted.headers[
|
|
"location"
|
|
]
|
|
assert trusted_query["code"][0] != trusted_code, trusted.headers["location"]
|
|
_assert_cleared_oauth_state_cookie(trusted, trusted_cookie)
|
|
|
|
wire_requests: Final = wire.drain()
|
|
wire_method_paths: Final = tuple((request.method, urlsplit(request.target).path) for request in wire_requests)
|
|
assert wire_method_paths == (("GET", "/.well-known/oauth-authorization-server"),), wire_method_paths
|
|
assert not any(method == "POST" and path == "/token" for method, path in wire_method_paths), wire_method_paths
|
|
|
|
|
|
@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
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _Hold:
|
|
armed: threading.Event = field(default_factory=threading.Event)
|
|
released: threading.Event = field(default_factory=threading.Event)
|
|
|
|
|
|
def _idp_upstream(origin: Callable[[], str], moved: threading.Event, hold: _Hold | None = None) -> _Upstream:
|
|
def issuer() -> str:
|
|
return origin() + ("/idp-after" if moved.is_set() else "/idp-before")
|
|
|
|
def respond(request: Request) -> Reply:
|
|
if "oauth-authorization-server" in request.target or "openid-configuration" in request.target:
|
|
current: Final = issuer()
|
|
return Reply(
|
|
body=json.dumps(
|
|
{
|
|
"issuer": current,
|
|
"authorization_endpoint": current + "/authorize",
|
|
"token_endpoint": current + "/token",
|
|
}
|
|
).encode()
|
|
)
|
|
if request.target.startswith("/.well-known/oauth-protected-resource"):
|
|
body: Final = json.dumps({"resource": origin() + "/mcp", "authorization_servers": [issuer()]}).encode()
|
|
if hold is not None and hold.armed.is_set():
|
|
assert hold.released.wait(timeout=15), "the held upstream metadata reply was never released"
|
|
return Reply(body=body)
|
|
return Reply(status=404, body=b'{"error":"unexpected"}')
|
|
|
|
return respond
|
|
|
|
|
|
def _register_pass_through(scenario: Scenario, wire: Wire, alias: str) -> str:
|
|
return register_mcp(scenario, McpPeer(wire.url + "/mcp", queue.Queue()), alias, auth_type="true_passthrough")
|
|
|
|
|
|
def _wire_requests(wire: Wire, seen: list[Request]) -> Callable[[], tuple[Request, ...]]:
|
|
def observed() -> tuple[Request, ...]:
|
|
seen.extend(wire.drain())
|
|
return tuple(seen)
|
|
|
|
return observed
|
|
|
|
|
|
def _registration_discovery_settled(requests: tuple[Request, ...]) -> bool:
|
|
return any(
|
|
"oauth-authorization-server" in item.target or "openid-configuration" in item.target for item in requests
|
|
)
|
|
|
|
|
|
def _advertised_authorization_servers(gateway: Gateway, alias: str) -> tuple[str, ...]:
|
|
response: Final = gateway.client.get(f"/.well-known/oauth-protected-resource/{alias}/mcp")
|
|
assert response.status_code == 200, response.text
|
|
return tuple(TypeAdapter(list[str]).validate_python(response.json()["authorization_servers"]))
|
|
|
|
|
|
def _eventually_advertises(gateway: Gateway, alias: str, issuer: str) -> None:
|
|
eventually(
|
|
lambda: gateway.client.get(f"/.well-known/oauth-protected-resource/{alias}/mcp"),
|
|
lambda response: response.status_code == 200 and response.json()["authorization_servers"] == [issuer],
|
|
seconds=40,
|
|
)
|
|
|
|
|
|
def test_saving_a_pass_through_server_refetches_its_upstream_oauth_metadata(gateway: Gateway) -> None:
|
|
moved: Final = threading.Event()
|
|
with wire_server(_idp_upstream(lambda: wire.url, moved)) as wire, gateway.scenario() as scenario:
|
|
alias: Final = "pt" + uuid.uuid4().hex[:8]
|
|
identity: Final = _register_pass_through(scenario, wire, alias)
|
|
assert _advertised_authorization_servers(gateway, alias) == (wire.url + "/idp-before",)
|
|
assert _advertised_authorization_servers(gateway, alias) == (wire.url + "/idp-before",)
|
|
moved.set()
|
|
wire.drain()
|
|
saved: Final = gateway.request("PUT", "/v1/mcp/server", {"server_id": identity, "description": "IdP moved"})
|
|
assert saved.status_code == 202, saved.text
|
|
assert _advertised_authorization_servers(gateway, alias) == (wire.url + "/idp-after",)
|
|
assert any(request.target.startswith("/.well-known/oauth-protected-resource") for request in wire.drain()), (
|
|
"the save must send protected-resource discovery back to the upstream"
|
|
)
|
|
|
|
|
|
def test_peer_worker_stops_advertising_the_old_idp_after_a_save_on_another_worker(
|
|
gateway: Gateway, peer: Gateway
|
|
) -> None:
|
|
moved: Final = threading.Event()
|
|
with wire_server(_idp_upstream(lambda: wire.url, moved)) as wire, gateway.scenario() as scenario:
|
|
alias: Final = "pt" + uuid.uuid4().hex[:8]
|
|
identity: Final = _register_pass_through(scenario, wire, alias)
|
|
assert _advertised_authorization_servers(gateway, alias) == (wire.url + "/idp-before",)
|
|
_eventually_advertises(peer, alias, wire.url + "/idp-before")
|
|
moved.set()
|
|
saved: Final = gateway.request("PUT", "/v1/mcp/server", {"server_id": identity, "description": "IdP moved"})
|
|
assert saved.status_code == 202, saved.text
|
|
assert _advertised_authorization_servers(gateway, alias) == (wire.url + "/idp-after",)
|
|
_eventually_advertises(peer, alias, wire.url + "/idp-after")
|
|
|
|
|
|
def test_metadata_fetched_before_a_save_cannot_repopulate_the_cache_after_it(gateway: Gateway) -> None:
|
|
moved: Final = threading.Event()
|
|
hold: Final = _Hold()
|
|
with (
|
|
wire_server(_idp_upstream(lambda: wire.url, moved, hold)) as wire,
|
|
gateway.scenario() as scenario,
|
|
ThreadPoolExecutor(max_workers=1) as pool,
|
|
):
|
|
alias: Final = "pt" + uuid.uuid4().hex[:8]
|
|
identity: Final = _register_pass_through(scenario, wire, alias)
|
|
seen: Final[list[Request]] = []
|
|
observed: Final = _wire_requests(wire, seen)
|
|
eventually(observed, _registration_discovery_settled, seconds=10)
|
|
settled: Final = len(seen)
|
|
hold.armed.set()
|
|
stale: Final = pool.submit(_advertised_authorization_servers, gateway, alias)
|
|
eventually(observed, lambda requests: len(requests) > settled, seconds=10)
|
|
assert seen[settled].target.startswith("/.well-known/oauth-protected-resource"), seen[settled:]
|
|
moved.set()
|
|
saved: Final = gateway.request("PUT", "/v1/mcp/server", {"server_id": identity, "description": "IdP moved"})
|
|
assert saved.status_code == 202, saved.text
|
|
hold.released.set()
|
|
assert stale.result(timeout=30) == (wire.url + "/idp-before",)
|
|
assert _advertised_authorization_servers(gateway, alias) == (wire.url + "/idp-after",)
|
|
|
|
|
|
@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
|