Cancel SSE keep-alive when instance is deleted (#1828)

b12cd1b1862fe368f5ee41c6b0df4a0a5af40c73 · 2026-04-08 16:14:28 +0100 · ciaranbor

## Motivation

When a model instance is deleted (e.g. node disconnect, manual
teardown), any in-flight SSE streaming connections for that instance
hang indefinitely. The API never closes the response stream, so clients
block forever waiting for more chunks.

## Changes

- Listen for `InstanceDeleted` events in the API event loop
- Add `_close_streams_for_instance()` to find and close any active
text/image generation queues tied to tasks on the deleted instance
- Add unit tests covering text gen, image gen, and
unrelated-instance-not-closed scenarios

## Why It Works

When an instance is deleted, we iterate `state.tasks` to find commands
running on that instance, then close and remove their send-side queue
handles. This causes the SSE generator to terminate, unblocking the
client.

## Test Plan

### Manual Testing
- This was causing issues for me on another branch (integration tests).
Including this fix solved the issue

### Automated Testing
- `test_instance_deleted_stream_cleanup.py`: 3 tests covering text gen
cleanup, image gen cleanup, and ensuring unrelated streams are not
affected

Files touched

M src/exo/api/main.py
A src/exo/api/tests/test_instance_deleted_stream_cleanup.py

Diff

commit b12cd1b1862fe368f5ee41c6b0df4a0a5af40c73
Author: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date:   Wed Apr 8 16:14:28 2026 +0100

    Cancel SSE keep-alive when instance is deleted (#1828)
    
    ## Motivation
    
    When a model instance is deleted (e.g. node disconnect, manual
    teardown), any in-flight SSE streaming connections for that instance
    hang indefinitely. The API never closes the response stream, so clients
    block forever waiting for more chunks.
    
    ## Changes
    
    - Listen for `InstanceDeleted` events in the API event loop
    - Add `_close_streams_for_instance()` to find and close any active
    text/image generation queues tied to tasks on the deleted instance
    - Add unit tests covering text gen, image gen, and
    unrelated-instance-not-closed scenarios
    
    ## Why It Works
    
    When an instance is deleted, we iterate `state.tasks` to find commands
    running on that instance, then close and remove their send-side queue
    handles. This causes the SSE generator to terminate, unblocking the
    client.
    
    ## Test Plan
    
    ### Manual Testing
    - This was causing issues for me on another branch (integration tests).
    Including this fix solved the issue
    
    ### Automated Testing
    - `test_instance_deleted_stream_cleanup.py`: 3 tests covering text gen
    cleanup, image gen cleanup, and ensuring unrelated streams are not
    affected
---
 src/exo/api/main.py                                | 26 ++++++
 .../tests/test_instance_deleted_stream_cleanup.py  | 93 ++++++++++++++++++++++
 2 files changed, 119 insertions(+)

diff --git a/src/exo/api/main.py b/src/exo/api/main.py
index 3c183f4f..73488ff0 100644
--- a/src/exo/api/main.py
+++ b/src/exo/api/main.py
@@ -172,10 +172,20 @@ from exo.shared.types.events import (
     ChunkGenerated,
     Event,
     IndexedEvent,
+    InstanceDeleted,
     TracesMerged,
 )
 from exo.shared.types.memory import Memory
 from exo.shared.types.state import State
+from exo.shared.types.tasks import (
+    ImageEdits as ImageEditsTask,
+)
+from exo.shared.types.tasks import (
+    ImageGeneration as ImageGenerationTask,
+)
+from exo.shared.types.tasks import (
+    TextGeneration as TextGenerationTask,
+)
 from exo.shared.types.text_generation import TextGenerationTaskParams
 from exo.shared.types.worker.downloads import DownloadCompleted
 from exo.shared.types.worker.instances import Instance, InstanceId, InstanceMeta
@@ -1818,9 +1828,25 @@ class API:
                             await queue.send(event.chunk)
                         except (BrokenResourceError, ClosedResourceError):
                             self._text_generation_queues.pop(event.command_id, None)
+                if isinstance(event, InstanceDeleted):
+                    self._close_streams_for_instance(event.instance_id)
                 if isinstance(event, TracesMerged):
                     self._save_merged_trace(event)
 
+    def _close_streams_for_instance(self, instance_id: InstanceId) -> None:
+        """Close any active generation streams for commands running on the given instance."""
+        for task in self.state.tasks.values():
+            if task.instance_id != instance_id:
+                continue
+            if not isinstance(
+                task, (TextGenerationTask, ImageGenerationTask, ImageEditsTask)
+            ):
+                continue
+            if sender := self._text_generation_queues.pop(task.command_id, None):
+                sender.close()
+            if sender := self._image_generation_queues.pop(task.command_id, None):
+                sender.close()
+
     def _save_merged_trace(self, event: TracesMerged) -> None:
         traces = [
             TraceEvent(
diff --git a/src/exo/api/tests/test_instance_deleted_stream_cleanup.py b/src/exo/api/tests/test_instance_deleted_stream_cleanup.py
new file mode 100644
index 00000000..843554b9
--- /dev/null
+++ b/src/exo/api/tests/test_instance_deleted_stream_cleanup.py
@@ -0,0 +1,93 @@
+# pyright: reportUnusedFunction=false, reportAny=false
+"""Tests that InstanceDeleted events close active generation streams."""
+
+from unittest.mock import MagicMock
+
+from exo.api.main import API
+from exo.api.types import ImageGenerationTaskParams
+from exo.shared.types.common import CommandId, ModelId
+from exo.shared.types.state import State
+from exo.shared.types.tasks import ImageGeneration, TextGeneration
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+from exo.shared.types.worker.instances import InstanceId
+
+
+def _make_api_with_state(state: State) -> API:
+    """Create a minimal API instance with pre-set state."""
+    api = object.__new__(API)
+    api.state = state
+    api._text_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    api._image_generation_queues = {}  # pyright: ignore[reportPrivateUsage]
+    return api
+
+
+def _make_text_gen_task(
+    instance_id: InstanceId, command_id: CommandId
+) -> TextGeneration:
+    return TextGeneration(
+        instance_id=instance_id,
+        command_id=command_id,
+        task_params=TextGenerationTaskParams(
+            model=ModelId("test-model"),
+            input=[InputMessage(role="user", content="hello")],
+        ),
+    )
+
+
+def test_close_streams_for_deleted_instance() -> None:
+    """Deleting an instance closes the text generation sender for commands on that instance."""
+    instance_id = InstanceId("inst-1")
+    command_id = CommandId("cmd-1")
+    task = _make_text_gen_task(instance_id, command_id)
+
+    state = State(tasks={task.task_id: task})
+    api = _make_api_with_state(state)
+
+    sender = MagicMock()
+    api._text_generation_queues[command_id] = sender  # pyright: ignore[reportPrivateUsage]
+
+    api._close_streams_for_instance(instance_id)  # pyright: ignore[reportPrivateUsage]
+
+    sender.close.assert_called_once()
+    assert command_id not in api._text_generation_queues  # pyright: ignore[reportPrivateUsage]
+
+
+def test_close_streams_ignores_unrelated_instances() -> None:
+    """Deleting an instance does NOT close streams for commands on other instances."""
+    target_id = InstanceId("inst-delete")
+    other_id = InstanceId("inst-keep")
+    other_cmd = CommandId("cmd-keep")
+    other_task = _make_text_gen_task(other_id, other_cmd)
+
+    state = State(tasks={other_task.task_id: other_task})
+    api = _make_api_with_state(state)
+
+    sender = MagicMock()
+    api._text_generation_queues[other_cmd] = sender  # pyright: ignore[reportPrivateUsage]
+
+    api._close_streams_for_instance(target_id)  # pyright: ignore[reportPrivateUsage]
+
+    sender.close.assert_not_called()
+    assert other_cmd in api._text_generation_queues  # pyright: ignore[reportPrivateUsage]
+
+
+def test_close_streams_for_deleted_instance_image_generation() -> None:
+    """Deleting an instance closes the image generation sender for commands on that instance."""
+    instance_id = InstanceId("inst-img")
+    command_id = CommandId("cmd-img")
+    task = ImageGeneration(
+        instance_id=instance_id,
+        command_id=command_id,
+        task_params=ImageGenerationTaskParams(prompt="a cat", model="test-model"),
+    )
+
+    state = State(tasks={task.task_id: task})
+    api = _make_api_with_state(state)
+
+    sender = MagicMock()
+    api._image_generation_queues[command_id] = sender  # pyright: ignore[reportPrivateUsage]
+
+    api._close_streams_for_instance(instance_id)  # pyright: ignore[reportPrivateUsage]
+
+    sender.close.assert_called_once()
+    assert command_id not in api._image_generation_queues  # pyright: ignore[reportPrivateUsage]

← 62570227 Catch ClosedResourceError when forwarding chunks to client q · back to Exo · Fix reasoning_tokens counting for multi-token thinking tag m e2e17eaf →