feat: SSE keepalive for chat completions streaming + fix cancellation propagation

src/exo/master/adapters/chat_completions.py

@@ -147,6 +147,40 @@ def chunk_to_response(
     )
 
 
+async def sse_with_keepalive(
+    sse_stream: AsyncGenerator[str, None],
+    interval: float = 15.0,
+) -> AsyncGenerator[str, None]:
+    import asyncio
+
+    queue: asyncio.Queue[str | None] = asyncio.Queue()
+
+    async def _producer() -> None:
+        try:
+            async for line in sse_stream:
+                await queue.put(line)
+        finally:
+            await queue.put(None)
+
+    task = asyncio.create_task(_producer())
+    try:
+        while True:
+            try:
+                item = await asyncio.wait_for(queue.get(), timeout=interval)
+                if item is None:
+                    return
+                yield item
+            except asyncio.TimeoutError:
+                yield ": keepalive\n\n"
+    finally:
+        task.cancel()
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass
+
+
+
 async def generate_chat_stream(
     command_id: CommandId,
     chunk_stream: AsyncGenerator[

src/exo/master/api.py

@@ -1,3 +1,4 @@
+import asyncio
 import base64
 import contextlib
 import json
@@ -25,6 +26,7 @@
     chat_request_to_text_generation,
     collect_chat_response,
     generate_chat_stream,
+    sse_with_keepalive,
 )
 from exo.master.adapters.claude import (
     claude_request_to_text_generation,
@@ -611,7 +613,7 @@ async def _token_chunk_stream(
                     if chunk.finish_reason is not None:
                         break
 
-        except anyio.get_cancelled_exc_class():
+        except (anyio.get_cancelled_exc_class(), asyncio.CancelledError):
             command = TaskCancelled(cancelled_command_id=command_id)
             with anyio.CancelScope(shield=True):
                 await self.command_sender.send(
@@ -712,9 +714,11 @@ async def chat_completions(
 
         if payload.stream:
             return StreamingResponse(
-                generate_chat_stream(
-                    command.command_id,
-                    self._token_chunk_stream(command.command_id),
+                sse_with_keepalive(
+                    generate_chat_stream(
+                        command.command_id,
+                        self._token_chunk_stream(command.command_id),
+                    ),
                 ),
                 media_type="text/event-stream",
                 headers={
@@ -965,7 +969,7 @@ async def _generate_image_stream(
                         del image_total_chunks[key]
                         del image_metadata[key]
 
-        except anyio.get_cancelled_exc_class():
+        except (anyio.get_cancelled_exc_class(), asyncio.CancelledError):
             command = TaskCancelled(cancelled_command_id=command_id)
             with anyio.CancelScope(shield=True):
                 await self.command_sender.send(
@@ -1051,7 +1055,7 @@ async def _collect_image_chunks(
                     )
 
             return (images, stats if capture_stats else None)
-        except anyio.get_cancelled_exc_class():
+        except (anyio.get_cancelled_exc_class(), asyncio.CancelledError):
             command = TaskCancelled(cancelled_command_id=command_id)
             with anyio.CancelScope(shield=True):
                 await self.command_sender.send(