Don't block the event loop on sync resource and prompt functions (#2380)
This commit is contained in:
parent
e6235d1667
commit
3ce0f76e6e
6 changed files with 107 additions and 13 deletions
|
|
@ -2,10 +2,12 @@
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import functools
|
||||||
import inspect
|
import inspect
|
||||||
from collections.abc import Awaitable, Callable, Sequence
|
from collections.abc import Awaitable, Callable, Sequence
|
||||||
from typing import TYPE_CHECKING, Any, Literal
|
from typing import TYPE_CHECKING, Any, Literal
|
||||||
|
|
||||||
|
import anyio.to_thread
|
||||||
import pydantic_core
|
import pydantic_core
|
||||||
from pydantic import BaseModel, Field, TypeAdapter, validate_call
|
from pydantic import BaseModel, Field, TypeAdapter, validate_call
|
||||||
|
|
||||||
|
|
@ -155,10 +157,10 @@ class Prompt(BaseModel):
|
||||||
# Add context to arguments if needed
|
# Add context to arguments if needed
|
||||||
call_args = inject_context(self.fn, arguments or {}, context, self.context_kwarg)
|
call_args = inject_context(self.fn, arguments or {}, context, self.context_kwarg)
|
||||||
|
|
||||||
# Call function and check if result is a coroutine
|
if inspect.iscoroutinefunction(self.fn):
|
||||||
result = self.fn(**call_args)
|
result = await self.fn(**call_args)
|
||||||
if inspect.iscoroutine(result):
|
else:
|
||||||
result = await result
|
result = await anyio.to_thread.run_sync(functools.partial(self.fn, **call_args))
|
||||||
|
|
||||||
# Validate messages
|
# Validate messages
|
||||||
if not isinstance(result, list | tuple):
|
if not isinstance(result, list | tuple):
|
||||||
|
|
|
||||||
|
|
@ -2,12 +2,14 @@
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import functools
|
||||||
import inspect
|
import inspect
|
||||||
import re
|
import re
|
||||||
from collections.abc import Callable
|
from collections.abc import Callable
|
||||||
from typing import TYPE_CHECKING, Any
|
from typing import TYPE_CHECKING, Any
|
||||||
from urllib.parse import unquote
|
from urllib.parse import unquote
|
||||||
|
|
||||||
|
import anyio.to_thread
|
||||||
from pydantic import BaseModel, Field, validate_call
|
from pydantic import BaseModel, Field, validate_call
|
||||||
|
|
||||||
from mcp.server.mcpserver.resources.types import FunctionResource, Resource
|
from mcp.server.mcpserver.resources.types import FunctionResource, Resource
|
||||||
|
|
@ -110,10 +112,10 @@ class ResourceTemplate(BaseModel):
|
||||||
# Add context to params if needed
|
# Add context to params if needed
|
||||||
params = inject_context(self.fn, params, context, self.context_kwarg)
|
params = inject_context(self.fn, params, context, self.context_kwarg)
|
||||||
|
|
||||||
# Call function and check if result is a coroutine
|
if inspect.iscoroutinefunction(self.fn):
|
||||||
result = self.fn(**params)
|
result = await self.fn(**params)
|
||||||
if inspect.iscoroutine(result):
|
else:
|
||||||
result = await result
|
result = await anyio.to_thread.run_sync(functools.partial(self.fn, **params))
|
||||||
|
|
||||||
return FunctionResource(
|
return FunctionResource(
|
||||||
uri=uri, # type: ignore
|
uri=uri, # type: ignore
|
||||||
|
|
|
||||||
|
|
@ -55,11 +55,10 @@ 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:
|
||||||
# Call the function first to see if it returns a coroutine
|
if inspect.iscoroutinefunction(self.fn):
|
||||||
result = self.fn()
|
result = await self.fn()
|
||||||
# If it's a coroutine, await it
|
else:
|
||||||
if inspect.iscoroutine(result):
|
result = await anyio.to_thread.run_sync(self.fn)
|
||||||
result = await result
|
|
||||||
|
|
||||||
if isinstance(result, Resource): # pragma: no cover
|
if isinstance(result, Resource): # pragma: no cover
|
||||||
return await result.read()
|
return await result.read()
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,4 @@
|
||||||
|
import threading
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
@ -190,3 +191,21 @@ class TestRenderPrompt:
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_sync_fn_runs_in_worker_thread():
|
||||||
|
"""Sync prompt functions must run in a worker thread, not the event loop."""
|
||||||
|
|
||||||
|
main_thread = threading.get_ident()
|
||||||
|
fn_thread: list[int] = []
|
||||||
|
|
||||||
|
def blocking_fn() -> str:
|
||||||
|
fn_thread.append(threading.get_ident())
|
||||||
|
return "hello"
|
||||||
|
|
||||||
|
prompt = Prompt.from_function(blocking_fn)
|
||||||
|
messages = await prompt.render(None, Context())
|
||||||
|
|
||||||
|
assert messages == [UserMessage(content=TextContent(type="text", text="hello"))]
|
||||||
|
assert fn_thread[0] != main_thread
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,7 @@
|
||||||
|
import threading
|
||||||
|
|
||||||
|
import anyio
|
||||||
|
import anyio.from_thread
|
||||||
import pytest
|
import pytest
|
||||||
from pydantic import BaseModel
|
from pydantic import BaseModel
|
||||||
|
|
||||||
|
|
@ -190,3 +194,51 @@ class TestFunctionResourceMetadata:
|
||||||
)
|
)
|
||||||
|
|
||||||
assert resource.meta is None
|
assert resource.meta is None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_sync_fn_runs_in_worker_thread():
|
||||||
|
"""Sync resource functions must run in a worker thread, not the event loop."""
|
||||||
|
|
||||||
|
main_thread = threading.get_ident()
|
||||||
|
fn_thread: list[int] = []
|
||||||
|
|
||||||
|
def blocking_fn() -> str:
|
||||||
|
fn_thread.append(threading.get_ident())
|
||||||
|
return "data"
|
||||||
|
|
||||||
|
resource = FunctionResource(uri="resource://test", name="test", fn=blocking_fn)
|
||||||
|
result = await resource.read()
|
||||||
|
|
||||||
|
assert result == "data"
|
||||||
|
assert fn_thread[0] != main_thread
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_sync_fn_does_not_block_event_loop():
|
||||||
|
"""A blocking sync resource function must not stall the event loop.
|
||||||
|
|
||||||
|
On regression (sync runs inline), anyio.from_thread.run_sync raises
|
||||||
|
RuntimeError because there is no worker-thread context, failing fast.
|
||||||
|
"""
|
||||||
|
handler_entered = anyio.Event()
|
||||||
|
release = threading.Event()
|
||||||
|
|
||||||
|
def blocking_fn() -> str:
|
||||||
|
anyio.from_thread.run_sync(handler_entered.set)
|
||||||
|
release.wait()
|
||||||
|
return "done"
|
||||||
|
|
||||||
|
resource = FunctionResource(uri="resource://test", name="test", fn=blocking_fn)
|
||||||
|
result: list[str | bytes] = []
|
||||||
|
|
||||||
|
async def run() -> None:
|
||||||
|
result.append(await resource.read())
|
||||||
|
|
||||||
|
with anyio.fail_after(5):
|
||||||
|
async with anyio.create_task_group() as tg:
|
||||||
|
tg.start_soon(run)
|
||||||
|
await handler_entered.wait()
|
||||||
|
release.set()
|
||||||
|
|
||||||
|
assert result == ["done"]
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,5 @@
|
||||||
import json
|
import json
|
||||||
|
import threading
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
@ -310,3 +311,22 @@ class TestResourceTemplateMetadata:
|
||||||
assert resource.meta == metadata
|
assert resource.meta == metadata
|
||||||
assert resource.meta["category"] == "inventory"
|
assert resource.meta["category"] == "inventory"
|
||||||
assert resource.meta["cacheable"] is True
|
assert resource.meta["cacheable"] is True
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_sync_fn_runs_in_worker_thread():
|
||||||
|
"""Sync template functions must run in a worker thread, not the event loop."""
|
||||||
|
|
||||||
|
main_thread = threading.get_ident()
|
||||||
|
fn_thread: list[int] = []
|
||||||
|
|
||||||
|
def blocking_fn(name: str) -> str:
|
||||||
|
fn_thread.append(threading.get_ident())
|
||||||
|
return f"hello {name}"
|
||||||
|
|
||||||
|
template = ResourceTemplate.from_function(fn=blocking_fn, uri_template="test://{name}")
|
||||||
|
resource = await template.create_resource("test://world", {"name": "world"}, Context())
|
||||||
|
|
||||||
|
assert isinstance(resource, FunctionResource)
|
||||||
|
assert await resource.read() == "hello world"
|
||||||
|
assert fn_thread[0] != main_thread
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue