Fix completions (modelcontextprotocol#1107)

Author: ihrpr

tests/server/test_session.py

@@ -13,12 +13,13 @@
     ClientNotification,
     Completion,
     CompletionArgument,
+    CompletionContext,
     CompletionsCapability,
     InitializedNotification,
     PromptReference,
     PromptsCapability,
-    ResourceReference,
     ResourcesCapability,
+    ResourceTemplateReference,
     ServerCapabilities,
 )
 
@@ -109,7 +110,11 @@ async def list_resources():
 
     # Add a complete handler
     @server.completion()
-    async def complete(ref: PromptReference | ResourceReference, argument: CompletionArgument):
+    async def complete(
+        ref: PromptReference | ResourceTemplateReference,
+        argument: CompletionArgument,
+        context: CompletionContext | None,
+    ) -> Completion | None:
         return Completion(
             values=["completion1", "completion2"],
         )