From f018e5462873783929c83e4a15d8b1434d615d55 Mon Sep 17 00:00:00 2001 From: dnyanesh1011 Date: Mon, 18 May 2026 14:39:02 +0530 Subject: [PATCH 1/3] Handle CancelledError in SSH interact session --- src/prompt_toolkit/contrib/ssh/server.py | 31 ++++++++++++++++-------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/src/prompt_toolkit/contrib/ssh/server.py b/src/prompt_toolkit/contrib/ssh/server.py index 7ce70f781..e4e3aad75 100644 --- a/src/prompt_toolkit/contrib/ssh/server.py +++ b/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, + ) \ No newline at end of file From 9a75cf746640fe993d01c934c2fab7c302258fb8 Mon Sep 17 00:00:00 2001 From: dnyanesh1011 Date: Mon, 18 May 2026 23:03:50 +0530 Subject: [PATCH 2/3] Fix trailing newline --- src/prompt_toolkit/contrib/ssh/server.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/prompt_toolkit/contrib/ssh/server.py b/src/prompt_toolkit/contrib/ssh/server.py index e4e3aad75..7393fb361 100644 --- a/src/prompt_toolkit/contrib/ssh/server.py +++ b/src/prompt_toolkit/contrib/ssh/server.py @@ -187,4 +187,5 @@ def session_requested(self) -> PromptToolkitSSHSession: return PromptToolkitSSHSession( self.interact, enable_cpr=self.enable_cpr, - ) \ No newline at end of file + ) + \ No newline at end of file From 6d54a90276440bb13254eb2814f57feda59dfee5 Mon Sep 17 00:00:00 2001 From: dnyanesh1011 Date: Mon, 18 May 2026 23:25:13 +0530 Subject: [PATCH 3/3] Handle CancelledError in SSH interaction --- src/prompt_toolkit/contrib/ssh/server.py | 23 +++++++++-------------- 1 file changed, 9 insertions(+), 14 deletions(-) diff --git a/src/prompt_toolkit/contrib/ssh/server.py b/src/prompt_toolkit/contrib/ssh/server.py index 7393fb361..92101700b 100644 --- a/src/prompt_toolkit/contrib/ssh/server.py +++ b/src/prompt_toolkit/contrib/ssh/server.py @@ -34,11 +34,13 @@ 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 data in manually.) + # (This is something that we can use in the prompt_toolkit event loop, + # but still write date 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: @@ -85,7 +87,8 @@ 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() @@ -106,16 +109,12 @@ async def _interact(self) -> None: pass except Exception: - # Unexpected application error. traceback.print_exc() finally: # Close the connection. - if self._chan is not None: - self._chan.close() - - if self._input is not None: - self._input.close() + self._chan.close() + self._input.close() def terminal_size_changed( self, width: int, height: int, pixwidth: object, pixheight: object @@ -184,8 +183,4 @@ def begin_auth(self, username: str) -> bool: return False def session_requested(self) -> PromptToolkitSSHSession: - return PromptToolkitSSHSession( - self.interact, - enable_cpr=self.enable_cpr, - ) - \ No newline at end of file + return PromptToolkitSSHSession(self.interact, enable_cpr=self.enable_cpr)