Fix: Properly close OpenAI SSE stream on interruption
When an InterruptionFrame causes task cancellation, the OpenAI AsyncStream was not being properly closed, which could lead to connection leaks and inconsistent stream state. This fix wraps the stream processing in a try/except/finally block to: - Catch asyncio.CancelledError from interruptions - Ensure the stream is always closed via chunk_stream.close() - Log the cancellation for debugging purposes This prevents SSE connection leaks when users interrupt the bot mid-response.
This commit is contained in:
@@ -339,6 +339,7 @@ class BaseOpenAILLMService(LLMService):
|
|||||||
else self._stream_chat_completions_universal_context(context)
|
else self._stream_chat_completions_universal_context(context)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
try:
|
||||||
async for chunk in chunk_stream:
|
async for chunk in chunk_stream:
|
||||||
if chunk.usage:
|
if chunk.usage:
|
||||||
cached_tokens = (
|
cached_tokens = (
|
||||||
@@ -425,6 +426,13 @@ class BaseOpenAILLMService(LLMService):
|
|||||||
)
|
)
|
||||||
|
|
||||||
await self.run_function_calls(function_calls)
|
await self.run_function_calls(function_calls)
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
# Handle cancellation gracefully (e.g., from InterruptionFrame)
|
||||||
|
logger.debug(f"{self}: Stream processing cancelled due to interruption")
|
||||||
|
raise
|
||||||
|
finally:
|
||||||
|
# Ensure the SSE stream is properly closed to avoid connection leaks
|
||||||
|
await chunk_stream.close()
|
||||||
|
|
||||||
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
||||||
"""Process frames for LLM completion requests.
|
"""Process frames for LLM completion requests.
|
||||||
|
|||||||
Reference in New Issue
Block a user