theroyallab/tabbyAPI
1
0
Fork 0
mirror of https://github.com/theroyallab/tabbyAPI.git synced 2026-03-14 15:57:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

API: Fix race condition when client disconnects

Browse Source
This commit is contained in:
turboderp
2025-10-05 21:23:02 +02:00
parent 52e093ae6c
commit d672dc2137
3 changed files with 6 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
endpoints/Kobold/utils/generation.py
View File

@@ -61,10 +61,7 @@ async def _stream_collector(data: GenerateRequest, request: Request):
async for generation in generator:
if disconnect_task.done():
abort_event.set()
handle_request_disconnect(
f"Kobold generation {data.genkey} cancelled by user."
)
raise CancelledError()
text = generation.get("text")
@@ -78,7 +75,7 @@ async def _stream_collector(data: GenerateRequest, request: Request):
break
except CancelledError:
# If the request disconnects, break out
if not disconnect_task.done():
if not abort_event.is_set():
abort_event.set()
handle_request_disconnect(
f"Kobold generation {data.genkey} cancelled by user."

7
endpoints/OAI/utils/chat_completion.py
View File

@@ -348,10 +348,7 @@ async def stream_generate_chat_completion(
# Consumer loop
while True:
if disconnect_task.done():
abort_event.set()
handle_request_disconnect(
f"Chat completion generation {request.state.id} cancelled by user."
)
raise CancelledError()
generation = await gen_queue.get()
@@ -401,7 +398,7 @@ async def stream_generate_chat_completion(
except CancelledError:
# Get out if the request gets disconnected
if not disconnect_task.done():
if not abort_event.is_set():
abort_event.set()
handle_request_disconnect("Chat completion generation cancelled by user.")
except Exception:

7
endpoints/OAI/utils/completion.py
View File

@@ -226,10 +226,7 @@ async def stream_generate_completion(
# Consumer loop
while True:
if disconnect_task.done():
abort_event.set()
handle_request_disconnect(
f"Completion generation {request.state.id} cancelled by user."
)
raise CancelledError()
generation = await gen_queue.get()
@@ -248,7 +245,7 @@ async def stream_generate_completion(
except CancelledError:
# Get out if the request gets disconnected
if not disconnect_task.done():
if not abort_event.is_set():
abort_event.set()
handle_request_disconnect(
f"Completion generation {request.state.id} cancelled by user."