interactive shell websockets

[gregsadetsky]

• bumps version to 0.28.0 -- this version is checked by the CLI/web UI before attempting to use this new shell endpoint
• one way to test is to checkout this new CLI branch and use disco shell -- I would recommend this
• another way to test it is using the python script below -- the terminal resizing is not great, but it works otherwise
• I've been using dev-deploy.sh to test this, it worked well!
• the UI has been updated (but not released yet) as well - https://github.com/letsdiscodev/ui/tree/interactive-shell

test_shell.py

#!/usr/bin/env python3
"""
Test client for shell WebSocket endpoint.
Usage: python test_shell.py <host> <project> --token TOKEN
Example: python test_shell.py example.ondis.co flask --token <api-key>
"""
import asyncio
import json
import os
import signal
import sys
import termios
import tty

import websockets


async def main():
    if len(sys.argv) < 3 or "--token" not in sys.argv:
        print("Usage: python test_shell.py <host> <project> --token TOKEN")
        print(
            "Example: python test_shell.py example.ondis.co flask --token <api-key>"
        )
        sys.exit(1)

    host = sys.argv[1]
    project = sys.argv[2]

    # Parse --token argument
    token_idx = sys.argv.index("--token")
    if token_idx + 1 >= len(sys.argv):
        print("Error: --token requires a value")
        sys.exit(1)
    token = sys.argv[token_idx + 1]

    # Determine scheme based on port or host
    if "localhost" in host or "127.0.0.1" in host:
        scheme = "ws"
    else:
        scheme = "wss"

    url = f"{scheme}://{host}/api/projects/{project}/shell"

    print(f"Connecting to {url}...")

    old_settings = termios.tcgetattr(sys.stdin.fileno())

    try:
        async with websockets.connect(url) as ws:
            print("Connected. Authenticating...")
            await ws.send(json.dumps({"token": token}))

            response = await ws.recv()
            msg = json.loads(response)
            if msg.get("type") != "connected":
                print(f"Unexpected response: {response}")
                return

            print(f"Authenticated! Container: {msg.get('container')}")
            print("Entering raw mode... (Ctrl+D to exit)\n")

            tty.setraw(sys.stdin.fileno())

            rows, cols = os.get_terminal_size()
            await ws.send(json.dumps({"type": "resize", "rows": rows, "cols": cols}))

            # Handle terminal resize (SIGWINCH)
            resize_event = asyncio.Event()
            loop = asyncio.get_event_loop()

            def on_sigwinch(signum, frame):
                loop.call_soon_threadsafe(resize_event.set)

            signal.signal(signal.SIGWINCH, on_sigwinch)

            exit_event = asyncio.Event()

            async def handle_resize():
                while not exit_event.is_set():
                    await asyncio.sleep(0.1)
                    if resize_event.is_set():
                        resize_event.clear()
                        rows, cols = os.get_terminal_size()
                        try:
                            await ws.send(json.dumps({"type": "resize", "rows": rows, "cols": cols}))
                        except Exception:
                            break

            async def read_stdin():
                reader = asyncio.StreamReader()
                await loop.connect_read_pipe(
                    lambda: asyncio.StreamReaderProtocol(reader), sys.stdin
                )
                while not exit_event.is_set():
                    try:
                        data = await asyncio.wait_for(reader.read(1024), timeout=0.5)
                        if data:
                            await ws.send(data)
                    except asyncio.TimeoutError:
                        continue
                    except Exception:
                        break

            async def read_ws():
                try:
                    async for msg in ws:
                        if isinstance(msg, bytes):
                            os.write(sys.stdout.fileno(), msg)
                except websockets.exceptions.ConnectionClosed:
                    pass
                finally:
                    exit_event.set()

            await asyncio.gather(read_stdin(), read_ws(), handle_resize())

    except websockets.exceptions.ConnectionClosed as e:
        print(f"\nConnection closed: {e.code} - {e.reason}")
    except KeyboardInterrupt:
        print("\nInterrupted")
    finally:
        termios.tcsetattr(sys.stdin.fileno(), termios.TCSADRAIN, old_settings)
        print()


if __name__ == "__main__":
    asyncio.run(main())

[gregsadetsky]

note/comments:

• maybe switch away from sleep(0.01)
• /bin/bash is hardcoded - maybe have a fallback of /bin/sh and/or also optionally allow passing the shell command via the CLI/and even ui...?
• maybe add a heartbeat/ping
• call docker stop/kill on the container rather than proc.terminate()?

EDIT: all done