Add SSE-keepalive to not time out on long prefill on clients (#1803)

## Motivation

<!-- Why is this change needed? What problem does it solve? -->
<!-- If it fixes an open issue, please link to the issue here -->

## Changes

<!-- Describe what you changed in detail -->

## Why It Works

<!-- Explain why your approach solves the problem -->

## Test Plan

### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
<!-- - -->

### Automated Testing
<!-- Describe changes to automated tests, or how existing tests cover
this change -->
<!-- - -->

src/exo/api/keepalive.py

@@ -0,0 +1,34 @@
+from collections.abc import AsyncIterator
+from typing import Final
+
+import anyio
+
+_DONE: Final = object()
+
+
+async def with_sse_keepalive(
+    generator: AsyncIterator[str],
+    keepalive_message: str = ": keep-alive\n\n",
+    interval: float = 10.0,
+) -> AsyncIterator[str]:
+    yield keepalive_message
+    send, recv = anyio.create_memory_object_stream[str | object]()
+
+    async def _consume() -> None:
+        async for item in generator:
+            await send.send(item)
+        await send.send(_DONE)
+
+    async with anyio.create_task_group() as tg:
+        tg.start_soon(_consume)
+        while True:
+            item: str | object | None = None
+            with anyio.move_on_after(interval):
+                item = await recv.receive()
+            if item is None:
+                yield keepalive_message
+            elif item is _DONE:
+                break
+            else:
+                assert isinstance(item, str)
+                yield item

src/exo/api/main.py

@@ -46,6 +46,7 @@ from exo.api.adapters.responses import (
     generate_responses_stream,
     responses_request_to_text_generation,
 )
+from exo.api.keepalive import with_sse_keepalive
 from exo.api.types import (
     AddCustomModelParams,
     AdvancedImageParams,
@@ -799,9 +800,11 @@ class API:
 
         if payload.stream:
             return StreamingResponse(
-                generate_chat_stream(
-                    command.command_id,
-                    self._token_chunk_stream(command.command_id),
+                with_sse_keepalive(
+                    generate_chat_stream(
+                        command.command_id,
+                        self._token_chunk_stream(command.command_id),
+                    ),
                 ),
                 media_type="text/event-stream",
                 headers={
@@ -1404,10 +1407,12 @@ class API:
 
         if payload.stream:
             return StreamingResponse(
-                generate_claude_stream(
-                    command.command_id,
-                    payload.model,
-                    self._token_chunk_stream(command.command_id),
+                with_sse_keepalive(
+                    generate_claude_stream(
+                        command.command_id,
+                        payload.model,
+                        self._token_chunk_stream(command.command_id),
+                    ),
                 ),
                 media_type="text/event-stream",
                 headers={
@@ -1438,10 +1443,12 @@ class API:
 
         if payload.stream:
             return StreamingResponse(
-                generate_responses_stream(
-                    command.command_id,
-                    payload.model,
-                    self._token_chunk_stream(command.command_id),
+                with_sse_keepalive(
+                    generate_responses_stream(
+                        command.command_id,
+                        payload.model,
+                        self._token_chunk_stream(command.command_id),
+                    ),
                 ),
                 media_type="text/event-stream",
                 headers={