Marcelo Trylesinski
210 lines
8.2 KiB
Python
210 lines
8.2 KiB
Python
"""Tests for completion handler with context functionality."""
|
|
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
import mcp.types as types
|
|
from mcp import Client
|
|
from mcp.server.lowlevel import Server
|
|
from mcp.server.session import ServerSession
|
|
from mcp.shared.context import RequestContext
|
|
from mcp.types import (
|
|
Completion,
|
|
CompletionArgument,
|
|
CompletionContext,
|
|
PromptReference,
|
|
ResourceTemplateReference,
|
|
)
|
|
|
|
|
|
@pytest.mark.anyio
|
|
async def test_completion_handler_receives_context():
|
|
"""Test that the completion handler receives context correctly."""
|
|
|
|
# Track what the handler receives
|
|
received_args: dict[str, Any] = {}
|
|
|
|
async def on_completion(
|
|
ctx: RequestContext[ServerSession, Any, Any],
|
|
params: types.CompleteRequestParams,
|
|
) -> types.CompleteResult:
|
|
ref = params.ref
|
|
argument = params.argument
|
|
context = params.context
|
|
received_args["ref"] = ref
|
|
received_args["argument"] = argument
|
|
received_args["context"] = context
|
|
|
|
# Return test completion
|
|
return types.CompleteResult(
|
|
completion=Completion(values=["test-completion"], total=1, has_more=False)
|
|
)
|
|
|
|
server = Server("test-server", on_completion=on_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_args["context"] is not None
|
|
assert received_args["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 on_completion(
|
|
ctx: RequestContext[ServerSession, Any, Any],
|
|
params: types.CompleteRequestParams,
|
|
) -> types.CompleteResult:
|
|
nonlocal context_was_none
|
|
context_was_none = params.context is None
|
|
|
|
return types.CompleteResult(
|
|
completion=Completion(values=["no-context-completion"], total=1, has_more=False)
|
|
)
|
|
|
|
server = Server("test-server", on_completion=on_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 on_completion(
|
|
ctx: RequestContext[ServerSession, Any, Any],
|
|
params: types.CompleteRequestParams,
|
|
) -> types.CompleteResult:
|
|
ref = params.ref
|
|
argument = params.argument
|
|
context = params.context
|
|
# Simulate database/table completion scenario
|
|
if isinstance(ref, ResourceTemplateReference):
|
|
if ref.uri == "db://{database}/{table}":
|
|
if argument.name == "database":
|
|
# Complete database names
|
|
return types.CompleteResult(
|
|
completion=Completion(
|
|
values=["users_db", "products_db", "analytics_db"], total=3, has_more=False
|
|
)
|
|
)
|
|
elif argument.name == "table":
|
|
# Complete table names based on selected database
|
|
if context and context.arguments:
|
|
db = context.arguments.get("database")
|
|
if db == "users_db":
|
|
return types.CompleteResult(
|
|
completion=Completion(
|
|
values=["users", "sessions", "permissions"], total=3, has_more=False
|
|
)
|
|
)
|
|
elif db == "products_db":
|
|
return types.CompleteResult(
|
|
completion=Completion(
|
|
values=["products", "categories", "inventory"], total=3, has_more=False
|
|
)
|
|
)
|
|
|
|
return types.CompleteResult(
|
|
completion=Completion(values=[], total=0, has_more=False)
|
|
) # pragma: no cover
|
|
|
|
server = Server("test-server", on_completion=on_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 on_completion(
|
|
ctx: RequestContext[ServerSession, Any, Any],
|
|
params: types.CompleteRequestParams,
|
|
) -> types.CompleteResult:
|
|
ref = params.ref
|
|
argument = params.argument
|
|
context = params.context
|
|
if isinstance(ref, ResourceTemplateReference):
|
|
if ref.uri == "db://{database}/{table}":
|
|
if argument.name == "table":
|
|
# Check if database context is provided
|
|
if not context or not context.arguments or "database" not in context.arguments:
|
|
# Raise an error instead of returning error as completion
|
|
raise ValueError("Please select a database first to see available tables")
|
|
# Normal completion if context is provided
|
|
db = context.arguments.get("database")
|
|
if db == "test_db":
|
|
return types.CompleteResult(
|
|
completion=Completion(values=["users", "orders", "products"], total=3, has_more=False)
|
|
)
|
|
|
|
return types.CompleteResult(
|
|
completion=Completion(values=[], total=0, has_more=False)
|
|
) # pragma: no cover
|
|
|
|
server = Server("test-server", on_completion=on_completion)
|
|
|
|
async with Client(server) 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"]
|