Add fiit_console a custom Jupyter kernel client

This patch add the new command 'fiit_console' to replace the Jupyter
Console command used in this project to connect a shell to a remote
Jupyter ipykernel hosted in the fiit runtime. This command subclass and
modify the internals of the `ZMQTerminalIPythonApp` class of the
jupyter/jupyter_console project and provides a number of useful features
in the context of this project.

- At startup, if the kernel is busy, the shell waits infinitely for the
  prompt to resume and displays the remote output (stdout/stderr
  streams).

- Echo potential long running command (`%emu_start`, `%step`, `%cont`)
  run by another kernel client.

- Optional full echo of commands run by another kernel client, via the
  environment variable `FULL_ECHO=true`.

- Stop the shell prompt when echo a command run by another kernel client
  in order to reflect the busy status of the remote kernel.

- Fix the long output latency of other kernel client output when Jupyter
  Console is configured with `include_other_output` to true, via a
  custom asynchronous polling of the iopub channel.

- Provide a consistent output for commands run by another client, rather
  than the bugged output provided by the default implementation.

- Provide a way to launch a shell by retrieving remote kernel
  information from `plugin_backend` without specify a backend port, if
  backend uses its default port.

- Provide two new commands for launching the custom Jupiter client:
    * via config: `fiit_console --existing [jupyter_json_kernel_config]`
    * via backend: `fiit console --help`

- It seems, but it is not certain, that this implementation solves the
  invalid signature problem encountered with the default Jupyter
  Console.

Previously, the `jupyter console` command from the Jupyter project was
used as frontend to provide a terminal connected to a remote Jupyter
ipykernel hosted in the fiit runtime. The shell implementation has
several issues and behaviours which are not convenient in the context of
this project.

When the shell client start, if the remote kernel is in 'busy' state
because it is executing a cell, the shell hang 60 seconds and quit, in
addition if the remote kernel stream output such as stdout/stderr the
shell display nothing. This behavior is very annoying, for example if a
remote kernel is busy because a process is emulating and debugging a
firmware executing a long run with a forward of the serial output to
stdout, the shell client displays nothing and can potentially quit
before the debugger resumes the shell prompt at a specific breakpoint.

When echoing commands of another kernel client, via the
`include_other_output` option, the output is inconsistent which is a bug
in the current implementation, see the old pending pull request
jupyter/jupyter_console#274. Furthermore, there is a long latency in the
output of the other commands caused by an inefficient pooling of the
iopub channel.

When the state of the remote kernel is 'busy', the shell client prompt
remains active. This is the default implementation of the Jupyter
console. If the `include_other_output` option is enable, the remote
output is displayed and nested with the shell prompt which is a bit
messy and makes the prompt unusable if the remote output is very
verbose. It is preferable to block the prompt and display the cell
output run by another kernel client for a subset of commands. Such as
the `%emu_start` `%step` and `%cont` commands, which are potentially
performing a long run in the remote kernel and provide important
information when the prompt is resumed, such ad the new current code
location in the emulated code, which is useful for all other kernel
client.

When several shell clients are connected to the remote kernel and after
several commands, all shells crash with the exception message 'Invalid
Signature'.

Author: VincentDary

setup.py

@@ -68,7 +68,8 @@
     },
     entry_points={
         'console_scripts': [
-            'fiit=fiit.scripts.fiit:main'
+            'fiit=fiit.scripts.fiit:main',
+            'fiit_console=fiit.frontend.fiit_console:fiit_console'
         ]
     },
     zip_safe=False,

src/fiit/frontend/fiit_console.py

