Add pyright strict mode on the whole project (#1254)

tests/server/fastmcp/test_elicitation.py

@@ -2,12 +2,17 @@
 Test the elicitation feature using stdio transport.
 """
 
+from typing import Any
+
 import pytest
 from pydantic import BaseModel, Field
 
+from mcp.client.session import ClientSession, ElicitationFnT
 from mcp.server.fastmcp import Context, FastMCP
+from mcp.server.session import ServerSession
+from mcp.shared.context import RequestContext
 from mcp.shared.memory import create_connected_server_and_client_session
-from mcp.types import ElicitResult, TextContent
+from mcp.types import ElicitRequestParams, ElicitResult, TextContent
 
 
 # Shared schema for basic tests
@@ -19,11 +24,8 @@ def create_ask_user_tool(mcp: FastMCP):
     """Create a standard ask_user tool that handles all elicitation responses."""
 
     @mcp.tool(description="A tool that uses elicitation")
-    async def ask_user(prompt: str, ctx: Context) -> str:
-        result = await ctx.elicit(
-            message=f"Tool wants to ask: {prompt}",
-            schema=AnswerSchema,
-        )
+    async def ask_user(prompt: str, ctx: Context[ServerSession, None]) -> str:
+        result = await ctx.elicit(message=f"Tool wants to ask: {prompt}", schema=AnswerSchema)
 
         if result.action == "accept" and result.data:
             return f"User answered: {result.data.answer}"
@@ -37,9 +39,9 @@ def create_ask_user_tool(mcp: FastMCP):
 
 async def call_tool_and_assert(
     mcp: FastMCP,
-    elicitation_callback,
+    elicitation_callback: ElicitationFnT,
     tool_name: str,
-    args: dict,
+    args: dict[str, Any],
     expected_text: str | None = None,
     text_contains: list[str] | None = None,
 ):
@@ -69,7 +71,7 @@ async def test_stdio_elicitation():
     create_ask_user_tool(mcp)
 
     # Create a custom handler for elicitation requests
-    async def elicitation_callback(context, params):
+    async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
         if params.message == "Tool wants to ask: What is your name?":
             return ElicitResult(action="accept", content={"answer": "Test User"})
         else:
@@ -86,7 +88,7 @@ async def test_stdio_elicitation_decline():
     mcp = FastMCP(name="StdioElicitationDeclineServer")
     create_ask_user_tool(mcp)
 
-    async def elicitation_callback(context, params):
+    async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
         return ElicitResult(action="decline")
 
     await call_tool_and_assert(
@@ -101,7 +103,7 @@ async def test_elicitation_schema_validation():
 
     def create_validation_tool(name: str, schema_class: type[BaseModel]):
         @mcp.tool(name=name, description=f"Tool testing {name}")
-        async def tool(ctx: Context) -> str:
+        async def tool(ctx: Context[ServerSession, None]) -> str:
             try:
                 await ctx.elicit(message="This should fail validation", schema=schema_class)
                 return "Should not reach here"
@@ -124,7 +126,7 @@ async def test_elicitation_schema_validation():
     create_validation_tool("nested_model", InvalidNestedSchema)
 
     # Dummy callback (won't be called due to validation failure)
-    async def elicitation_callback(context, params):
+    async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
         return ElicitResult(action="accept", content={})
 
     async with create_connected_server_and_client_session(
@@ -153,7 +155,7 @@ async def test_elicitation_with_optional_fields():
         subscribe: bool | None = Field(default=False, description="Subscribe to newsletter?")
 
     @mcp.tool(description="Tool with optional fields")
-    async def optional_tool(ctx: Context) -> str:
+    async def optional_tool(ctx: Context[ServerSession, None]) -> str:
         result = await ctx.elicit(message="Please provide your information", schema=OptionalSchema)
 
         if result.action == "accept" and result.data:
@@ -168,7 +170,7 @@ async def test_elicitation_with_optional_fields():
             return f"User {result.action}"
 
     # Test cases with different field combinations
-    test_cases = [
+    test_cases: list[tuple[dict[str, Any], str]] = [
         (
             # All fields provided
             {"required_name": "John Doe", "optional_age": 30, "optional_email": "john@example.com", "subscribe": True},
@@ -183,7 +185,7 @@ async def test_elicitation_with_optional_fields():
 
     for content, expected in test_cases:
 
-        async def callback(context, params):
+        async def callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
             return ElicitResult(action="accept", content=content)
 
         await call_tool_and_assert(mcp, callback, "optional_tool", {}, expected)
@@ -194,16 +196,19 @@ async def test_elicitation_with_optional_fields():
         optional_list: list[str] | None = Field(default=None, description="Invalid optional list")
 
     @mcp.tool(description="Tool with invalid optional field")
-    async def invalid_optional_tool(ctx: Context) -> str:
+    async def invalid_optional_tool(ctx: Context[ServerSession, None]) -> str:
         try:
             await ctx.elicit(message="This should fail", schema=InvalidOptionalSchema)
             return "Should not reach here"
         except TypeError as e:
             return f"Validation failed: {str(e)}"
 
+    async def elicitation_callback(context: RequestContext[ClientSession, None], params: ElicitRequestParams):
+        return ElicitResult(action="accept", content={})
+
     await call_tool_and_assert(
         mcp,
-        lambda c, p: ElicitResult(action="accept", content={}),
+        elicitation_callback,
         "invalid_optional_tool",
         {},
         text_contains=["Validation failed:", "optional_list"],