From c3d47a0ed28443240b92be788a56c4982feebc5d Mon Sep 17 00:00:00 2001
From: ihrpr <inna@anthropic.com>
Date: Wed, 9 Jul 2025 09:37:13 +0100
Subject: [PATCH] fix completions

---
 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 d00eda875..154c3a368 100644
--- a/tests/server/test_session.py
+++ b/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"],
         )