@@ -0,0 +1,285 @@
+################################################################################
+#
+# Copyright 2022-2025 Vincent Dary
+#
+# This file is part of fiit.
+#
+# fiit is free software: you can redistribute it and/or modify it under the
+# terms of the GNU General Public License as published by the Free Software
+# Foundation, either version 3 of the License, or (at your option) any later
+# version.
+#
+# fiit is distributed in the hope that it will be useful, but WITHOUT ANY
+# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
+# A PARTICULAR PURPOSE. See the GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along with
+# fiit. If not, see <https://www.gnu.org/licenses/>.
+#
+################################################################################
+import uuid
+import dataclasses
+import tempfile
+import signal
+import sys
+import os
+from typing import Optional
+import asyncio
+
+import zmq
+
+from prompt_toolkit.application import get_app_or_none
+
+from jupyter_console.ptshell import ZMQTerminalInteractiveShell, ask_yes_no
+from jupyter_console.app import ZMQTerminalIPythonApp
+from jupyter_client.consoleapp import JupyterConsoleApp
+
+from fiit.plugins.backend import BACKEND_REQ_GET_BACKEND_DATA, BackendRequest
+
+
+class RemoteKernelSync(Exception):
+    pass
+
+
+class SynchronizedZmqTerminal(ZMQTerminalInteractiveShell):
+    """
+    Extend the ZMQTerminalInteractiveShell
+    """
+    include_other_output = True
+
+    ECHO_FILTER = ['%emu_start', '%es', '%step', '%s', '%cont', '%c']
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        # cache
+        self._iopub_msg_cache: Optional[dict] = None
+
+        # Warning: asyncio.Event must be initialized in the event loop
+        self._is_running_cell: Optional[asyncio.Event] = None
+        self._interact_lock: Optional[asyncio.Event] = None
+        self._is_interact_loop_stopped: Optional[asyncio.Event] = None
+
+        self._full_echo = os.getenv('FULL_ECHO', False)
+        self._other_is_running_cell_with_echo = False
+        self._other_is_running_cell = False
+
+    def init_kernel_info(self):
+        """ Subclassed to print stdout/stderr stream if kernel is busy. """
+        self.client.hb_channel.unpause()
+        msg_id = self.client.kernel_info()
+        iopub_socket = self.client.iopub_channel.socket
+        shell_socket = self.client.shell_channel.socket
+        socket_poller = zmq.Poller()
+        socket_poller.register(iopub_socket, zmq.POLLIN)
+        socket_poller.register(shell_socket, zmq.POLLIN)
+
+        while True:
+            socks = dict(socket_poller.poll())
+
+            if socks.get(shell_socket) == zmq.POLLIN:
+                reply = self.client.get_shell_msg()
+                if reply['parent_header'].get('msg_id') == msg_id:
+                    self.kernel_info = reply['content']
+                    return
+
+            elif socks.get(iopub_socket) == zmq.POLLIN:
+                msg = self.client.iopub_channel.get_msg()
+                if msg['header']['msg_type'] == 'stream':
+                    if msg['content']['name'] == "stdout":
+                        print(msg['content']['text'], end='', flush=True)
+                    elif msg['content']['name'] == 'stderr':
+                        print(msg['content']['text'], end='', flush=True)
+
+    def _init_events(self):
+        self._is_running_cell = asyncio.Event()
+        self._is_interact_loop_stopped = asyncio.Event()
+        self._interact_lock = asyncio.Event()
+        self._is_running_cell.clear()
+        self._is_interact_loop_stopped.clear()
+        self._interact_lock.clear()
+
+    @property
+    def interact_loop_is_locked(self) -> bool:
+        return self._is_interact_loop_stopped.is_set()
+
+    @property
+    def cell_is_running(self) -> bool:
+        return self._is_running_cell.is_set()
+
+    async def lock_interact_loop(self):
+        if not self.interact_loop_is_locked and not self.cell_is_running:
+            app = get_app_or_none()
+            if not app.is_done and app.is_running:
+                self._interact_lock.clear()
+                app.exit(exception=RemoteKernelSync('kernel sync'))
+                await self._is_interact_loop_stopped.wait()
+
+    def unlock_interact_loop(self):
+        if self.interact_loop_is_locked:
+            self._interact_lock.set()
+
+    async def interact(self, loop=None, display_banner=None):
+        """ Override to allow prompt freezing via `RemoteKernelSync`. """
+        while self.keep_running:
+            print('\n', end='')
+
+            try:
+                code = await self.prompt_for_code()
+            except EOFError:
+                if (not self.confirm_exit
+                        or ask_yes_no('Do you really want to exit ([y]/n)?',
+                                      'y', 'n')):
+                    self.ask_exit()
+            except RemoteKernelSync:
+                # Can fix ghost side effects of asynchronous prompt not yet exited
+                # await asyncio.sleep(0.1)
+                self._is_interact_loop_stopped.set()
+                await self._interact_lock.wait()
+                self._is_interact_loop_stopped.clear()
+
+            else:
+                if code:
+                    self._is_running_cell.set()
+                    self.run_cell(code, store_history=True)
+                    self._is_running_cell.clear()
+
+    async def handle_external_iopub(self, loop=None):
+        """
+        Override to fix inefficient and slow manual polling in parent method,
+        and allow post jupyter message render with asynchronous capability in
+        same event loop (for exemple for asynchronous event sync).
+        """
+        self._init_events()
+        poller = zmq.asyncio.Poller()
+        poller.register(self.client.iopub_channel.socket, zmq.POLLIN)
+
+        while self.keep_running:
+            events = dict(await poller.poll(0.5))
+
+            if self.client.iopub_channel.socket in events:
+                self.handle_iopub()
+                await self._post_jupyter_message_render(self._iopub_msg_cache)
+
+    async def _post_jupyter_message_render(self, msg: dict) -> None:
+        msg = self._iopub_msg_cache
+        msg_type = msg['header']['msg_type']
+
+        if (msg_type == 'execute_input'
+                and self._other_is_running_cell_with_echo):
+            await self.lock_interact_loop()
+            content = self._iopub_msg_cache['content']
+            ec = content.get('execution_count',
+                             self.execution_count - 1)
+
+            if self._pending_clearoutput:
+                print("\r", end="")
+                sys.stdout.flush()
+                sys.stdout.flush()
+                self._pending_clearoutput = False
+
+            sys.stdout.write(f'Remote In [{ec}]: {content["code"]}\n')
+            sys.stdout.flush()
+
+        elif not self._other_is_running_cell_with_echo:
+            self.unlock_interact_loop()
+
+    def _include_output(self, msg: dict) -> bool:
+        self._set_terminal_states(msg)
+        msg_type = msg['header']['msg_type']
+
+        if self._other_is_running_cell_with_echo and msg_type == 'execute_input':
+            return False  # input render from handle_iopub() is bugged for other
+        elif self._other_is_running_cell and not self._other_is_running_cell_with_echo:
+            return False  # no render for other cell running without echo
+
+        return super().include_output(msg)
+
+    def _msg_cache_wrapper(self, msg: dict) -> bool:
+        ret = self._include_output(msg)
+        self._iopub_msg_cache = msg
+        return ret
+
+    def include_output(self, msg: dict) -> bool:
+        """
+        `Include_output()` is the best place to capture iopub message since this
+        method is called just after read message on the channel iopub channel
+        in `handle_iopub()`.
+        """
+        return self._msg_cache_wrapper(msg)
+
+    def _set_terminal_states(self, msg: dict) -> None:
+        """
+        Warning:
+        This methods set the states only for this terminal layer before
+        `handle_iopub()` set states, so `ZMQTerminalInteractiveShell` states
+        are the past states (t-1), minus the `execution_count` counter which is
+        synchronized before this method call.
+        """
+        msg_type = msg['header']['msg_type']
+        from_here = self.from_here(msg)
+
+        if (self.include_other_output
+                and not from_here
+                and self._execution_state == 'busy'
+                and msg_type == 'execute_input'
+                and (self._full_echo or msg['content']['code'] in self.ECHO_FILTER)):
+            self._other_is_running_cell_with_echo = True
+        elif (self.include_other_output
+              and not from_here
+              and self._execution_state == 'busy'
+              and msg_type == 'status'
+              and msg['content']['execution_state'] == 'idle'
+              and (self._full_echo or self._other_is_running_cell_with_echo)):
+            self._other_is_running_cell_with_echo = False
+
+        if (self.include_other_output
+                and not from_here
+                and self._execution_state == 'busy'
+                and msg_type == 'execute_input'):
+            self._other_is_running_cell = True
+        elif (self.include_other_output
+              and not from_here
+              and self._execution_state == 'busy'
+              and msg_type == 'status'
+              and msg['content']['execution_state'] == 'idle'):
+            self._other_is_running_cell = False
+
+
+class SynchronizedTerminalApp(ZMQTerminalIPythonApp):
+    classes = [SynchronizedZmqTerminal] + JupyterConsoleApp.classes
+
+    def init_shell(self):
+        JupyterConsoleApp.initialize(self)
+        # relay sigint to kernel
+        signal.signal(signal.SIGINT, self.handle_sigint)
+        self.shell = SynchronizedZmqTerminal.instance(
+            parent=self,
+            manager=self.kernel_manager,
+            client=self.kernel_client,
+            confirm_exit=self.confirm_exit,
+        )
+        self.shell.own_kernel = not self.existing
+
+
+fiit_console = SynchronizedTerminalApp.launch_instance
+
+
+def fiit_console_from_backend(backend_ip: str, backend_port: str) -> None:
+    zmq_context = zmq.Context()
+    sock = zmq_context.socket(zmq.REQ)
+    sock.connect(f'tcp://{backend_ip}:{backend_port}')
+    req = BackendRequest(method=BACKEND_REQ_GET_BACKEND_DATA, id=uuid.uuid1().hex)
+    sock.send_json(dataclasses.asdict(req))
+    res = sock.recv_json()
+    sock.close()
+
+    if res.get('error') is not None:
+        print(f'error: {res["error"]["message"]}')
+        sys.exit(1)
+
+    f = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False)
+    f.write(res['result']['jupyter_client_json_config'])
+    f.close()
+    print(f'[i] Jupyter console configuration file dropped to "{f.name}".')
+    SynchronizedTerminalApp.launch_instance(['--existing', f.name])

