Skip to content

Commit 032b630

Browse files
committed
add elicitation test using create_client_server_memory_streams
1 parent 3599de8 commit 032b630

File tree

3 files changed

+63
-1
lines changed

3 files changed

+63
-1
lines changed

src/mcp/shared/memory.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import mcp.types as types
1414
from mcp.client.session import (
1515
ClientSession,
16+
ElicitationFnT,
1617
ListRootsFnT,
1718
LoggingFnT,
1819
MessageHandlerFnT,
@@ -68,6 +69,7 @@ async def create_connected_server_and_client_session(
6869
message_handler: MessageHandlerFnT | None = None,
6970
client_info: types.Implementation | None = None,
7071
raise_exceptions: bool = False,
72+
elicitation_callback: ElicitationFnT | None = None,
7173
) -> AsyncGenerator[ClientSession, None]:
7274
"""Creates a ClientSession that is connected to a running MCP server."""
7375
async with create_client_server_memory_streams() as (
@@ -98,6 +100,7 @@ async def create_connected_server_and_client_session(
98100
logging_callback=logging_callback,
99101
message_handler=message_handler,
100102
client_info=client_info,
103+
elicitation_callback=elicitation_callback,
101104
) as client_session:
102105
await client_session.initialize()
103106
yield client_session

tests/server/fastmcp/test_integration.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
from mcp.client.session import ClientSession
1717
from mcp.client.sse import sse_client
1818
from mcp.server.fastmcp import Context, FastMCP
19-
from mcp.types import InitializeResult, TextContent, ElicitResult
19+
from mcp.types import ElicitResult, InitializeResult, TextContent
2020

2121

2222
@pytest.fixture
Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
"""
2+
Test the elicitation feature using stdio transport.
3+
"""
4+
5+
import pytest
6+
7+
from mcp.server.fastmcp import Context, FastMCP
8+
from mcp.shared.memory import create_connected_server_and_client_session
9+
from mcp.types import ElicitResult, TextContent
10+
11+
12+
@pytest.mark.anyio
13+
async def test_stdio_elicitation():
14+
"""Test the elicitation feature using stdio transport."""
15+
16+
# Create a FastMCP server with a tool that uses elicitation
17+
mcp = FastMCP(name="StdioElicitationServer")
18+
19+
@mcp.tool(description="A tool that uses elicitation")
20+
async def ask_user(prompt: str, ctx: Context) -> str:
21+
schema = {
22+
"type": "object",
23+
"properties": {
24+
"answer": {"type": "string"},
25+
},
26+
"required": ["answer"],
27+
}
28+
29+
response = await ctx.elicit(
30+
message=f"Tool wants to ask: {prompt}",
31+
requestedSchema=schema,
32+
)
33+
return f"User answered: {response['answer']}"
34+
35+
# Create a custom handler for elicitation requests
36+
async def elicitation_callback(context, params):
37+
# Verify the elicitation parameters
38+
if params.message == "Tool wants to ask: What is your name?":
39+
return ElicitResult(response={"answer": "Test User"})
40+
else:
41+
raise ValueError(f"Unexpected elicitation message: {params.message}")
42+
43+
# Use memory-based session to test with stdio transport
44+
async with create_connected_server_and_client_session(
45+
mcp._mcp_server, elicitation_callback=elicitation_callback
46+
) as client_session:
47+
# First initialize the session
48+
result = await client_session.initialize()
49+
assert result.serverInfo.name == "StdioElicitationServer"
50+
51+
# Call the tool that uses elicitation
52+
tool_result = await client_session.call_tool(
53+
"ask_user", {"prompt": "What is your name?"}
54+
)
55+
56+
# Verify the result
57+
assert len(tool_result.content) == 1
58+
assert isinstance(tool_result.content[0], TextContent)
59+
assert tool_result.content[0].text == "User answered: Test User"

0 commit comments

Comments
 (0)