mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 15:19:10 +02:00
138 lines
4.4 KiB
Python
138 lines
4.4 KiB
Python
"""Tests for FastMCP Progress dependency (SEP-2663 tasks)."""
|
|
|
|
import asyncio
|
|
import json
|
|
|
|
from mcp_types import TextContent
|
|
|
|
from fastmcp import FastMCP
|
|
from fastmcp.server.dependencies import Progress
|
|
from fastmcp_tasks import TasksExtension
|
|
from tests.tasks.task_helpers import (
|
|
call_tool_without_optin,
|
|
running_task_server,
|
|
submit_task,
|
|
wait_for_task,
|
|
)
|
|
|
|
|
|
async def test_progress_in_immediate_execution():
|
|
"""Progress dependency works when a tool runs synchronously."""
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool
|
|
async def test_tool(progress: Progress = Progress()) -> str:
|
|
await progress.set_total(10)
|
|
await progress.increment()
|
|
await progress.set_message("Testing")
|
|
return "done"
|
|
|
|
result = await call_tool_without_optin(mcp, "test_tool", {})
|
|
assert isinstance(result.content[0], TextContent)
|
|
assert result.content[0].text == "done"
|
|
|
|
|
|
async def test_progress_in_background_task():
|
|
"""Progress dependency works inside a background task."""
|
|
mcp = FastMCP("test")
|
|
mcp.add_extension(TasksExtension())
|
|
|
|
@mcp.tool(task=True)
|
|
async def test_task(progress: Progress = Progress()) -> str:
|
|
await progress.set_total(5)
|
|
await progress.increment()
|
|
await progress.set_message("Step 1")
|
|
return "done"
|
|
|
|
async with running_task_server(mcp):
|
|
created = await submit_task(mcp, "test_task", {})
|
|
final = await wait_for_task(mcp, created.task_id)
|
|
assert final.status == "completed"
|
|
assert final.result is not None
|
|
assert final.result["structuredContent"] == {"result": "done"}
|
|
|
|
|
|
async def test_progress_tracks_multiple_increments():
|
|
"""Progress correctly tracks multiple increment calls."""
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool
|
|
async def count_to_ten(progress: Progress = Progress()) -> str:
|
|
await progress.set_total(10)
|
|
for _ in range(10):
|
|
await progress.increment()
|
|
return "counted"
|
|
|
|
result = await call_tool_without_optin(mcp, "count_to_ten", {})
|
|
assert isinstance(result.content[0], TextContent)
|
|
assert result.content[0].text == "counted"
|
|
|
|
|
|
async def test_progress_status_message_in_background_task():
|
|
"""A working task surfaces the current progress message as statusMessage."""
|
|
mcp = FastMCP("test")
|
|
mcp.add_extension(TasksExtension())
|
|
release = asyncio.Event()
|
|
|
|
@mcp.tool(task=True)
|
|
async def task_with_progress(progress: Progress = Progress()) -> str:
|
|
await progress.set_total(3)
|
|
await progress.set_message("Step 1 of 3")
|
|
await progress.increment()
|
|
await release.wait()
|
|
await progress.set_message("Step 2 of 3")
|
|
await progress.increment()
|
|
return "done"
|
|
|
|
async with running_task_server(mcp):
|
|
created = await submit_task(mcp, "task_with_progress", {})
|
|
|
|
# The task parks on `release` while working; its statusMessage should
|
|
# reflect the progress message (or be None, depending on the poll race).
|
|
working = await wait_for_task(
|
|
mcp, created.task_id, target_states=frozenset({"working"})
|
|
)
|
|
msg = working.status_message
|
|
assert msg is None or msg.startswith("Step")
|
|
|
|
release.set()
|
|
final = await wait_for_task(mcp, created.task_id)
|
|
assert final.status == "completed"
|
|
assert final.result is not None
|
|
assert final.result["structuredContent"] == {"result": "done"}
|
|
|
|
|
|
async def test_inmemory_progress_state():
|
|
"""In-memory progress stores and returns state correctly."""
|
|
mcp = FastMCP("test")
|
|
|
|
@mcp.tool
|
|
async def test_tool(progress: Progress = Progress()) -> dict:
|
|
assert progress.current is None
|
|
assert progress.total == 1
|
|
assert progress.message is None
|
|
|
|
await progress.set_total(10)
|
|
assert progress.total == 10
|
|
|
|
await progress.increment()
|
|
assert progress.current == 1
|
|
|
|
await progress.increment(2)
|
|
assert progress.current == 3
|
|
|
|
await progress.set_message("Testing")
|
|
assert progress.message == "Testing"
|
|
|
|
return {
|
|
"current": progress.current,
|
|
"total": progress.total,
|
|
"message": progress.message,
|
|
}
|
|
|
|
result = await call_tool_without_optin(mcp, "test_tool", {})
|
|
assert isinstance(result.content[0], TextContent)
|
|
state = json.loads(result.content[0].text)
|
|
assert state["current"] == 3
|
|
assert state["total"] == 10
|
|
assert state["message"] == "Testing"
|