Fix resource leak in WebSocket handler by properly awaiting cancelled tasks #253
Conversation
📝 WalkthroughWalkthroughThe WebSocket chat handler in Changes
Estimated code review effort🎯 1 (Trivial) | ⏱️ ~5 minutes Poem
Pre-merge checks and finishing touches❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
Backend/app/routes/chat.py (1)
51-51: Improve error handling and logging.Two suggestions for better error handling:
Line 51: While catching broad
Exceptionis flagged by Ruff (BLE001), it's acceptable here to ensure cleanup. However, consider whether to re-raise the exception after cleanup to propagate it appropriately.Line 59: Replace
print()with proper logging using Python'sloggingmodule for production-ready error tracking.🔎 Proposed improvements
+import logging + +logger = logging.getLogger(__name__)Then update the exception handler:
except Exception as e: listener_task.cancel() try: await listener_task except asyncio.CancelledError: pass await chat_service.disconnect(user_id, redis, db) - # Optionally log the error - print(f"Error in websocket for user {user_id}: {e}") + logger.error(f"Error in websocket for user {user_id}: {e}", exc_info=True) + # Optionally re-raise if you want the error to propagate + # raiseAlso applies to: 59-59
📜 Review details
Configuration used: defaults
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
Backend/app/routes/chat.py
🧰 Additional context used
🧬 Code graph analysis (1)
Backend/app/routes/chat.py (1)
Backend/app/services/chat_services.py (1)
disconnect(48-60)
🪛 Ruff (0.14.10)
Backend/app/routes/chat.py
51-51: Do not catch blind exception: Exception
(BLE001)
🔇 Additional comments (1)
Backend/app/routes/chat.py (1)
43-49: Excellent fix for the resource leak!The proper asyncio task cancellation pattern is correctly implemented:
cancel()schedules the cancellationawait listener_taskensures cleanup code executes (e.g., Redis pubsub unsubscribe/close inlisten_to_channel)CancelledErroris suppressed as expectedThis ensures no resource leaks when WebSocket disconnects.
| except Exception as e: | ||
| listener_task.cancel() | ||
| try: | ||
| await listener_task | ||
| except asyncio.CancelledError: | ||
| pass | ||
| await chat_service.disconnect(user_id, redis, db) |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# First, let's look at the file structure
find Backend -name "chat.py" -type f | head -20Repository: AOSSIE-Org/InPactAI
Length of output: 116
🏁 Script executed:
#!/bin/bash
# Look at the specific lines mentioned (51-57) with context
if [ -f "Backend/app/routes/chat.py" ]; then
echo "=== Lines 40-70 of Backend/app/routes/chat.py ==="
sed -n '40,70p' Backend/app/routes/chat.py | cat -n
fiRepository: AOSSIE-Org/InPactAI
Length of output: 1250
🏁 Script executed:
#!/bin/bash
# Search for listen_to_channel function
rg -n "listen_to_channel" Backend/app --type py -A 5 -B 2Repository: AOSSIE-Org/InPactAI
Length of output: 1558
🏁 Script executed:
#!/bin/bash
# Get the full listen_to_channel function implementation
rg -n "async def listen_to_channel" Backend/app/services/chat_pubsub.py -A 30Repository: AOSSIE-Org/InPactAI
Length of output: 530
🏁 Script executed:
#!/bin/bash
# Also check the full file to understand structure better
wc -l Backend/app/services/chat_pubsub.pyRepository: AOSSIE-Org/InPactAI
Length of output: 101
🏁 Script executed:
#!/bin/bash
# View the full Backend/app/services/chat_pubsub.py file if it's not too large
cat -n Backend/app/services/chat_pubsub.pyRepository: AOSSIE-Org/InPactAI
Length of output: 689
🏁 Script executed:
#!/bin/bash
# Check what chat_service.disconnect does
rg -n "def disconnect" Backend/app/services --type py -A 10Repository: AOSSIE-Org/InPactAI
Length of output: 959
Task cancellation and pubsub cleanup are correct.
The exception handler properly cancels the listener task and awaits it with CancelledError suppression. The listen_to_channel function includes a finally block that unsubscribes and closes the Redis pubsub connection, so cleanup will execute when the task is cancelled.
However, two improvements are needed:
-
Explicitly close the websocket after an exception: Currently,
chat_service.disconnect()only removes the connection from tracking and updates user status—it does not explicitly close the websocket. After a generic exception, the websocket may be left in an inconsistent state. Consider callingawait websocket.close()before cleanup. -
Replace
print()with proper logging: Use a logger instead ofprint()for production-ready error reporting.
🧰 Tools
🪛 Ruff (0.14.10)
51-51: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
In Backend/app/routes/chat.py around lines 51 to 57, update the exception
handler to (1) explicitly close the websocket before doing service cleanup by
awaiting websocket.close() so the socket is left in a consistent state, and (2)
replace any print() calls with the module logger (e.g. logger.error(...)) and
include exception information (the exception object/traceback) in the log
message; then proceed to cancel/await the listener task and call
chat_service.disconnect(user_id, redis, db) as before.
1 participant
Fixes #252
Added proper asyncio task cancellation pattern by:
This ensures the Redis pubsub cleanup in chat_pubsub.py (unsubscribe() and close()) executes properly on every WebSocket disconnection.
✅ Checklist
Summary by CodeRabbit
✏️ Tip: You can customize this high-level summary in your review settings.