[v2] ClientSession runs on JSONRPCDispatcher; BaseSession removed (#2838)

tests/server/mcpserver/test_elicitation.py

@@ -6,9 +6,9 @@ import pytest
 from pydantic import BaseModel, Field
 
 from mcp import Client, types
-from mcp.client.session import ClientSession, ElicitationFnT
+from mcp.client import ClientRequestContext
+from mcp.client.session import ElicitationFnT
 from mcp.server.mcpserver import Context, MCPServer
-from mcp.shared._context import RequestContext
 from mcp.types import ElicitRequestParams, ElicitResult, TextContent
 
 
@@ -64,7 +64,7 @@ async def test_elicitation_accept_returns_the_users_answer_to_the_tool():
     create_ask_user_tool(mcp)
 
     # Create a custom handler for elicitation requests
-    async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
+    async def elicitation_callback(context: ClientRequestContext, params: ElicitRequestParams):
         if params.message == "Tool wants to ask: What is your name?":
             return ElicitResult(action="accept", content={"answer": "Test User"})
         else:  # pragma: no cover
@@ -81,7 +81,7 @@ async def test_elicitation_decline_reaches_the_tool_without_content():
     mcp = MCPServer(name="ElicitationDeclineServer")
     create_ask_user_tool(mcp)
 
-    async def elicitation_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
+    async def elicitation_callback(context: ClientRequestContext, params: ElicitRequestParams):
         return ElicitResult(action="decline")
 
     await call_tool_and_assert(
@@ -119,9 +119,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: RequestContext[ClientSession], params: ElicitRequestParams
-    ):  # pragma: no cover
+    async def elicitation_callback(context: ClientRequestContext, params: ElicitRequestParams):  # pragma: no cover
         return ElicitResult(action="accept", content={})
 
     async with Client(mcp, elicitation_callback=elicitation_callback) as client:
@@ -176,7 +174,7 @@ async def test_elicitation_with_optional_fields():
 
     for content, expected in test_cases:
 
-        async def callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
+        async def callback(context: ClientRequestContext, params: ElicitRequestParams):
             return ElicitResult(action="accept", content=content)
 
         await call_tool_and_assert(mcp, callback, "optional_tool", {}, expected)
@@ -194,9 +192,7 @@ async def test_elicitation_with_optional_fields():
         except TypeError as e:
             return f"Validation failed: {str(e)}"
 
-    async def elicitation_callback(
-        context: RequestContext[ClientSession], params: ElicitRequestParams
-    ):  # pragma: no cover
+    async def elicitation_callback(context: ClientRequestContext, params: ElicitRequestParams):  # pragma: no cover
         return ElicitResult(action="accept", content={})
 
     await call_tool_and_assert(
@@ -219,7 +215,7 @@ async def test_elicitation_with_optional_fields():
             return f"Name: {result.data.name}, Tags: {', '.join(result.data.tags)}"
         return f"User {result.action}"  # pragma: no cover
 
-    async def multiselect_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
+    async def multiselect_callback(context: ClientRequestContext, params: ElicitRequestParams):
         if "Please provide tags" in params.message:
             return ElicitResult(action="accept", content={"name": "Test", "tags": ["tag1", "tag2"]})
         return ElicitResult(action="decline")  # pragma: no cover
@@ -239,7 +235,7 @@ async def test_elicitation_with_optional_fields():
             return f"Name: {result.data.name}, Tags: {tags_str}"
         return f"User {result.action}"  # pragma: no cover
 
-    async def optional_multiselect_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
+    async def optional_multiselect_callback(context: ClientRequestContext, params: ElicitRequestParams):
         if "Please provide optional tags" in params.message:
             return ElicitResult(action="accept", content={"name": "Test", "tags": ["tag1", "tag2"]})
         return ElicitResult(action="decline")  # pragma: no cover
@@ -273,7 +269,7 @@ async def test_elicitation_with_default_values():
             return f"User {result.action}"
 
     # First verify that defaults are present in the JSON schema sent to clients
-    async def callback_schema_verify(context: RequestContext[ClientSession], params: ElicitRequestParams):
+    async def callback_schema_verify(context: ClientRequestContext, params: ElicitRequestParams):
         # Verify the schema includes defaults
         assert isinstance(params, types.ElicitRequestFormParams), "Expected form mode elicitation"
         schema = params.requested_schema
@@ -295,7 +291,7 @@ async def test_elicitation_with_default_values():
     )
 
     # Test overriding defaults
-    async def callback_override(context: RequestContext[ClientSession], params: ElicitRequestParams):
+    async def callback_override(context: ClientRequestContext, params: ElicitRequestParams):
         return ElicitResult(
             action="accept", content={"email": "john@example.com", "name": "John", "age": 25, "subscribe": False}
         )
@@ -371,7 +367,7 @@ async def test_elicitation_with_enum_titles():
             return f"User: {result.data.user_name}, Color: {result.data.color}"
         return f"User {result.action}"  # pragma: no cover
 
-    async def enum_callback(context: RequestContext[ClientSession], params: ElicitRequestParams):
+    async def enum_callback(context: ClientRequestContext, params: ElicitRequestParams):
         if "colors" in params.message and "legacy" not in params.message:
             return ElicitResult(action="accept", content={"user_name": "Bob", "favorite_colors": ["red", "green"]})
         elif "color" in params.message: