ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Preserve audio stream disconnect cancellation

Browse source
This commit is contained in:
wasimysaid 2026-06-19 21:16:29 +02:00
commit 3ca278cb79
2 changed files with 31 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

4
studio/backend/routes/inference.py
View file

@ -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(

27
tests/studio/test_stream_cancel_registration_timing.py
View file

@ -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 = {