86 lines
2.8 KiB
Python
86 lines
2.8 KiB
Python
"""Simple task server demonstrating MCP tasks over streamable HTTP."""
|
|
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
|
|
import anyio
|
|
import click
|
|
import uvicorn
|
|
from mcp import types
|
|
from mcp.server import Server, ServerRequestContext
|
|
from mcp.server.experimental.task_context import ServerTaskContext
|
|
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
|
|
from starlette.applications import Starlette
|
|
from starlette.routing import Mount
|
|
|
|
|
|
async def handle_list_tools(
|
|
ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
|
|
) -> types.ListToolsResult:
|
|
return types.ListToolsResult(
|
|
tools=[
|
|
types.Tool(
|
|
name="long_running_task",
|
|
description="A task that takes a few seconds to complete with status updates",
|
|
input_schema={"type": "object", "properties": {}},
|
|
execution=types.ToolExecution(task_support=types.TASK_REQUIRED),
|
|
)
|
|
]
|
|
)
|
|
|
|
|
|
async def handle_call_tool(
|
|
ctx: ServerRequestContext, params: types.CallToolRequestParams
|
|
) -> types.CallToolResult | types.CreateTaskResult:
|
|
"""Dispatch tool calls to their handlers."""
|
|
if params.name == "long_running_task":
|
|
ctx.experimental.validate_task_mode(types.TASK_REQUIRED)
|
|
|
|
async def work(task: ServerTaskContext) -> types.CallToolResult:
|
|
await task.update_status("Starting work...")
|
|
await anyio.sleep(1)
|
|
|
|
await task.update_status("Processing step 1...")
|
|
await anyio.sleep(1)
|
|
|
|
await task.update_status("Processing step 2...")
|
|
await anyio.sleep(1)
|
|
|
|
return types.CallToolResult(content=[types.TextContent(type="text", text="Task completed!")])
|
|
|
|
return await ctx.experimental.run_task(work)
|
|
|
|
return types.CallToolResult(
|
|
content=[types.TextContent(type="text", text=f"Unknown tool: {params.name}")],
|
|
is_error=True,
|
|
)
|
|
|
|
|
|
server = Server(
|
|
"simple-task-server",
|
|
on_list_tools=handle_list_tools,
|
|
on_call_tool=handle_call_tool,
|
|
)
|
|
|
|
# One-line setup: auto-registers get_task, get_task_result, list_tasks, cancel_task
|
|
server.experimental.enable_tasks()
|
|
|
|
|
|
@click.command()
|
|
@click.option("--port", default=8000, help="Port to listen on")
|
|
def main(port: int) -> int:
|
|
session_manager = StreamableHTTPSessionManager(app=server)
|
|
|
|
@asynccontextmanager
|
|
async def app_lifespan(app: Starlette) -> AsyncIterator[None]:
|
|
async with session_manager.run():
|
|
yield
|
|
|
|
starlette_app = Starlette(
|
|
routes=[Mount("/mcp", app=session_manager.handle_request)],
|
|
lifespan=app_lifespan,
|
|
)
|
|
|
|
print(f"Starting server on http://localhost:{port}/mcp")
|
|
uvicorn.run(starlette_app, host="127.0.0.1", port=port)
|
|
return 0
|