Handle CancelledError in SSH interact session

Author: dnyanesh1011

src/prompt_toolkit/contrib/ssh/server.py

@@ -34,13 +34,11 @@ def __init__(
         self.app_session: AppSession | None = None
 
         # PipInput object, for sending input in the CLI.
-        # (This is something that we can use in the prompt_toolkit event loop,
-        # but still write date in manually.)
+        # (This is something that we can use in the prompt_toolkit event loop, but still write data in manually.)
         self._input: PipeInput | None = None
         self._output: Vt100_Output | None = None
 
-        # Output object. Don't render to the real stdout, but write everything
-        # in the SSH channel.
+        # Output object. Don't render to the real stdout, but write everything in the SSH channel.
         class Stdout:
             def write(s, data: str) -> None:
                 try:
@@ -87,8 +85,7 @@ async def _interact(self) -> None:
             raise Exception("`_interact` called before `connection_made`.")
 
         if hasattr(self._chan, "set_line_mode") and self._chan._editor is not None:
-            # Disable the line editing provided by asyncssh. Prompt_toolkit
-            # provides the line editing.
+            # Disable the line editing provided by asyncssh. Prompt_toolkit provides the line editing.
             self._chan.set_line_mode(False)
 
         term = self._chan.get_terminal_type()
@@ -100,14 +97,25 @@ async def _interact(self) -> None:
         with create_pipe_input() as self._input:
             with create_app_session(input=self._input, output=self._output) as session:
                 self.app_session = session
+
                 try:
                     await self.interact(self)
-                except BaseException:
+
+                except asyncio.CancelledError:
+                    # Expected during disconnect/shutdown.
+                    pass
+
+                except Exception:
+                    # Unexpected application error.
                     traceback.print_exc()
+
                 finally:
                     # Close the connection.
-                    self._chan.close()
-                    self._input.close()
+                    if self._chan is not None:
+                        self._chan.close()
+
+                    if self._input is not None:
+                        self._input.close()
 
     def terminal_size_changed(
         self, width: int, height: int, pixwidth: object, pixheight: object
@@ -176,4 +184,7 @@ def begin_auth(self, username: str) -> bool:
         return False
 
     def session_requested(self) -> PromptToolkitSSHSession:
-        return PromptToolkitSSHSession(self.interact, enable_cpr=self.enable_cpr)
+        return PromptToolkitSSHSession(
+            self.interact,
+            enable_cpr=self.enable_cpr,
+        )