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:
James Hush
2025-10-22 09:25:44 +08:00
parent 788465cb04
commit 434d463a3f

View File

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