From c067ef291e9f73baff6e40b9e6895e00eff8f6e3 Mon Sep 17 00:00:00 2001
From: Inna Harper <inna.hrpr@gmail.com>
Date: Wed, 9 Jul 2025 11:15:49 +0100
Subject: [PATCH] Fix completions (#1107)

---
 tests/server/test_session.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/tests/server/test_session.py b/tests/server/test_session.py
index d00eda87..154c3a36 100644
--- a/tests/server/test_session.py
+++ b/tests/server/test_session.py
@@ -13,12 +13,13 @@ from mcp.types import (
     ClientNotification,
     Completion,
     CompletionArgument,
+    CompletionContext,
     CompletionsCapability,
     InitializedNotification,
     PromptReference,
     PromptsCapability,
-    ResourceReference,
     ResourcesCapability,
+    ResourceTemplateReference,
     ServerCapabilities,
 )
 
@@ -109,7 +110,11 @@ async def test_server_capabilities():
 
     # 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"],
         )