Marcelo Trylesinski
148 lines
5.3 KiB
Python
148 lines
5.3 KiB
Python
"""Simple interactive task server demonstrating elicitation and sampling.
|
|
|
|
This example shows the simplified task API where:
|
|
- server.experimental.enable_tasks() sets up all infrastructure
|
|
- ctx.experimental.run_task() handles task lifecycle automatically
|
|
- ServerTaskContext.elicit() and ServerTaskContext.create_message() queue requests properly
|
|
"""
|
|
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any
|
|
|
|
import click
|
|
import mcp.types as types
|
|
import uvicorn
|
|
from mcp.server.experimental.task_context import ServerTaskContext
|
|
from mcp.server.lowlevel import Server
|
|
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
|
|
from starlette.applications import Starlette
|
|
from starlette.routing import Mount
|
|
|
|
server = Server("simple-task-interactive")
|
|
|
|
# Enable task support - this auto-registers all handlers
|
|
server.experimental.enable_tasks()
|
|
|
|
|
|
@server.list_tools()
|
|
async def list_tools() -> list[types.Tool]:
|
|
return [
|
|
types.Tool(
|
|
name="confirm_delete",
|
|
description="Asks for confirmation before deleting (demonstrates elicitation)",
|
|
input_schema={
|
|
"type": "object",
|
|
"properties": {"filename": {"type": "string"}},
|
|
},
|
|
execution=types.ToolExecution(task_support=types.TASK_REQUIRED),
|
|
),
|
|
types.Tool(
|
|
name="write_haiku",
|
|
description="Asks LLM to write a haiku (demonstrates sampling)",
|
|
input_schema={"type": "object", "properties": {"topic": {"type": "string"}}},
|
|
execution=types.ToolExecution(task_support=types.TASK_REQUIRED),
|
|
),
|
|
]
|
|
|
|
|
|
async def handle_confirm_delete(arguments: dict[str, Any]) -> types.CreateTaskResult:
|
|
"""Handle the confirm_delete tool - demonstrates elicitation."""
|
|
ctx = server.request_context
|
|
ctx.experimental.validate_task_mode(types.TASK_REQUIRED)
|
|
|
|
filename = arguments.get("filename", "unknown.txt")
|
|
print(f"\n[Server] confirm_delete called for '{filename}'")
|
|
|
|
async def work(task: ServerTaskContext) -> types.CallToolResult:
|
|
print(f"[Server] Task {task.task_id} starting elicitation...")
|
|
|
|
result = await task.elicit(
|
|
message=f"Are you sure you want to delete '{filename}'?",
|
|
requested_schema={
|
|
"type": "object",
|
|
"properties": {"confirm": {"type": "boolean"}},
|
|
"required": ["confirm"],
|
|
},
|
|
)
|
|
|
|
print(f"[Server] Received elicitation response: action={result.action}, content={result.content}")
|
|
|
|
if result.action == "accept" and result.content:
|
|
confirmed = result.content.get("confirm", False)
|
|
text = f"Deleted '{filename}'" if confirmed else "Deletion cancelled"
|
|
else:
|
|
text = "Deletion cancelled"
|
|
|
|
print(f"[Server] Completing task with result: {text}")
|
|
return types.CallToolResult(content=[types.TextContent(type="text", text=text)])
|
|
|
|
return await ctx.experimental.run_task(work)
|
|
|
|
|
|
async def handle_write_haiku(arguments: dict[str, Any]) -> types.CreateTaskResult:
|
|
"""Handle the write_haiku tool - demonstrates sampling."""
|
|
ctx = server.request_context
|
|
ctx.experimental.validate_task_mode(types.TASK_REQUIRED)
|
|
|
|
topic = arguments.get("topic", "nature")
|
|
print(f"\n[Server] write_haiku called for topic '{topic}'")
|
|
|
|
async def work(task: ServerTaskContext) -> types.CallToolResult:
|
|
print(f"[Server] Task {task.task_id} starting sampling...")
|
|
|
|
result = await task.create_message(
|
|
messages=[
|
|
types.SamplingMessage(
|
|
role="user",
|
|
content=types.TextContent(type="text", text=f"Write a haiku about {topic}"),
|
|
)
|
|
],
|
|
max_tokens=50,
|
|
)
|
|
|
|
haiku = "No response"
|
|
if isinstance(result.content, types.TextContent):
|
|
haiku = result.content.text
|
|
|
|
print(f"[Server] Received sampling response: {haiku[:50]}...")
|
|
return types.CallToolResult(content=[types.TextContent(type="text", text=f"Haiku:\n{haiku}")])
|
|
|
|
return await ctx.experimental.run_task(work)
|
|
|
|
|
|
@server.call_tool()
|
|
async def handle_call_tool(name: str, arguments: dict[str, Any]) -> types.CallToolResult | types.CreateTaskResult:
|
|
"""Dispatch tool calls to their handlers."""
|
|
if name == "confirm_delete":
|
|
return await handle_confirm_delete(arguments)
|
|
elif name == "write_haiku":
|
|
return await handle_write_haiku(arguments)
|
|
else:
|
|
return types.CallToolResult(
|
|
content=[types.TextContent(type="text", text=f"Unknown tool: {name}")],
|
|
is_error=True,
|
|
)
|
|
|
|
|
|
def create_app(session_manager: StreamableHTTPSessionManager) -> Starlette:
|
|
@asynccontextmanager
|
|
async def app_lifespan(app: Starlette) -> AsyncIterator[None]:
|
|
async with session_manager.run():
|
|
yield
|
|
|
|
return Starlette(
|
|
routes=[Mount("/mcp", app=session_manager.handle_request)],
|
|
lifespan=app_lifespan,
|
|
)
|
|
|
|
|
|
@click.command()
|
|
@click.option("--port", default=8000, help="Port to listen on")
|
|
def main(port: int) -> int:
|
|
session_manager = StreamableHTTPSessionManager(app=server)
|
|
starlette_app = create_app(session_manager)
|
|
print(f"Starting server on http://localhost:{port}/mcp")
|
|
uvicorn.run(starlette_app, host="127.0.0.1", port=port)
|
|
return 0
|