Fix/broad exception handling #6280

python/samples/agentchat_fastapi/app_team.py

@@ -98,9 +98,14 @@ async def chat(websocket: WebSocket):
 
     # User input function used by the team.
     async def _user_input(prompt: str, cancellation_token: CancellationToken | None) -> str:
-        data = await websocket.receive_json()
-        message = TextMessage.model_validate(data)
-        return message.content
+        try:
+            data = await websocket.receive_json()
+            message = TextMessage.model_validate(data)
+            return message.content
+        except WebSocketDisconnect:
+            # Client disconnected while waiting for input - this is the root cause of the issue
+            logger.info("Client disconnected while waiting for user input")
+            raise  # Let WebSocketDisconnect propagate to be handled by outer try/except
 
     try:
         while True:
@@ -130,21 +135,33 @@ async def _user_input(prompt: str, cancellation_token: CancellationToken | None)
                 async with aiofiles.open(history_path, "w") as file:
                     await file.write(json.dumps(history))
 
+            except WebSocketDisconnect:
+                # Client disconnected during message processing - exit gracefully
+                logger.info("Client disconnected during message processing")
+                break
             except Exception as e:
                 # Send error message to client
                 error_message = {
                     "type": "error",
                     "content": f"Error: {str(e)}",
                     "source": "system"
                 }
-                await websocket.send_json(error_message)
-                # Re-enable input after error
-                await websocket.send_json({
-                    "type": "UserInputRequestedEvent",
-                    "content": "An error occurred. Please try again.",
-                    "source": "system"
-                })
-
+                try:
+                    await websocket.send_json(error_message)
+                    # Re-enable input after error
+                    await websocket.send_json({
+                        "type": "UserInputRequestedEvent",
+                        "content": "An error occurred. Please try again.",
+                        "source": "system"
+                    })
+                except WebSocketDisconnect:
+                    # Client disconnected while sending error - exit gracefully
+                    logger.info("Client disconnected while sending error message")
+                    break
+                except Exception as send_error:
+                    logger.error(f"Failed to send error message: {str(send_error)}")
+                    break
+
     except WebSocketDisconnect:
         logger.info("Client disconnected")
     except Exception as e:
@@ -155,7 +172,12 @@ async def _user_input(prompt: str, cancellation_token: CancellationToken | None)
                 "content": f"Unexpected error: {str(e)}",
                 "source": "system"
             })
-        except:
+        except WebSocketDisconnect:
+            # Client already disconnected - no need to send
+            logger.info("Client disconnected before error could be sent")
+        except Exception:
+            # Failed to send error message - connection likely broken
+            logger.error("Failed to send error message to client")
             pass