Soak up protocol responses after a cancel to avoid outputting garbage to the shell after the kitten exits

kittens/transfer/main.py

@@ -292,6 +292,7 @@ class SendState(NameReprEnum):
     waiting_for_permission = auto()
     permission_granted = auto()
     permission_denied = auto()
+    canceled = auto()
     finished = auto()
 
 
@@ -385,6 +386,7 @@ def __init__(self, cli_opts: TransferCLIOptions, manager: SendManager):
         self.manager = manager
         self.cli_opts = cli_opts
         self.transmit_started = False
+        self.file_metadata_sent = False
 
     def send_payload(self, payload: str) -> None:
         self.write(f'\x1b]{FILE_TRANSFER_CODE};id={self.manager.request_id};')
@@ -394,6 +396,11 @@ def send_payload(self, payload: str) -> None:
     def on_file_transfer_response(self, ftc: FileTransmissionCommand) -> None:
         if ftc.id != self.manager.request_id:
             return
+        if ftc.status == 'CANCELED':
+            self.quit_loop(1)
+            return
+        if self.manager.state in (SendState.finished, SendState.canceled):
+            return
         before = self.manager.state
         self.manager.on_file_transfer_response(ftc)
         if before == SendState.waiting_for_permission:
@@ -405,6 +412,7 @@ def on_file_transfer_response(self, ftc: FileTransmissionCommand) -> None:
             if self.manager.state == SendState.permission_granted:
                 self.cmd.styled('Permission granted for this transfer', fg='green')
                 self.print()
+                self.send_file_metadata()
         self.loop_tick()
 
     def check_for_transmit_ok(self) -> None:
@@ -435,6 +443,8 @@ def on_writing_finished(self) -> None:
             self.loop_tick()
 
     def loop_tick(self) -> None:
+        if self.manager.state == SendState.waiting_for_permission:
+            return
         if self.transmit_started:
             self.transmit_next_chunk()
         else:
@@ -442,17 +452,34 @@ def loop_tick(self) -> None:
 
     def initialize(self) -> None:
         self.send_payload(self.manager.start_transfer())
-        for payload in self.manager.send_file_metadata():
-            self.send_payload(payload)
+        if self.cli_opts.permissions_password:
+            # dont wait for permission, not needed with a password and
+            # avoids a roundtrip
+            self.send_file_metadata()
+
+    def send_file_metadata(self) -> None:
+        if not self.file_metadata_sent:
+            for payload in self.manager.send_file_metadata():
+                self.send_payload(payload)
+            self.file_metadata_sent = True
+
+    def on_term(self) -> None:
+        self.cmd.styled('Terminate requested, cancelling transfer, transferred files are in undefined state', fg='red')
+        self.print()
+        self.abort_transfer(delay=2)
 
     def on_interrupt(self) -> None:
+        if self.manager.state is SendState.canceled:
+            self.print('Waiting for canceled acknowledgement from terminal, will abort in a few seconds if no response received')
+            return
         self.cmd.styled('Interrupt requested, cancelling transfer, transferred files are in undefined state', fg='red')
         self.print()
         self.abort_transfer()
 
-    def abort_transfer(self) -> None:
+    def abort_transfer(self, delay: float = 5) -> None:
         self.send_payload(FileTransmissionCommand(action=Action.cancel).serialize())
-        self.quit_loop(1)
+        self.manager.state = SendState.canceled
+        self.asyncio_loop.call_later(delay, self.quit_loop, 1)
 
 
 def send_main(cli_opts: TransferCLIOptions, args: List[str]) -> None:

kitty/file_transmission.py

@@ -61,7 +61,7 @@ class TransmissionType(NameReprEnum):
     rsync = auto()
 
 
-ErrorCode = Enum('ErrorCode', 'OK STARTED EINVAL EPERM EISDIR')
+ErrorCode = Enum('ErrorCode', 'OK STARTED CANCELED EINVAL EPERM EISDIR')
 
 
 class TransmissionError(Exception):
@@ -428,13 +428,13 @@ def handle_receive_cmd(self, cmd: FileTransmissionCommand) -> None:
                 self.drop_receive(cmd.id)
                 return
             if not ar.accepted:
-                log_error(f'File transmission command received for pending id: {cmd.id}, aborting')
+                log_error(f'File transmission command {cmd.action} received for pending id: {cmd.id}, aborting')
                 self.drop_receive(cmd.id)
                 return
             ar.last_activity_at = monotonic()
         else:
             if cmd.action is not Action.send:
-                log_error(f'File transmission command received for unknown or rejected id: {cmd.id}, ignoring')
+                log_error(f'File transmission command {cmd.action} received for unknown or rejected id: {cmd.id}, ignoring')
                 return
             if len(self.active_receives) >= MAX_ACTIVE_RECEIVES:
                 log_error('New File transmission send with too many active receives, ignoring')
@@ -445,6 +445,8 @@ def handle_receive_cmd(self, cmd: FileTransmissionCommand) -> None:
 
         if cmd.action is Action.cancel:
             self.drop_receive(ar.id)
+            if ar.send_acknowledgements:
+                self.send_status_response(ErrorCode.CANCELED, request_id=ar.id)
         elif cmd.action is Action.file:
             try:
                 df = ar.start_file(cmd)

kitty_tests/file_transmission.py

@@ -118,7 +118,7 @@ def test_file_put(self):
         ft.handle_serialized_command(serialized_cmd(action='data', data='abcd'))
         self.assertTrue(os.path.exists(dest))
         ft.handle_serialized_command(serialized_cmd(action='cancel'))
-        self.ae(ft.test_responses, [response(status='OK'), response(status='STARTED', name=dest)])
+        self.ae(ft.test_responses, [response(status='OK'), response(status='STARTED', name=dest), response(status='CANCELED')])
         self.assertFalse(ft.active_receives)
         # compress with zlib
         ft = FileTransmission()