src/fiit/frontend/jupyter/jupyter_console.py

@@ -49,7 +49,7 @@ class EventQueueInfo:
 @dataclasses.dataclass
 class BackendData:
     event_queue_info: EventQueueInfo = dataclasses.field(default_factory=EventQueueInfo)
-    jupiter_client_json_config: Optional[str] = None
+    jupyter_client_json_config: Optional[str] = None
 
 
 class ThreadSafeSingleton(type):
@@ -360,6 +360,8 @@ def run_backend_request_loop(self) -> None:
 # Backend Plugin
 ################################################################################
 
+BACKEND_REQUEST_DEFAULT_PORT = 2560
+
 
 class PluginBackend(FiitPlugin):
     NAME = 'plugin_backend'
@@ -372,7 +374,8 @@ class PluginBackend(FiitPlugin):
             'schema': {
                 'allow_remote_connection': {'type': 'boolean', 'default': True,
                                             'required': False},
-                'request_port': {'type': 'integer', 'required': True},
+                'request_port': {'type': 'integer', 'required': False,
+                                 'default': BACKEND_REQUEST_DEFAULT_PORT},
                 'event_pub_port': {'type': 'integer', 'required': True},
             }
         }
@@ -390,7 +393,7 @@ def plugin_load(
         with BackendDataContext() as backend_data:
             if emulator_shell := plugins_context.get(ctx_conf.SHELL.name):
                 if emulator_shell._remote_ipykernel:
-                    backend_data.jupiter_client_json_config = \
+                    backend_data.jupyter_client_json_config = \
                         emulator_shell.get_remote_ipkernel_client_config()
 
         backend.run_backend_request_loop()