forked from modelcontextprotocol/python-sdk
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_completion.py
More file actions
137 lines (102 loc) · 5.75 KB
/
Copy pathtest_completion.py
File metadata and controls
137 lines (102 loc) · 5.75 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
"""Completion interactions against the low-level Server, driven through the public Client API."""
import mcp_types as types
import pytest
from inline_snapshot import snapshot
from mcp_types import (
INVALID_PARAMS,
METHOD_NOT_FOUND,
CompleteResult,
Completion,
ErrorData,
PromptReference,
ResourceTemplateReference,
)
from mcp import MCPError
from mcp.server import Server, ServerRequestContext
from tests._stamp import Unstamp
from tests.interaction._connect import Connect
from tests.interaction._requirements import requirement
pytestmark = pytest.mark.anyio
@requirement("completion:prompt-arg")
@requirement("completion:result-shape")
async def test_complete_prompt_argument(connect: Connect, unstamped: Unstamp) -> None:
"""Completing a prompt argument delivers the ref, argument name, and current value to the handler.
The returned values are filtered by the argument's value, proving the value reached the handler.
"""
async def completion(ctx: ServerRequestContext, params: types.CompleteRequestParams) -> CompleteResult:
assert isinstance(params.ref, PromptReference)
assert params.ref.name == "code_review"
assert params.argument.name == "language"
candidates = ["python", "pytorch", "ruby"]
matches = [candidate for candidate in candidates if candidate.startswith(params.argument.value)]
return CompleteResult(completion=Completion(values=matches, total=len(matches), has_more=False))
server = Server("completer", on_completion=completion)
async with connect(server) as client:
result = await client.complete(
PromptReference(name="code_review"), argument={"name": "language", "value": "py"}
)
assert unstamped(result) == snapshot(
CompleteResult(completion=Completion(values=["python", "pytorch"], total=2, has_more=False))
)
@requirement("completion:resource-template-arg")
async def test_complete_resource_template_variable(connect: Connect, unstamped: Unstamp) -> None:
"""Completing a URI template variable delivers the template URI and variable name to the handler."""
async def completion(ctx: ServerRequestContext, params: types.CompleteRequestParams) -> CompleteResult:
assert isinstance(params.ref, ResourceTemplateReference)
assert params.ref.uri == "github://repos/{owner}/{repo}"
assert params.argument.name == "owner"
return CompleteResult(completion=Completion(values=[f"{params.argument.value}contextprotocol"]))
server = Server("completer", on_completion=completion)
async with connect(server) as client:
result = await client.complete(
ResourceTemplateReference(uri="github://repos/{owner}/{repo}"),
argument={"name": "owner", "value": "model"},
)
assert unstamped(result) == snapshot(CompleteResult(completion=Completion(values=["modelcontextprotocol"])))
@requirement("completion:context-arguments")
async def test_complete_receives_context_arguments(connect: Connect, unstamped: Unstamp) -> None:
"""Previously-resolved arguments passed as completion context reach the handler.
The returned value is derived from the context, proving it arrived.
"""
async def completion(ctx: ServerRequestContext, params: types.CompleteRequestParams) -> CompleteResult:
assert params.argument.name == "repo"
assert params.context is not None
assert params.context.arguments is not None
return CompleteResult(completion=Completion(values=[f"{params.context.arguments['owner']}/python-sdk"]))
server = Server("completer", on_completion=completion)
async with connect(server) as client:
result = await client.complete(
ResourceTemplateReference(uri="github://repos/{owner}/{repo}"),
argument={"name": "repo", "value": ""},
context_arguments={"owner": "modelcontextprotocol"},
)
assert unstamped(result) == snapshot(
CompleteResult(completion=Completion(values=["modelcontextprotocol/python-sdk"]))
)
@requirement("completion:error:invalid-ref")
async def test_completion_against_an_unknown_ref_is_rejected_with_invalid_params(connect: Connect) -> None:
"""completion/complete with a ref naming an unknown prompt is answered with -32602 Invalid params.
The lowlevel server does not validate refs itself (it has no prompt/template registry to check
against); rejecting an unknown ref is the handler's job, and this test pins the spec-recommended
way to do it.
"""
async def completion(ctx: ServerRequestContext, params: types.CompleteRequestParams) -> CompleteResult:
assert isinstance(params.ref, PromptReference)
raise MCPError(code=INVALID_PARAMS, message=f"Unknown prompt: {params.ref.name!r}")
server = Server("completer", on_completion=completion)
async with connect(server) as client:
with pytest.raises(MCPError) as exc_info:
await client.complete(PromptReference(name="ghost"), argument={"name": "x", "value": ""})
assert exc_info.value.error.code == INVALID_PARAMS
@requirement("completion:complete:not-supported")
@requirement("protocol:error:method-not-found")
async def test_complete_without_handler_is_method_not_found(connect: Connect) -> None:
"""A server with no completion handler advertises no completions capability and rejects the request."""
server = Server("incomplete")
async with connect(server) as client:
assert client.server_capabilities.completions is None
with pytest.raises(MCPError) as exc_info:
await client.complete(PromptReference(name="anything"), argument={"name": "topic", "value": ""})
assert exc_info.value.error == snapshot(
ErrorData(code=METHOD_NOT_FOUND, message="Method not found", data="completion/complete")
)