77
88import asyncio
99import logging
10- from typing import Any
1110
1211import click
1312from mcp .server .fastmcp import Context , FastMCP
1413from mcp .server .fastmcp .prompts .base import UserMessage
1514from mcp .server .session import ServerSession
1615from mcp .types import (
1716 AudioContent ,
17+ Completion ,
1818 CompletionArgument ,
1919 CompletionContext ,
2020 EmbeddedResource ,
2525 TextContent ,
2626 TextResourceContents ,
2727)
28- from pydantic import BaseModel , Field
28+ from pydantic import AnyUrl , BaseModel , Field
2929
3030# Configure logging
3131logger = logging .getLogger (__name__ )
@@ -79,7 +79,7 @@ def test_embedded_resource() -> list[EmbeddedResource]:
7979 EmbeddedResource (
8080 type = "resource" ,
8181 resource = TextResourceContents (
82- uri = "test://embedded-resource" ,
82+ uri = AnyUrl ( "test://embedded-resource" ) ,
8383 mimeType = "text/plain" ,
8484 text = "This is an embedded resource content." ,
8585 ),
@@ -96,7 +96,7 @@ def test_multiple_content_types() -> list[TextContent | ImageContent | EmbeddedR
9696 EmbeddedResource (
9797 type = "resource" ,
9898 resource = TextResourceContents (
99- uri = "test://mixed-content-resource" ,
99+ uri = AnyUrl ( "test://mixed-content-resource" ) ,
100100 mimeType = "application/json" ,
101101 text = '{"test": "data", "value": 123}' ,
102102 ),
@@ -105,7 +105,7 @@ def test_multiple_content_types() -> list[TextContent | ImageContent | EmbeddedR
105105
106106
107107@mcp .tool ()
108- async def test_tool_with_logging (ctx : Context ) -> str :
108+ async def test_tool_with_logging (ctx : Context [ ServerSession , None ] ) -> str :
109109 """Tests tool that emits log messages during execution"""
110110 await ctx .info ("Tool execution started" )
111111 await asyncio .sleep (0.05 )
@@ -118,7 +118,7 @@ async def test_tool_with_logging(ctx: Context) -> str:
118118
119119
120120@mcp .tool ()
121- async def test_tool_with_progress (ctx : Context ) -> str :
121+ async def test_tool_with_progress (ctx : Context [ ServerSession , None ] ) -> str :
122122 """Tests tool that reports progress notifications"""
123123 await ctx .report_progress (progress = 0 , total = 100 , message = "Completed step 0 of 100" )
124124 await asyncio .sleep (0.05 )
@@ -178,9 +178,13 @@ async def test_elicitation(message: str, ctx: Context[ServerSession, None]) -> s
178178 # Request user input from client
179179 result = await ctx .elicit (message = message , schema = UserResponse )
180180
181- return (
182- f"User response: action={ result .action } , content={ result .data .model_dump_json () if result .data else '{}' } "
183- )
181+ # Type-safe discriminated union narrowing using action field
182+ if result .action == "accept" :
183+ content = result .data .model_dump_json ()
184+ else : # decline or cancel
185+ content = "{}"
186+
187+ return f"User response: action={ result .action } , content={ content } "
184188 except Exception as e :
185189 return f"Elicitation not supported or error: { str (e )} "
186190
@@ -244,7 +248,7 @@ def test_prompt_with_embedded_resource(resourceUri: str) -> list[UserMessage]:
244248 content = EmbeddedResource (
245249 type = "resource" ,
246250 resource = TextResourceContents (
247- uri = resourceUri ,
251+ uri = AnyUrl ( resourceUri ) ,
248252 mimeType = "text/plain" ,
249253 text = "Embedded resource content for testing." ,
250254 ),
@@ -278,34 +282,32 @@ async def handle_set_logging_level(level: str) -> None:
278282 # For conformance testing, we just acknowledge the request
279283
280284 # Handler for resources/subscribe
281- async def handle_subscribe (uri : str ) -> dict [ str , Any ] :
285+ async def handle_subscribe (uri : AnyUrl ) -> None :
282286 """Handle resource subscription"""
283- resource_subscriptions .add (uri )
287+ resource_subscriptions .add (str ( uri ) )
284288 logger .info (f"Subscribed to resource: { uri } " )
285- return {}
286289
287290 # Handler for resources/unsubscribe
288- async def handle_unsubscribe (uri : str ) -> dict [ str , Any ] :
291+ async def handle_unsubscribe (uri : AnyUrl ) -> None :
289292 """Handle resource unsubscription"""
290- resource_subscriptions .discard (uri )
293+ resource_subscriptions .discard (str ( uri ) )
291294 logger .info (f"Unsubscribed from resource: { uri } " )
292- return {}
293295
294296 # Register subscription handlers
295297 mcp ._mcp_server .subscribe_resource ()(handle_subscribe )
296298 mcp ._mcp_server .unsubscribe_resource ()(handle_unsubscribe )
297299
298300 # Handler for completion/complete
299- @mcp ._mcp_server . completion ()
301+ @mcp .completion ()
300302 async def handle_completion (
301303 ref : PromptReference | ResourceTemplateReference ,
302304 argument : CompletionArgument ,
303305 context : CompletionContext | None ,
304- ) -> dict [ str , Any ] :
306+ ) -> Completion :
305307 """Handle completion requests"""
306308 # Basic completion support - returns empty array for conformance
307309 # Real implementations would provide contextual suggestions
308- return { " values" : [], " total" : 0 , " hasMore" : False }
310+ return Completion ( values = [], total = 0 , hasMore = False )
309311
310312
311313# Set up custom handlers when module is loaded
0 commit comments