Max Isbey
98bd125ddd
The 2026-07-28 protocol did not work over stdio: subscriptions/listen was hard-refused, a legacy initialize arriving during an in-flight modern request was accepted and re-locked the connection, and a peer cancel produced a trailing "Request cancelled" frame. All three share one root cause: the connection's era was derived from which requests had completed instead of being decided once, in wire order, from how the client opened the connection. Replace serve_dual_era_loop and serve_loop with a single serve_stream driver that decides the era synchronously in the dispatcher's read loop, before the request body is spawned. initialize (or any envelope-less request) opens the legacy era, an enveloped request opens the modern era, server/discover is answered without pinning, and a stray leading notification opens nothing. A conflicting era claim on a committed connection is refused (-32022 or -32600) rather than silently switching. Cancel silence is structural: each request answers through a one-shot channel whose write target becomes a powerless void on a peer cancel, so there is no cancelled-check at any write site. The generic JSON-RPC dispatcher loses the code-0 cancel frame, the code-0 str(exc) catch-all, and the inline_methods knob, and documents that handlers are invoked synchronously in receive order with the returned awaitable as the body. Add a Posture enum (DUAL default, LEGACY_ONLY, MODERN_ONLY) on the Server and MCPServer constructors, honoured by the stream driver and the streamable-HTTP manager alike. Server.run(read, write) now stands alone, Server.lifespan() is a bound context manager, and serve_listener / newline_json_transport / close_subscriptions() give a straightforward path for custom transports. See docs/migration.md for the full list of observable changes.
163 lines
6.8 KiB
Python
163 lines
6.8 KiB
Python
"""Tests for completion handler with context functionality."""
|
|
|
|
import pytest
|
|
from mcp_types import (
|
|
INVALID_PARAMS,
|
|
CompleteRequestParams,
|
|
CompleteResult,
|
|
Completion,
|
|
PromptReference,
|
|
ResourceTemplateReference,
|
|
)
|
|
|
|
from mcp import Client, MCPError
|
|
from mcp.server import Server, ServerRequestContext
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_completion_handler_receives_context():
|
|
"""Test that the completion handler receives context correctly."""
|
|
# Track what the handler receives
|
|
received_params: CompleteRequestParams | None = None
|
|
|
|
async def handle_completion(ctx: ServerRequestContext, params: CompleteRequestParams) -> CompleteResult:
|
|
nonlocal received_params
|
|
received_params = params
|
|
return CompleteResult(completion=Completion(values=["test-completion"], total=1, has_more=False))
|
|
|
|
server = Server("test-server", on_completion=handle_completion)
|
|
|
|
async with Client(server) as client:
|
|
# Test with context
|
|
result = await client.complete(
|
|
ref=ResourceTemplateReference(type="ref/resource", uri="test://resource/{param}"),
|
|
argument={"name": "param", "value": "test"},
|
|
context_arguments={"previous": "value"},
|
|
)
|
|
|
|
# Verify handler received the context
|
|
assert received_params is not None
|
|
assert received_params.context is not None
|
|
assert received_params.context.arguments == {"previous": "value"}
|
|
assert result.completion.values == ["test-completion"]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_completion_backward_compatibility():
|
|
"""Test that completion works without context (backward compatibility)."""
|
|
context_was_none = False
|
|
|
|
async def handle_completion(ctx: ServerRequestContext, params: CompleteRequestParams) -> CompleteResult:
|
|
nonlocal context_was_none
|
|
context_was_none = params.context is None
|
|
return CompleteResult(completion=Completion(values=["no-context-completion"], total=1, has_more=False))
|
|
|
|
server = Server("test-server", on_completion=handle_completion)
|
|
|
|
async with Client(server) as client:
|
|
# Test without context
|
|
result = await client.complete(
|
|
ref=PromptReference(type="ref/prompt", name="test-prompt"), argument={"name": "arg", "value": "val"}
|
|
)
|
|
|
|
# Verify context was None
|
|
assert context_was_none
|
|
assert result.completion.values == ["no-context-completion"]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_dependent_completion_scenario():
|
|
"""Test a real-world scenario with dependent completions."""
|
|
|
|
async def handle_completion(ctx: ServerRequestContext, params: CompleteRequestParams) -> CompleteResult:
|
|
# Simulate database/table completion scenario
|
|
assert isinstance(params.ref, ResourceTemplateReference)
|
|
assert params.ref.uri == "db://{database}/{table}"
|
|
|
|
if params.argument.name == "database":
|
|
return CompleteResult(
|
|
completion=Completion(values=["users_db", "products_db", "analytics_db"], total=3, has_more=False)
|
|
)
|
|
|
|
assert params.argument.name == "table"
|
|
assert params.context and params.context.arguments
|
|
db = params.context.arguments.get("database")
|
|
if db == "users_db":
|
|
return CompleteResult(
|
|
completion=Completion(values=["users", "sessions", "permissions"], total=3, has_more=False)
|
|
)
|
|
else:
|
|
assert db == "products_db"
|
|
return CompleteResult(
|
|
completion=Completion(values=["products", "categories", "inventory"], total=3, has_more=False)
|
|
)
|
|
|
|
server = Server("test-server", on_completion=handle_completion)
|
|
|
|
async with Client(server) as client:
|
|
# First, complete database
|
|
db_result = await client.complete(
|
|
ref=ResourceTemplateReference(type="ref/resource", uri="db://{database}/{table}"),
|
|
argument={"name": "database", "value": ""},
|
|
)
|
|
assert "users_db" in db_result.completion.values
|
|
assert "products_db" in db_result.completion.values
|
|
|
|
# Then complete table with database context
|
|
table_result = await client.complete(
|
|
ref=ResourceTemplateReference(type="ref/resource", uri="db://{database}/{table}"),
|
|
argument={"name": "table", "value": ""},
|
|
context_arguments={"database": "users_db"},
|
|
)
|
|
assert table_result.completion.values == ["users", "sessions", "permissions"]
|
|
|
|
# Different database gives different tables
|
|
table_result2 = await client.complete(
|
|
ref=ResourceTemplateReference(type="ref/resource", uri="db://{database}/{table}"),
|
|
argument={"name": "table", "value": ""},
|
|
context_arguments={"database": "products_db"},
|
|
)
|
|
assert table_result2.completion.values == ["products", "categories", "inventory"]
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_completion_error_on_missing_context():
|
|
"""Test that server can raise error when required context is missing."""
|
|
|
|
async def handle_completion(ctx: ServerRequestContext, params: CompleteRequestParams) -> CompleteResult:
|
|
assert isinstance(params.ref, ResourceTemplateReference)
|
|
assert params.ref.uri == "db://{database}/{table}"
|
|
assert params.argument.name == "table"
|
|
|
|
if not params.context or not params.context.arguments or "database" not in params.context.arguments:
|
|
# A message the client should see travels as an MCPError; a bare
|
|
# exception would reach the wire only as an opaque internal error.
|
|
raise MCPError(code=INVALID_PARAMS, message="Please select a database first to see available tables")
|
|
|
|
db = params.context.arguments.get("database")
|
|
assert db == "test_db"
|
|
return CompleteResult(completion=Completion(values=["users", "orders", "products"], total=3, has_more=False))
|
|
|
|
server = Server("test-server", on_completion=handle_completion)
|
|
|
|
async with Client(server, mode="legacy") as client:
|
|
# Try to complete table without database context - should raise error
|
|
with pytest.raises(Exception) as exc_info:
|
|
await client.complete(
|
|
ref=ResourceTemplateReference(type="ref/resource", uri="db://{database}/{table}"),
|
|
argument={"name": "table", "value": ""},
|
|
)
|
|
|
|
# Verify error message
|
|
assert "Please select a database first" in str(exc_info.value)
|
|
|
|
# Now complete with proper context - should work normally
|
|
result_with_context = await client.complete(
|
|
ref=ResourceTemplateReference(type="ref/resource", uri="db://{database}/{table}"),
|
|
argument={"name": "table", "value": ""},
|
|
context_arguments={"database": "test_db"},
|
|
)
|
|
|
|
# Should get normal completions
|
|
assert result_with_context.completion.values == ["users", "orders", "products"]
|