From 3ca278cb794c57c601de502a02e2b620e8b8c93e Mon Sep 17 00:00:00 2001
From: wasimysaid <wasimysdev@gmail.com>
Date: Fri, 19 Jun 2026 21:16:29 +0200
Subject: [PATCH] Preserve audio stream disconnect cancellation

---
 studio/backend/routes/inference.py            |  4 +++
 .../test_stream_cancel_registration_timing.py | 27 +++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py
index da8f98b1f5..fb5426047b 100644
--- a/studio/backend/routes/inference.py
+++ b/studio/backend/routes/inference.py
@@ -4697,6 +4697,9 @@ async def openai_chat_completions(
                 _tracker.__enter__()
 
                 async def audio_input_stream():
+                    disconnect_watcher = asyncio.create_task(
+                        _await_disconnect_then_cancel(request, cancel_event)
+                    )
                     try:
                         first_chunk = ChatCompletionChunk(
                             id = completion_id,
@@ -4758,6 +4761,7 @@ async def openai_chat_completions(
                         api_monitor.fail(monitor_id, _friendly_error(e))
                         yield f"data: {json.dumps({'error': {'message': _friendly_error(e), 'type': 'server_error'}})}\n\n"
                     finally:
+                        await _stop_local_disconnect_cancel_watcher(disconnect_watcher)
                         _tracker.__exit__(None, None, None)
 
                 return _SameTaskStreamingResponse(
diff --git a/tests/studio/test_stream_cancel_registration_timing.py b/tests/studio/test_stream_cancel_registration_timing.py
index 9df54ba23c..bc1688fef4 100644
--- a/tests/studio/test_stream_cancel_registration_timing.py
+++ b/tests/studio/test_stream_cancel_registration_timing.py
@@ -211,6 +211,33 @@ def test_direct_llama_server_streams_install_disconnect_watcher():
     )
 
 
+def test_audio_input_stream_installs_disconnect_watcher():
+    audio = _async_function("audio_input_stream")
+    has_watcher = False
+    has_cleanup = False
+    for sub in ast.walk(audio):
+        if isinstance(sub, ast.Call):
+            fn = sub.func
+            if (
+                isinstance(fn, ast.Attribute)
+                and fn.attr == "create_task"
+                and isinstance(fn.value, ast.Name)
+                and fn.value.id == "asyncio"
+                and sub.args
+                and isinstance(sub.args[0], ast.Call)
+                and isinstance(sub.args[0].func, ast.Name)
+                and sub.args[0].func.id == "_await_disconnect_then_cancel"
+            ):
+                has_watcher = True
+            if isinstance(fn, ast.Name) and fn.id == "_stop_local_disconnect_cancel_watcher":
+                has_cleanup = True
+    assert has_watcher, (
+        "audio_input_stream must install a disconnect watcher so client "
+        "disconnects set cancel_event while asyncio.to_thread(next, ...) is blocked"
+    )
+    assert has_cleanup, "audio_input_stream must stop its disconnect watcher in finally"
+
+
 # ── Behavioral helpers ───────────────────────────────────────
 
 _WANTED = {