gabriel / musehub public
sse.py python
193 lines 5.7 KB
Raw
sha256:c73281cffffbc70487969720af074e54b34326d59ec6aff0827fa16512a4e512 docs(mwp-2): mark all acceptance criteria met; add closing … Sonnet 4.6 84 days ago
1 """SSE (Server-Sent Events) helpers for MCP Streamable HTTP transport.
2
3 Provides event formatting per the HTML Living Standard:
4 https://html.spec.whatwg.org/multipage/server-sent-events.html
5
6 Every outbound MCP server-to-client message (elicitation/create, progress
7 notifications, tool responses) is formatted with :func:`sse_event` before being
8 pushed into session SSE queues or streamed via :func:`sse_stream_response`.
9
10 Design notes:
11 - ``data:`` lines must not contain bare newlines — each embedded newline
12 is split into a separate ``data:`` line per the spec.
13 - ``event_id`` monotonically increments per session using a simple counter;
14 callers that need replay support should pass the counter value.
15 - The heartbeat (``": heartbeat"`` comment) keeps proxies alive without
16 affecting the event stream semantics.
17 """
18
19 import json
20 import logging
21 from typing import AsyncIterator
22
23 from musehub.types.json_types import JSONObject, JSONValue
24
25 logger = logging.getLogger(__name__)
26
27 # SSE content type required by the spec.
28 SSE_CONTENT_TYPE = "text/event-stream"
29
30 # Heartbeat comment text — sent every N seconds on GET /mcp streams.
31 _HEARTBEAT = ": heartbeat\n\n"
32
33
34 def sse_event(
35 data: JSONObject,
36 *,
37 event_id: str | None = None,
38 event_type: str | None = None,
39 retry_ms: int | None = None,
40 ) -> str:
41 """Format a JSON object as an SSE event string.
42
43 Args:
44 data: The JSON-serialisable payload. Encoded as compact JSON on a
45 single ``data:`` line (embedded newlines are split per the spec).
46 event_id: Optional ``id:`` field. Clients use this as ``Last-Event-ID``
47 on reconnection.
48 event_type: Optional ``event:`` field (default stream type if omitted).
49 retry_ms: Optional ``retry:`` reconnection timeout in milliseconds.
50
51 Returns:
52 A fully-formatted SSE event string ending with ``\\n\\n``.
53 """
54 parts: list[str] = []
55
56 if event_id is not None:
57 parts.append(f"id: {event_id}")
58
59 if event_type is not None:
60 parts.append(f"event: {event_type}")
61
62 if retry_ms is not None:
63 parts.append(f"retry: {retry_ms}")
64
65 # Encode data as compact JSON, splitting on newlines per spec.
66 encoded = json.dumps(data, separators=(",", ":"), default=str)
67 for line in encoded.split("\n"):
68 parts.append(f"data: {line}")
69
70 body = "\n".join(parts)
71 return f"{body}\n\n"
72
73
74 def sse_heartbeat() -> str:
75 """Return the SSE heartbeat comment string."""
76 return _HEARTBEAT
77
78
79 def sse_notification(
80 method: str,
81 params: JSONObject | None = None,
82 *,
83 event_id: str | None = None,
84 ) -> str:
85 """Format an MCP JSON-RPC notification as an SSE event.
86
87 Args:
88 method: JSON-RPC method name (e.g. ``"notifications/progress"``).
89 params: Optional parameters dict.
90 event_id: Optional SSE event ID.
91
92 Returns:
93 Formatted SSE event string.
94 """
95 payload: JSONObject = {"jsonrpc": "2.0", "method": method}
96 if params is not None:
97 payload["params"] = params
98 return sse_event(payload, event_id=event_id)
99
100
101 def sse_request(
102 req_id: str | int,
103 method: str,
104 params: JSONObject | None = None,
105 *,
106 event_id: str | None = None,
107 ) -> str:
108 """Format an MCP JSON-RPC request (server→client) as an SSE event.
109
110 Used for server-initiated requests such as ``elicitation/create``.
111
112 Args:
113 req_id: JSON-RPC request ID. The client echoes this in its response.
114 method: JSON-RPC method name.
115 params: Optional parameters dict.
116 event_id: Optional SSE event ID.
117
118 Returns:
119 Formatted SSE event string.
120 """
121 payload: JSONObject = {
122 "jsonrpc": "2.0",
123 "id": req_id,
124 "method": method,
125 }
126 if params is not None:
127 payload["params"] = params
128 return sse_event(payload, event_id=event_id)
129
130
131 def sse_response(
132 req_id: str | int | None,
133 result: JSONObject,
134 *,
135 event_id: str | None = None,
136 ) -> str:
137 """Format an MCP JSON-RPC success response as an SSE event.
138
139 Used to stream tool call results back to the client when the POST response
140 is a ``text/event-stream`` rather than a single JSON body.
141
142 Args:
143 req_id: JSON-RPC request ID from the original client request.
144 result: The result payload dict.
145 event_id: Optional SSE event ID.
146
147 Returns:
148 Formatted SSE event string.
149 """
150 payload: JSONObject = {"jsonrpc": "2.0", "id": req_id, "result": result}
151 return sse_event(payload, event_id=event_id)
152
153
154 async def heartbeat_stream(
155 event_stream: AsyncIterator[str],
156 *,
157 interval_seconds: float = 15.0,
158 ) -> AsyncIterator[str]:
159 """Interleave heartbeat SSE comments into an existing async event stream.
160
161 Yields events from ``event_stream`` unchanged and injects a heartbeat
162 comment if no event has been sent for ``interval_seconds``.
163
164 Args:
165 event_stream: Upstream async iterator of SSE event strings.
166 interval_seconds: Maximum silence duration before injecting a heartbeat.
167
168 Yields:
169 SSE event strings, including injected heartbeat comments.
170 """
171 import asyncio
172
173 aiter = event_stream.__aiter__()
174 pending: asyncio.Task[str] | None = None
175
176 try:
177 while True:
178 if pending is None:
179 pending = asyncio.ensure_future(aiter.__anext__())
180
181 try:
182 event = await asyncio.wait_for(
183 asyncio.shield(pending), timeout=interval_seconds
184 )
185 pending = None
186 yield event
187 except asyncio.TimeoutError:
188 yield _HEARTBEAT
189 except StopAsyncIteration:
190 return
191 finally:
192 if pending is not None and not pending.done():
193 pending.cancel()
File History 2 commits
sha256:c73281cffffbc70487969720af074e54b34326d59ec6aff0827fa16512a4e512 docs(mwp-2): mark all acceptance criteria met; add closing … Sonnet 4.6 84 days ago
sha256:4dd2a937f66f8e36d9aa59bd1bf3cb880ca1d503fef67300a0cba3b482784081 test(mwp2): Phase 0 RED — reproduction tests for _walk_comm… Sonnet 4.6 84 days ago