patx/micropie
added a patch to make sure async generators stop when client disconnects. fixes #10
Commit 1ea95de · patx · 2025-07-08T23:54:15-04:00
Comments
No comments yet.
Diff
diff --git a/examples/streaming/sse/chat.py b/examples/streaming/sse/chat.py
index 4ccc400..5a47192 100644
--- a/examples/streaming/sse/chat.py
+++ b/examples/streaming/sse/chat.py
@@ -1,10 +1,3 @@
-"""
-Example chat application using server sent events.
-This works with `uvicorn` but works better with `daphne`.
-This kind of application is better suited for websockets,
-but is a good example nonetheless.
-"""
-
from micropie import App
from collections import deque
import json
@@ -13,55 +6,82 @@ import asyncio
class ChatApp(App):
def __init__(self):
super().__init__()
- # Store messages in a deque with a max length
self.messages = deque(maxlen=100)
- # Store connected clients
self.clients = set()
async def index(self):
- """Serve the chat interface."""
- return await self._render_template('chat.html')
+ # Inline HTML for demo purposes
+ html = """
+ <!doctype html>
+ <title>Micropie Chat SSE</title>
+ <h1>Micropie Chat SSE Demo</h1>
+ <form id="chat-form">
+ Name: <input id="username" value="Anonymous" size="10">
+ <input id="message" autocomplete="off" size="40">
+ <button>Send</button>
+ </form>
+ <pre id="log"></pre>
+ <script>
+ const log = document.getElementById('log');
+ const usernameInput = document.getElementById('username');
+ const messageInput = document.getElementById('message');
+ const form = document.getElementById('chat-form');
+
+ const evtSource = new EventSource('/events');
+ evtSource.onmessage = function(event) {
+ const data = JSON.parse(event.data);
+ log.textContent += data.username + ": " + data.message + "\\n";
+ };
+
+ form.onsubmit = async function(e) {
+ e.preventDefault();
+ await fetch('/send', {
+ method: 'POST',
+ headers: {'Content-Type': 'application/json'},
+ body: JSON.stringify({
+ username: usernameInput.value,
+ message: messageInput.value,
+ })
+ });
+ messageInput.value = "";
+ };
+ </script>
+ """
+ return 200, html, [('Content-Type', 'text/html')]
async def send(self):
data = self.request.get_json
- username = data.get('username', 'Anonymous')
- message = data.get('message', '')
-
+ username = data.get("username", "Anonymous")
+ message = data.get("message", "")
if message.strip():
- msg = {'username': username, 'message': message}
+ msg = {"username": username, "message": message}
self.messages.append(msg)
-
# Broadcast to all connected clients
for client in self.clients.copy():
try:
await client.put(json.dumps(msg))
except Exception:
self.clients.discard(client)
-
- return 200, {'status': 'success'}
+ return 200, {"status": "success"}
async def events(self):
- """Stream chat messages via server-sent events."""
async def stream():
- client_queue = asyncio.Queue()
- self.clients.add(client_queue)
-
+ queue = asyncio.Queue()
+ setattr(stream, "_micropie_disconnect_queue", queue)
+ self.clients.add(queue)
try:
- # Send existing messages first
for msg in self.messages:
yield f"data: {json.dumps(msg)}\n\n"
-
- # Keep connection open and send new messages
while True:
- msg = await client_queue.get()
- yield f"data: {msg}\n\n"
-
- except asyncio.CancelledError:
- pass
-
+ try:
+ msg = await queue.get()
+ if msg is None:
+ break
+ yield f"data: {msg}\n\n"
+ except asyncio.CancelledError:
+ break
finally:
- self.clients.discard(client_queue)
-
+ self.clients.discard(queue)
return 200, stream(), [
('Content-Type', 'text/event-stream'),
('Cache-Control', 'no-cache'),
diff --git a/examples/streaming/sse/templates/chat.html b/examples/streaming/sse/templates/chat.html
deleted file mode 100644
index 1af0549..0000000
--- a/examples/streaming/sse/templates/chat.html
+++ /dev/null
@@ -1,97 +0,0 @@
-<!DOCTYPE html>
-<html lang="en">
-<head>
- <meta charset="UTF-8">
- <meta name="viewport" content="width=device-width, initial-scale=1.0">
- <title>Chat Application</title>
- <script src="https://cdn.tailwindcss.com"></script>
- <style>
- #chat-messages {
- height: 400px;
- overflow-y: auto;
- scroll-behavior: smooth;
- }
- #chat-messages::-webkit-scrollbar {
- width: 8px;
- }
- #chat-messages::-webkit-scrollbar-thumb {
- background: #9ca3af;
- border-radius: 4px;
- }
- </style>
-</head>
-<body class="bg-gray-100 flex items-center justify-center min-h-screen">
- <div class="w-full max-w-2xl bg-white rounded-lg shadow-lg p-6">
- <h1 class="text-2xl font-bold text-gray-800 mb-4">Chat Room</h1>
- <div id="chat-messages" class="bg-gray-50 p-4 rounded-lg mb-4 border border-gray-200"></div>
- <div class="flex gap-2">
- <input
- id="username"
- type="text"
- placeholder="Your username"
- class="flex-1 p-2 border border-gray-300 rounded-lg focus:outline-none focus:ring-2 focus:ring-blue-500"
- >
- <input
- id="message"
- type="text"
- placeholder="Type a message..."
- class="flex-1 p-2 border border-gray-300 rounded-lg focus:outline-none focus:ring-2 focus:ring-blue-500"
- >
- <button
- id="send-button"
- class="px-4 py-2 bg-blue-500 text-white rounded-lg hover:bg-blue-600 focus:outline-none focus:ring-2 focus:ring-blue-500"
- >
- Send
- </button>
- </div>
- </div>
-
- <script>
- // Initialize EventSource for server-sent events
- const source = new EventSource('/events');
-
- source.onmessage = function(event) {
- const data = JSON.parse(event.data);
- const messageDiv = document.createElement('div');
- messageDiv.className = 'mb-2';
- messageDiv.innerHTML = `<span class="font-bold text-blue-600">${data.username}:</span> ${data.message}`;
- const chatMessages = document.getElementById('chat-messages');
- chatMessages.appendChild(messageDiv);
- chatMessages.scrollTop = chatMessages.scrollHeight;
- };
-
- source.onerror = function() {
- console.error('EventSource failed. Attempting to reconnect...');
- };
-
- // Send message on button click or Enter key
- const sendButton = document.getElementById('send-button');
- const messageInput = document.getElementById('message');
- const usernameInput = document.getElementById('username');
-
- async function sendMessage() {
- const username = usernameInput.value.trim() || 'Anonymous';
- const message = messageInput.value.trim();
- if (message) {
- try {
- await fetch('/send', {
- method: 'POST',
- headers: { 'Content-Type': 'application/json' },
- body: JSON.stringify({ username, message })
- });
- messageInput.value = '';
- } catch (error) {
- console.error('Error sending message:', error);
- }
- }
- }
-
- sendButton.addEventListener('click', sendMessage);
- messageInput.addEventListener('keypress', (e) => {
- if (e.key === 'Enter') {
- sendMessage();
- }
- });
- </script>
-</body>
-</html>
diff --git a/micropie.py b/micropie.py
index 9b1d823..a317541 100644
--- a/micropie.py
+++ b/micropie.py
@@ -571,7 +571,76 @@ class App:
result.get("headers", extra_headers)
)
- await self._send_response(send, status_code, response_body, extra_headers)
+ # --- BEGIN PATCH FOR SSE CLIENT DISCONNECT ---
+ if hasattr(response_body, "__aiter__"):
+ sanitized_headers: List[Tuple[str, str]] = []
+ for k, v in extra_headers:
+ if "\n" in k or "\r" in k or "\n" in v or "\r" in v:
+ print(f"Header injection attempt detected: {k}: {v}")
+ continue
+ sanitized_headers.append((k, v))
+ if not any(h[0].lower() == "content-type" for h in sanitized_headers):
+ sanitized_headers.append(("Content-Type", "text/html; charset=utf-8"))
+ await send({
+ "type": "http.response.start",
+ "status": status_code,
+ "headers": [(k.encode("latin-1"), v.encode("latin-1")) for k, v in sanitized_headers],
+ })
+
+ gen = response_body
+
+ async def streamer():
+ try:
+ async for chunk in gen:
+ if isinstance(chunk, str):
+ chunk = chunk.encode("utf-8")
+ await send({
+ "type": "http.response.body",
+ "body": chunk,
+ "more_body": True
+ })
+ await send({
+ "type": "http.response.body",
+ "body": b"",
+ "more_body": False
+ })
+ except asyncio.CancelledError:
+ # Optional: do any extra logging here
+ raise
+ finally:
+ if hasattr(gen, "aclose"):
+ await gen.aclose()
+
+ streaming_task = asyncio.create_task(streamer())
+ try:
+ while True:
+ msg_task = asyncio.create_task(receive())
+ done, pending = await asyncio.wait(
+ [streaming_task, msg_task],
+ return_when=asyncio.FIRST_COMPLETED
+ )
+ if streaming_task in done:
+ break
+ if msg_task in done:
+ msg = msg_task.result()
+ if msg["type"] == "http.disconnect":
+ streaming_task.cancel()
+ try:
+ await streaming_task
+ except asyncio.CancelledError:
+ pass
+ break
+ finally:
+ if not streaming_task.done():
+ streaming_task.cancel()
+ try:
+ await streaming_task
+ except asyncio.CancelledError:
+ pass
+ return
+ # --- END PATCH FOR SSE CLIENT DISCONNECT ---
+ else:
+ await self._send_response(send, status_code, response_body, extra_headers)
finally:
current_request.reset(token)