← back to Exo

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

e5cb7b80d04d180829e4a342a8b3d19d07130bdc · 2026-03-30 12:18:38 +0100 · rltakashige

## 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 -->
<!-- - -->

Files touched

• A src/exo/api/keepalive.py
• M src/exo/api/main.py

Diff

commit e5cb7b80d04d180829e4a342a8b3d19d07130bdc
Author: rltakashige <rl.takashige@gmail.com>
Date:   Mon Mar 30 12:18:38 2026 +0100

    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 | 34 ++++++++++++++++++++++++++++++++++
 src/exo/api/main.py      | 29 ++++++++++++++++++-----------
 2 files changed, 52 insertions(+), 11 deletions(-)

diff --git a/src/exo/api/keepalive.py b/src/exo/api/keepalive.py
new file mode 100644
index 00000000..96a357e0
--- /dev/null
+++ b/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
diff --git a/src/exo/api/main.py b/src/exo/api/main.py
index cc3a159b..3021d689 100644
--- a/src/exo/api/main.py
+++ b/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={

← 635801d5 Add multimodality! (#1802)  ·  back to Exo  ·  Integrations helpers (#1810) 39c39e81 →