Add SSE polling support (SEP-1699) (#1654)
This commit is contained in:
Felix Weinberger
GitHub
parent
2cd178a962
commit
281fd4765e
1 changed files with 1461 additions and 45 deletions
177
examples/servers/sse-polling-demo/mcp_sse_polling_demo/server.py
Normal file
177
examples/servers/sse-polling-demo/mcp_sse_polling_demo/server.py
Normal file
|
|
@ -0,0 +1,177 @@
|
|||
"""
|
||||
SSE Polling Demo Server
|
||||
|
||||
Demonstrates the SSE polling pattern with close_sse_stream() for long-running tasks.
|
||||
|
||||
Features demonstrated:
|
||||
- Priming events (automatic with EventStore)
|
||||
- Server-initiated stream close via close_sse_stream callback
|
||||
- Client auto-reconnect with Last-Event-ID
|
||||
- Progress notifications during long-running tasks
|
||||
|
||||
Run with:
|
||||
uv run mcp-sse-polling-demo --port 3000
|
||||
"""
|
||||
|
||||
import contextlib
|
||||
import logging
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
import anyio
|
||||
import click
|
||||
import mcp.types as types
|
||||
from mcp.server.lowlevel import Server
|
||||
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
|
||||
from starlette.applications import Starlette
|
||||
from starlette.routing import Mount
|
||||
from starlette.types import Receive, Scope, Send
|
||||
|
||||
from .event_store import InMemoryEventStore
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.option("--port", default=3000, help="Port to listen on")
|
||||
@click.option(
|
||||
"--log-level",
|
||||
default="INFO",
|
||||
help="Logging level (DEBUG, INFO, WARNING, ERROR)",
|
||||
)
|
||||
@click.option(
|
||||
"--retry-interval",
|
||||
default=100,
|
||||
help="SSE retry interval in milliseconds (sent to client)",
|
||||
)
|
||||
def main(port: int, log_level: str, retry_interval: int) -> int:
|
||||
"""Run the SSE Polling Demo server."""
|
||||
logging.basicConfig(
|
||||
level=getattr(logging, log_level.upper()),
|
||||
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
||||
)
|
||||
|
||||
# Create the lowlevel server
|
||||
app = Server("sse-polling-demo")
|
||||
|
||||
@app.call_tool()
|
||||
async def call_tool(name: str, arguments: dict[str, Any]) -> list[types.ContentBlock]:
|
||||
"""Handle tool calls."""
|
||||
ctx = app.request_context
|
||||
|
||||
if name == "process_batch":
|
||||
items = arguments.get("items", 10)
|
||||
checkpoint_every = arguments.get("checkpoint_every", 3)
|
||||
|
||||
if items < 1 or items > 100:
|
||||
return [types.TextContent(type="text", text="Error: items must be between 1 and 100")]
|
||||
if checkpoint_every < 1 or checkpoint_every > 20:
|
||||
return [types.TextContent(type="text", text="Error: checkpoint_every must be between 1 and 20")]
|
||||
|
||||
await ctx.session.send_log_message(
|
||||
level="info",
|
||||
data=f"Starting batch processing of {items} items...",
|
||||
logger="process_batch",
|
||||
related_request_id=ctx.request_id,
|
||||
)
|
||||
|
||||
for i in range(1, items + 1):
|
||||
# Simulate work
|
||||
await anyio.sleep(0.5)
|
||||
|
||||
# Report progress
|
||||
await ctx.session.send_log_message(
|
||||
level="info",
|
||||
data=f"[{i}/{items}] Processing item {i}",
|
||||
logger="process_batch",
|
||||
related_request_id=ctx.request_id,
|
||||
)
|
||||
|
||||
# Checkpoint: close stream to trigger client reconnect
|
||||
if i % checkpoint_every == 0 and i < items:
|
||||
await ctx.session.send_log_message(
|
||||
level="info",
|
||||
data=f"Checkpoint at item {i} - closing SSE stream for polling",
|
||||
logger="process_batch",
|
||||
related_request_id=ctx.request_id,
|
||||
)
|
||||
if ctx.close_sse_stream:
|
||||
logger.info(f"Closing SSE stream at checkpoint {i}")
|
||||
await ctx.close_sse_stream()
|
||||
# Wait for client to reconnect (must be > retry_interval of 100ms)
|
||||
await anyio.sleep(0.2)
|
||||
|
||||
return [
|
||||
types.TextContent(
|
||||
type="text",
|
||||
text=f"Successfully processed {items} items with checkpoints every {checkpoint_every} items",
|
||||
)
|
||||
]
|
||||
|
||||
return [types.TextContent(type="text", text=f"Unknown tool: {name}")]
|
||||
|
||||
@app.list_tools()
|
||||
async def list_tools() -> list[types.Tool]:
|
||||
"""List available tools."""
|
||||
return [
|
||||
types.Tool(
|
||||
name="process_batch",
|
||||
description=(
|
||||
"Process a batch of items with periodic checkpoints. "
|
||||
"Demonstrates SSE polling where server closes stream periodically."
|
||||
),
|
||||
inputSchema={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"items": {
|
||||
"type": "integer",
|
||||
"description": "Number of items to process (1-100)",
|
||||
"default": 10,
|
||||
},
|
||||
"checkpoint_every": {
|
||||
"type": "integer",
|
||||
"description": "Close stream after this many items (1-20)",
|
||||
"default": 3,
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
]
|
||||
|
||||
# Create event store for resumability
|
||||
event_store = InMemoryEventStore()
|
||||
|
||||
# Create session manager with event store and retry interval
|
||||
session_manager = StreamableHTTPSessionManager(
|
||||
app=app,
|
||||
event_store=event_store,
|
||||
retry_interval=retry_interval,
|
||||
)
|
||||
|
||||
async def handle_streamable_http(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
await session_manager.handle_request(scope, receive, send)
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def lifespan(starlette_app: Starlette) -> AsyncIterator[None]:
|
||||
async with session_manager.run():
|
||||
logger.info(f"SSE Polling Demo server started on port {port}")
|
||||
logger.info("Try: POST /mcp with tools/call for 'process_batch'")
|
||||
yield
|
||||
logger.info("Server shutting down...")
|
||||
|
||||
starlette_app = Starlette(
|
||||
debug=True,
|
||||
routes=[
|
||||
Mount("/mcp", app=handle_streamable_http),
|
||||
],
|
||||
lifespan=lifespan,
|
||||
)
|
||||
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(starlette_app, host="127.0.0.1", port=port)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Reference in a new issue