Fix async resource functions not being awaited (#1092)
Co-authored-by: Marcelo Trylesinski <marcelotryle@gmail.com>
This commit is contained in:
parent
3fa7bc7b61
commit
3d67e723e5
2 changed files with 47 additions and 2 deletions
|
|
@ -54,7 +54,12 @@ class FunctionResource(Resource):
|
||||||
async def read(self) -> str | bytes:
|
async def read(self) -> str | bytes:
|
||||||
"""Read the resource by calling the wrapped function."""
|
"""Read the resource by calling the wrapped function."""
|
||||||
try:
|
try:
|
||||||
result = await self.fn() if inspect.iscoroutinefunction(self.fn) else self.fn()
|
# Call the function first to see if it returns a coroutine
|
||||||
|
result = self.fn()
|
||||||
|
# If it's a coroutine, await it
|
||||||
|
if inspect.iscoroutine(result):
|
||||||
|
result = await result
|
||||||
|
|
||||||
if isinstance(result, Resource):
|
if isinstance(result, Resource):
|
||||||
return await result.read()
|
return await result.read()
|
||||||
elif isinstance(result, bytes):
|
elif isinstance(result, bytes):
|
||||||
|
|
|
||||||
|
|
@ -1,12 +1,13 @@
|
||||||
import anyio
|
import anyio
|
||||||
import pytest
|
import pytest
|
||||||
|
from pydantic import AnyUrl
|
||||||
|
|
||||||
from mcp.server.fastmcp import FastMCP
|
from mcp.server.fastmcp import FastMCP
|
||||||
from mcp.shared.memory import create_connected_server_and_client_session as create_session
|
from mcp.shared.memory import create_connected_server_and_client_session as create_session
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.anyio
|
@pytest.mark.anyio
|
||||||
async def test_messages_are_executed_concurrently():
|
async def test_messages_are_executed_concurrently_tools():
|
||||||
server = FastMCP("test")
|
server = FastMCP("test")
|
||||||
event = anyio.Event()
|
event = anyio.Event()
|
||||||
tool_started = anyio.Event()
|
tool_started = anyio.Event()
|
||||||
|
|
@ -44,3 +45,42 @@ async def test_messages_are_executed_concurrently():
|
||||||
"trigger_end",
|
"trigger_end",
|
||||||
"tool_end",
|
"tool_end",
|
||||||
], f"Expected concurrent execution, but got: {call_order}"
|
], f"Expected concurrent execution, but got: {call_order}"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_messages_are_executed_concurrently_tools_and_resources():
|
||||||
|
server = FastMCP("test")
|
||||||
|
event = anyio.Event()
|
||||||
|
tool_started = anyio.Event()
|
||||||
|
call_order = []
|
||||||
|
|
||||||
|
@server.tool("sleep")
|
||||||
|
async def sleep_tool():
|
||||||
|
call_order.append("waiting_for_event")
|
||||||
|
tool_started.set()
|
||||||
|
await event.wait()
|
||||||
|
call_order.append("tool_end")
|
||||||
|
return "done"
|
||||||
|
|
||||||
|
@server.resource("slow://slow_resource")
|
||||||
|
async def slow_resource():
|
||||||
|
# Wait for tool to start before setting the event
|
||||||
|
await tool_started.wait()
|
||||||
|
event.set()
|
||||||
|
call_order.append("resource_end")
|
||||||
|
return "slow"
|
||||||
|
|
||||||
|
async with create_session(server._mcp_server) as client_session:
|
||||||
|
# First tool will wait on event, second will set it
|
||||||
|
async with anyio.create_task_group() as tg:
|
||||||
|
# Start the tool first (it will wait on event)
|
||||||
|
tg.start_soon(client_session.call_tool, "sleep")
|
||||||
|
# Then the resource (it will set the event)
|
||||||
|
tg.start_soon(client_session.read_resource, AnyUrl("slow://slow_resource"))
|
||||||
|
|
||||||
|
# Verify that both ran concurrently
|
||||||
|
assert call_order == [
|
||||||
|
"waiting_for_event",
|
||||||
|
"resource_end",
|
||||||
|
"tool_end",
|
||||||
|
], f"Expected concurrent execution, but got: {call_order}"
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue