mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
test(integration): cover mcp_runtime endpoint contracts (#44671)
* test(integration): cover mcp_runtime endpoint contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(integration): ruff format and import order in mcp runtime tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): assert out-of-path calls and untrusted forwarded headers are refused Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): move the untrusted forwarded-header discovery case into its own test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): hash the token row lookup, widen spend polling, and union open servers after listing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): tighten mcp runtime contracts per review (sdk client, toolset scope, forwarded trust, callback html, session tools/call) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): assert the sdk unknown-path error without pinning the sdk's exception wrappers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover trailing-path and group csv mcp paths, rfc 9207 iss-less and error-iss callbacks, and id_token exchange Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover lite logout spelling, team-mapped and unmapped IdP subjects Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover per-server mcp-session resource and root /token spelling Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): unpack the sealed client code in the step-up tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
595cf066cd
commit
27209b92b3
3 changed files with 1987 additions and 24 deletions
File diff suppressed because it is too large
Load diff
|
|
@ -1,27 +1,122 @@
|
|||
import builtins
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from mcp.shared.exceptions import MCPError
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.mcp import (
|
||||
ENTRY_POINTS,
|
||||
INITIALIZE,
|
||||
PEER_KINDS,
|
||||
EntryPoint,
|
||||
McpCaller,
|
||||
Outcome,
|
||||
PeerKind,
|
||||
_outcome_from_rpc,
|
||||
_parse_rpc_body,
|
||||
mcp_peer,
|
||||
official_client_outcomes,
|
||||
peer_of,
|
||||
register_mcp,
|
||||
tool_calls,
|
||||
)
|
||||
from integration._support.mcp_grants import create_toolset
|
||||
|
||||
ADD: Final = {"http": "add", "sse": "add", "stdio": "add", "openapi": "getpet"}
|
||||
ARGUMENTS: Final = {"add": {"a": 3, "b": 4}, "getpet": {"petId": "7"}}
|
||||
EXPECTED: Final = {"add": "7", "getpet": json.dumps({"id": "7", "name": "integration-pet"})}
|
||||
|
||||
|
||||
def _streamable_rpc(gateway: Gateway, path: str, key: str, method: str, params: Mapping[str, object]) -> httpx.Response:
|
||||
return gateway.client.post(
|
||||
path,
|
||||
json={"jsonrpc": "2.0", "id": 1, "method": method, "params": params},
|
||||
headers={"Authorization": f"Bearer {key}", "Accept": "application/json, text/event-stream"},
|
||||
)
|
||||
|
||||
|
||||
def _streamable_tools(gateway: Gateway, path: str, key: str) -> tuple[httpx.Response, Outcome]:
|
||||
response: Final = _streamable_rpc(gateway, path, key, "tools/list", {})
|
||||
return response, _outcome_from_rpc(response)
|
||||
|
||||
|
||||
def _list_and_call(
|
||||
gateway: Gateway, client: str, path: str, key: str, name: str, arguments: Mapping[str, object]
|
||||
) -> tuple[Outcome, Outcome]:
|
||||
if client == "official-sdk":
|
||||
return official_client_outcomes(gateway, key, path, name, dict(arguments))
|
||||
initialized: Final = _streamable_rpc(gateway, path, key, "initialize", dict(INITIALIZE))
|
||||
assert initialized.status_code == 200, initialized.text
|
||||
listed_response, listed = _streamable_tools(gateway, path, key)
|
||||
assert listed_response.status_code == 200, listed_response.text
|
||||
called_response: Final = _streamable_rpc(
|
||||
gateway, path, key, "tools/call", {"name": name, "arguments": dict(arguments)}
|
||||
)
|
||||
assert called_response.status_code == 200, called_response.text
|
||||
return listed, _outcome_from_rpc(called_response)
|
||||
|
||||
|
||||
def _refused_call(
|
||||
gateway: Gateway,
|
||||
client: str,
|
||||
path: str,
|
||||
key: str,
|
||||
name: str,
|
||||
arguments: Mapping[str, object],
|
||||
text: str = "Error: User not allowed to call this tool.",
|
||||
) -> Outcome:
|
||||
expected: Final = text
|
||||
if client == "official-sdk":
|
||||
_, refused = official_client_outcomes(gateway, key, path, name, dict(arguments))
|
||||
assert (refused.status, refused.error, refused.text) == (200, expected, expected), refused
|
||||
return refused
|
||||
response: Final = _streamable_rpc(gateway, path, key, "tools/call", {"name": name, "arguments": dict(arguments)})
|
||||
assert response.status_code == 200, response.text
|
||||
assert _parse_rpc_body(response) == {
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"result": {
|
||||
"content": [{"text": expected, "type": "text"}],
|
||||
"isError": True,
|
||||
},
|
||||
}, response.text
|
||||
return _outcome_from_rpc(response)
|
||||
|
||||
|
||||
def _leaf_exceptions(error: BaseException) -> tuple[BaseException, ...]:
|
||||
if isinstance(error, builtins.BaseExceptionGroup):
|
||||
return tuple(leaf for child in error.exceptions for leaf in _leaf_exceptions(child))
|
||||
return (error,)
|
||||
|
||||
|
||||
def _assert_upstream_call(call: Mapping[str, object], name: str, arguments: Mapping[str, object]) -> None:
|
||||
body: Final = call["body"]
|
||||
assert isinstance(body, Mapping), call
|
||||
params: Final = body["params"]
|
||||
assert isinstance(params, Mapping), call
|
||||
assert set(params) == {"name", "arguments", "_meta"}, call
|
||||
assert params["name"] == name, call
|
||||
assert params["arguments"] == arguments, call
|
||||
metadata: Final = params["_meta"]
|
||||
assert isinstance(metadata, Mapping), call
|
||||
assert set(metadata) == {"progressToken"}, call
|
||||
|
||||
|
||||
def _open_aliases() -> frozenset[str]:
|
||||
rows: Final = read_rows('SELECT alias FROM "LiteLLM_MCPServerTable" WHERE allow_all_keys', ())
|
||||
return frozenset(str(row["alias"]) for row in rows)
|
||||
|
||||
|
||||
def _without_foreign_open_servers(tools: tuple[str, ...], open_aliases: frozenset[str]) -> set[str]:
|
||||
prefixes: Final = tuple(f"{alias}-" for alias in open_aliases)
|
||||
return {tool for tool in tools if not tool.startswith(prefixes)}
|
||||
|
||||
|
||||
def _peer_saw_call(peer_kind: PeerKind, observed: tuple[dict[str, object], ...], tool: str) -> bool:
|
||||
if peer_kind == "openapi":
|
||||
return any(item.get("path") == "/pets/7" and item.get("method") == "GET" for item in observed)
|
||||
|
|
@ -95,6 +190,306 @@ def test_official_client_session_lists_and_calls_through_gateway(
|
|||
assert _peer_saw_call(peer_kind, peer.drain(), "add")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client", ("raw-jsonrpc", "official-sdk"))
|
||||
def test_standard_per_server_path_scopes_the_session_to_the_named_servers(gateway: Gateway, client: str) -> None:
|
||||
open_before: Final = _open_aliases()
|
||||
with mcp_peer() as peer_a, mcp_peer() as peer_b, mcp_peer() as peer_c, gateway.scenario() as scenario:
|
||||
alias_a: Final = "rt1a" + uuid.uuid4().hex[:10]
|
||||
alias_b: Final = "rt1b" + uuid.uuid4().hex[:10]
|
||||
alias_c: Final = "rt1c" + uuid.uuid4().hex[:10]
|
||||
server_a: Final = register_mcp(scenario, peer_a, alias_a)
|
||||
server_b: Final = register_mcp(scenario, peer_b, alias_b)
|
||||
server_c: Final = register_mcp(scenario, peer_c, alias_c)
|
||||
key: Final = scenario.key(object_permission={"mcp_servers": [server_a, server_b, server_c]})
|
||||
open_aliases: Final = open_before | _open_aliases()
|
||||
|
||||
first_path: Final = f"/mcp/{alias_a}"
|
||||
listed, called = _list_and_call(gateway, client, first_path, key, f"{alias_a}-add", {"a": 20, "b": 22})
|
||||
assert listed.ok, listed.raw
|
||||
assert _without_foreign_open_servers(listed.tools, open_aliases) == {
|
||||
f"{alias_a}-add",
|
||||
f"{alias_a}-multiply",
|
||||
f"{alias_a}-fail",
|
||||
}, listed.tools
|
||||
assert called.ok and called.text == "42", called.raw
|
||||
_refused_call(
|
||||
gateway,
|
||||
client,
|
||||
first_path,
|
||||
key,
|
||||
f"{alias_b}-add",
|
||||
{"a": 20, "b": 22},
|
||||
f"Error: Tool {alias_b}-add not found",
|
||||
)
|
||||
calls_a: Final = tool_calls(peer_a.drain())
|
||||
assert len(calls_a) == 1, calls_a
|
||||
_assert_upstream_call(calls_a[0], "add", {"a": 20, "b": 22})
|
||||
assert tool_calls(peer_b.drain()) == ()
|
||||
assert tool_calls(peer_c.drain()) == ()
|
||||
|
||||
combined_path: Final = f"/mcp/{alias_a},{alias_b}"
|
||||
combined_listed, combined_called = _list_and_call(
|
||||
gateway, client, combined_path, key, f"{alias_b}-multiply", {"a": 6, "b": 7}
|
||||
)
|
||||
assert combined_listed.ok, combined_listed.raw
|
||||
assert _without_foreign_open_servers(combined_listed.tools, open_aliases) == {
|
||||
f"{alias_a}-add",
|
||||
f"{alias_a}-multiply",
|
||||
f"{alias_a}-fail",
|
||||
f"{alias_b}-add",
|
||||
f"{alias_b}-multiply",
|
||||
f"{alias_b}-fail",
|
||||
}, combined_listed.tools
|
||||
assert combined_called.ok and combined_called.text == "42", combined_called.raw
|
||||
_refused_call(gateway, client, combined_path, key, f"{alias_c}-add", {"a": 20, "b": 22})
|
||||
assert tool_calls(peer_a.drain()) == ()
|
||||
calls_b: Final = tool_calls(peer_b.drain())
|
||||
assert len(calls_b) == 1, calls_b
|
||||
_assert_upstream_call(calls_b[0], "multiply", {"a": 6, "b": 7})
|
||||
assert tool_calls(peer_c.drain()) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client", ("raw-jsonrpc", "official-sdk"))
|
||||
def test_standard_comma_path_with_a_trailing_segment_scopes_a_proxy_admin_session(
|
||||
gateway: Gateway, client: str
|
||||
) -> None:
|
||||
open_before: Final = _open_aliases()
|
||||
with mcp_peer() as peer_a, mcp_peer() as peer_b, mcp_peer() as peer_c, gateway.scenario() as scenario:
|
||||
alias_a: Final = "rt1pa" + uuid.uuid4().hex[:10]
|
||||
alias_b: Final = "rt1pb" + uuid.uuid4().hex[:10]
|
||||
alias_c: Final = "rt1pc" + uuid.uuid4().hex[:10]
|
||||
register_mcp(scenario, peer_a, alias_a)
|
||||
register_mcp(scenario, peer_b, alias_b)
|
||||
register_mcp(scenario, peer_c, alias_c)
|
||||
open_aliases: Final = open_before | _open_aliases()
|
||||
path: Final = f"/mcp/{alias_a},{alias_c}/tools"
|
||||
|
||||
listed, called = _list_and_call(gateway, client, path, gateway.key, f"{alias_c}-add", {"a": 20, "b": 22})
|
||||
assert listed.ok, listed.raw
|
||||
assert _without_foreign_open_servers(listed.tools, open_aliases) == {
|
||||
f"{alias_a}-add",
|
||||
f"{alias_a}-multiply",
|
||||
f"{alias_a}-fail",
|
||||
f"{alias_c}-add",
|
||||
f"{alias_c}-multiply",
|
||||
f"{alias_c}-fail",
|
||||
}, listed.tools
|
||||
assert called.ok and called.text == "42", called.raw
|
||||
_refused_call(gateway, client, path, gateway.key, f"{alias_b}-add", {"a": 20, "b": 22})
|
||||
assert tool_calls(peer_a.drain()) == ()
|
||||
assert tool_calls(peer_b.drain()) == ()
|
||||
calls_c: Final = tool_calls(peer_c.drain())
|
||||
assert len(calls_c) == 1, calls_c
|
||||
_assert_upstream_call(calls_c[0], "add", {"a": 20, "b": 22})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client", ("raw-jsonrpc", "official-sdk"))
|
||||
def test_standard_comma_path_with_a_trailing_segment_admits_a_scoped_key(gateway: Gateway, client: str) -> None:
|
||||
pytest.skip(
|
||||
"BUG: a non-admin key gets 401 'Only proxy admin can be used...' on /mcp/{a},{b}/tools because "
|
||||
"mcp_inference_routes only admits the single-segment /mcp/{subpath}, while "
|
||||
"_get_mcp_servers_in_path parses the trailing path"
|
||||
)
|
||||
open_before: Final = _open_aliases()
|
||||
with mcp_peer() as peer_a, mcp_peer() as peer_b, mcp_peer() as peer_c, gateway.scenario() as scenario:
|
||||
alias_a: Final = "rt1sa" + uuid.uuid4().hex[:10]
|
||||
alias_b: Final = "rt1sb" + uuid.uuid4().hex[:10]
|
||||
alias_c: Final = "rt1sc" + uuid.uuid4().hex[:10]
|
||||
server_a: Final = register_mcp(scenario, peer_a, alias_a)
|
||||
server_b: Final = register_mcp(scenario, peer_b, alias_b)
|
||||
server_c: Final = register_mcp(scenario, peer_c, alias_c)
|
||||
key: Final = scenario.key(object_permission={"mcp_servers": [server_a, server_b, server_c]})
|
||||
open_aliases: Final = open_before | _open_aliases()
|
||||
path: Final = f"/mcp/{alias_a},{alias_c}/tools"
|
||||
|
||||
listed, called = _list_and_call(gateway, client, path, key, f"{alias_c}-add", {"a": 20, "b": 22})
|
||||
assert listed.ok, listed.raw
|
||||
assert _without_foreign_open_servers(listed.tools, open_aliases) == {
|
||||
f"{alias_a}-add",
|
||||
f"{alias_a}-multiply",
|
||||
f"{alias_a}-fail",
|
||||
f"{alias_c}-add",
|
||||
f"{alias_c}-multiply",
|
||||
f"{alias_c}-fail",
|
||||
}, listed.tools
|
||||
assert called.ok and called.text == "42", called.raw
|
||||
_refused_call(gateway, client, path, key, f"{alias_b}-add", {"a": 20, "b": 22})
|
||||
assert tool_calls(peer_a.drain()) == ()
|
||||
assert tool_calls(peer_b.drain()) == ()
|
||||
calls_c: Final = tool_calls(peer_c.drain())
|
||||
assert len(calls_c) == 1, calls_c
|
||||
_assert_upstream_call(calls_c[0], "add", {"a": 20, "b": 22})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client", ("raw-jsonrpc", "official-sdk"))
|
||||
def test_legacy_server_path_resolves_lists_groups_and_fails_closed_on_unknown_names(
|
||||
gateway: Gateway, client: str
|
||||
) -> None:
|
||||
open_before: Final = _open_aliases()
|
||||
with mcp_peer() as peer_a, mcp_peer() as peer_b, mcp_peer() as peer_c, gateway.scenario() as scenario:
|
||||
alias_a: Final = "rt2a" + uuid.uuid4().hex[:10]
|
||||
alias_b: Final = "rt2b" + uuid.uuid4().hex[:10]
|
||||
alias_c: Final = "rt2c" + uuid.uuid4().hex[:10]
|
||||
group: Final = "rt2g" + uuid.uuid4().hex[:10]
|
||||
server_a: Final = register_mcp(scenario, peer_a, alias_a, mcp_access_groups=[group])
|
||||
server_b: Final = register_mcp(scenario, peer_b, alias_b, mcp_access_groups=[group])
|
||||
server_c: Final = register_mcp(scenario, peer_c, alias_c)
|
||||
key: Final = scenario.key(object_permission={"mcp_servers": [server_a, server_b, server_c]})
|
||||
open_aliases: Final = open_before | _open_aliases()
|
||||
|
||||
direct_path: Final = f"/{alias_a},{alias_b}/mcp"
|
||||
direct_listed, direct_called = _list_and_call(
|
||||
gateway, client, direct_path, key, f"{alias_a}-add", {"a": 20, "b": 22}
|
||||
)
|
||||
assert direct_listed.ok, direct_listed.raw
|
||||
assert _without_foreign_open_servers(direct_listed.tools, open_aliases) == {
|
||||
f"{alias_a}-add",
|
||||
f"{alias_a}-multiply",
|
||||
f"{alias_a}-fail",
|
||||
f"{alias_b}-add",
|
||||
f"{alias_b}-multiply",
|
||||
f"{alias_b}-fail",
|
||||
}, direct_listed.tools
|
||||
assert direct_called.ok and direct_called.text == "42", direct_called.raw
|
||||
direct_calls_a: Final = tool_calls(peer_a.drain())
|
||||
assert len(direct_calls_a) == 1, direct_calls_a
|
||||
_assert_upstream_call(direct_calls_a[0], "add", {"a": 20, "b": 22})
|
||||
assert tool_calls(peer_b.drain()) == ()
|
||||
assert tool_calls(peer_c.drain()) == ()
|
||||
|
||||
mixed_path: Final = f"/{alias_a},%20{alias_a}%20,nope{uuid.uuid4().hex}/mcp"
|
||||
mixed_listed, mixed_called = _list_and_call(
|
||||
gateway, client, mixed_path, key, f"{alias_a}-multiply", {"a": 6, "b": 7}
|
||||
)
|
||||
assert mixed_listed.ok, mixed_listed.raw
|
||||
assert _without_foreign_open_servers(mixed_listed.tools, open_aliases) == {
|
||||
f"{alias_a}-add",
|
||||
f"{alias_a}-multiply",
|
||||
f"{alias_a}-fail",
|
||||
}, mixed_listed.tools
|
||||
assert mixed_called.ok and mixed_called.text == "42", mixed_called.raw
|
||||
mixed_calls_a: Final = tool_calls(peer_a.drain())
|
||||
assert len(mixed_calls_a) == 1, mixed_calls_a
|
||||
_assert_upstream_call(mixed_calls_a[0], "multiply", {"a": 6, "b": 7})
|
||||
assert tool_calls(peer_b.drain()) == ()
|
||||
assert tool_calls(peer_c.drain()) == ()
|
||||
|
||||
unknown_path: Final = f"/nope1{uuid.uuid4().hex},nope2{uuid.uuid4().hex}/mcp"
|
||||
assert peer_a.drain() == ()
|
||||
assert peer_b.drain() == ()
|
||||
assert peer_c.drain() == ()
|
||||
if client == "official-sdk":
|
||||
with pytest.raises(builtins.ExceptionGroup) as caught:
|
||||
official_client_outcomes(gateway, key, unknown_path, f"{alias_a}-add", {"a": 20, "b": 22})
|
||||
leaves: Final = _leaf_exceptions(caught.value)
|
||||
assert len(leaves) == 1, repr(leaves)
|
||||
leaf: Final = leaves[0]
|
||||
assert isinstance(leaf, MCPError), repr(leaves)
|
||||
assert leaf.error.code == -32601, repr(leaves)
|
||||
assert leaf.error.message == "Not Found" and leaf.error.data is None, repr(leaves)
|
||||
else:
|
||||
unknown_initialized: Final = _streamable_rpc(gateway, unknown_path, key, "initialize", dict(INITIALIZE))
|
||||
assert unknown_initialized.status_code == 404, unknown_initialized.text
|
||||
unknown_listed: Final = _streamable_rpc(gateway, unknown_path, key, "tools/list", {})
|
||||
assert unknown_listed.status_code == 404, unknown_listed.text
|
||||
assert peer_a.drain() == ()
|
||||
assert peer_b.drain() == ()
|
||||
assert peer_c.drain() == ()
|
||||
|
||||
group_path: Final = f"/{group}/mcp"
|
||||
group_listed, group_called = _list_and_call(
|
||||
gateway, client, group_path, key, f"{alias_b}-multiply", {"a": 6, "b": 7}
|
||||
)
|
||||
assert group_listed.ok, group_listed.raw
|
||||
assert _without_foreign_open_servers(group_listed.tools, open_aliases) == {
|
||||
f"{alias_a}-add",
|
||||
f"{alias_a}-multiply",
|
||||
f"{alias_a}-fail",
|
||||
f"{alias_b}-add",
|
||||
f"{alias_b}-multiply",
|
||||
f"{alias_b}-fail",
|
||||
}, group_listed.tools
|
||||
assert group_called.ok and group_called.text == "42", group_called.raw
|
||||
assert tool_calls(peer_a.drain()) == ()
|
||||
group_calls_b: Final = tool_calls(peer_b.drain())
|
||||
assert len(group_calls_b) == 1, group_calls_b
|
||||
_assert_upstream_call(group_calls_b[0], "multiply", {"a": 6, "b": 7})
|
||||
assert tool_calls(peer_c.drain()) == ()
|
||||
|
||||
mixed_group_path: Final = f"/{group},{alias_c}/mcp"
|
||||
mixed_group_listed, mixed_group_called = _list_and_call(
|
||||
gateway, client, mixed_group_path, key, f"{alias_c}-add", {"a": 20, "b": 22}
|
||||
)
|
||||
assert mixed_group_listed.ok, mixed_group_listed.raw
|
||||
assert _without_foreign_open_servers(mixed_group_listed.tools, open_aliases) == {
|
||||
f"{alias_a}-add",
|
||||
f"{alias_a}-multiply",
|
||||
f"{alias_a}-fail",
|
||||
f"{alias_b}-add",
|
||||
f"{alias_b}-multiply",
|
||||
f"{alias_b}-fail",
|
||||
f"{alias_c}-add",
|
||||
f"{alias_c}-multiply",
|
||||
f"{alias_c}-fail",
|
||||
}, mixed_group_listed.tools
|
||||
assert mixed_group_called.ok and mixed_group_called.text == "42", mixed_group_called.raw
|
||||
assert tool_calls(peer_a.drain()) == ()
|
||||
assert tool_calls(peer_b.drain()) == ()
|
||||
mixed_group_calls_c: Final = tool_calls(peer_c.drain())
|
||||
assert len(mixed_group_calls_c) == 1, mixed_group_calls_c
|
||||
_assert_upstream_call(mixed_group_calls_c[0], "add", {"a": 20, "b": 22})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("client", ("raw-jsonrpc", "official-sdk"))
|
||||
def test_legacy_toolset_path_scopes_the_session_to_the_toolsets_tools(gateway: Gateway, client: str) -> None:
|
||||
pytest.skip(
|
||||
"BUG: a /{toolset}/mcp session lists every granted server's tools instead of the toolset's "
|
||||
"tools (the toolset's tool filter is not applied on the legacy path)"
|
||||
)
|
||||
open_before: Final = _open_aliases()
|
||||
with mcp_peer() as peer_a, mcp_peer() as peer_b, mcp_peer() as peer_c, gateway.scenario() as scenario:
|
||||
alias_a: Final = "rt2ta" + uuid.uuid4().hex[:10]
|
||||
alias_b: Final = "rt2tb" + uuid.uuid4().hex[:10]
|
||||
alias_c: Final = "rt2tc" + uuid.uuid4().hex[:10]
|
||||
server_a: Final = register_mcp(scenario, peer_a, alias_a)
|
||||
server_b: Final = register_mcp(scenario, peer_b, alias_b)
|
||||
server_c: Final = register_mcp(scenario, peer_c, alias_c)
|
||||
toolset_name: Final = "rt2t" + uuid.uuid4().hex[:10]
|
||||
toolset_id: Final = create_toolset(scenario, ((server_a, "add"),), toolset_name=toolset_name)
|
||||
toolset_key: Final = scenario.key(
|
||||
object_permission={
|
||||
"mcp_servers": [server_a, server_b, server_c],
|
||||
"mcp_toolsets": [toolset_id],
|
||||
}
|
||||
)
|
||||
open_aliases: Final = open_before | _open_aliases()
|
||||
|
||||
toolset_path: Final = f"/{toolset_name}/mcp"
|
||||
toolset_listed, toolset_called = _list_and_call(
|
||||
gateway, client, toolset_path, toolset_key, f"{alias_a}-add", {"a": 20, "b": 22}
|
||||
)
|
||||
assert toolset_listed.ok, toolset_listed.raw
|
||||
assert _without_foreign_open_servers(toolset_listed.tools, open_aliases) == {f"{alias_a}-add"}, (
|
||||
toolset_listed.tools
|
||||
)
|
||||
assert toolset_called.ok and toolset_called.text == "42", toolset_called.raw
|
||||
_refused_call(
|
||||
gateway,
|
||||
client,
|
||||
toolset_path,
|
||||
toolset_key,
|
||||
f"{alias_a}-multiply",
|
||||
{"a": 6, "b": 7},
|
||||
f"Error: Tool 'multiply' is not allowed for your key/team on server '{alias_a}'. Contact proxy admin for access.",
|
||||
)
|
||||
_refused_call(gateway, client, toolset_path, toolset_key, f"{alias_b}-add", {"a": 20, "b": 22})
|
||||
toolset_calls_a: Final = tool_calls(peer_a.drain())
|
||||
assert len(toolset_calls_a) == 1, toolset_calls_a
|
||||
_assert_upstream_call(toolset_calls_a[0], "add", {"a": 20, "b": 22})
|
||||
assert tool_calls(peer_b.drain()) == ()
|
||||
assert tool_calls(peer_c.drain()) == ()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("peer_kind", ("http", "sse", "stdio"))
|
||||
def test_prompts_resources_and_templates_are_proxied_from_rich_peer(gateway: Gateway, peer_kind: PeerKind) -> None:
|
||||
with peer_of(peer_kind, rich=True) as peer, gateway.scenario() as scenario:
|
||||
|
|
@ -177,14 +572,19 @@ def test_pinned_revision_pairs_list_and_call_through_gateway(
|
|||
key: Final = scenario.key(object_permission={"mcp_servers": [identity]})
|
||||
endpoint: Final = str(gateway.client.base_url).rstrip("/") + ("/mcp/sse" if ingress == "sse" else "/mcp")
|
||||
client: Final = MCPClient(
|
||||
server_url=endpoint, transport_type=MCPTransport(ingress), protocol_version=downstream,
|
||||
extra_headers={"Authorization": f"Bearer {key}", "x-mcp-servers": identity}, timeout=15,
|
||||
server_url=endpoint,
|
||||
transport_type=MCPTransport(ingress),
|
||||
protocol_version=downstream,
|
||||
extra_headers={"Authorization": f"Bearer {key}", "x-mcp-servers": identity},
|
||||
timeout=15,
|
||||
)
|
||||
|
||||
async def exercise() -> None:
|
||||
tools: Final = await client.list_tools(raise_on_error=True)
|
||||
assert f"{alias}-add" in tuple(tool.name for tool in tools)
|
||||
result: Final = await client.call_tool(CallToolRequestParams(name=f"{alias}-add", arguments={"a": 3, "b": 4}))
|
||||
result: Final = await client.call_tool(
|
||||
CallToolRequestParams(name=f"{alias}-add", arguments={"a": 3, "b": 4})
|
||||
)
|
||||
assert result.is_error is False
|
||||
assert result.content[0].text == "7"
|
||||
|
||||
|
|
|
|||
|
|
@ -9,15 +9,263 @@ 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
|
||||
|
||||
_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")
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue