kiwi-code 0.0.4__py3-none-any.whl

kiwi_runtime/main.py

@@ -0,0 +1,989 @@
+#!/usr/bin/env python3
+"""
+Kiwi AI CLI Agent — connects to the server via WebSocket and executes
+terminal commands on behalf of the LLM agent.
+
+Authenticates with email and password.
+
+Usage:
+    kiwi connect --server local
+    kiwi connect --server dev
+    kiwi connect --server wss://custom.server.com
+    kiwi connect --server dev --scope full
+    kiwi connect --server dev --allow /path/to/extra/dir
+"""
+
+import argparse
+import asyncio
+import getpass
+import json
+import os
+import re
+import shlex
+import signal
+import sys
+
+IS_WINDOWS = sys.platform == "win32"
+
+if not IS_WINDOWS:
+    import fcntl
+    import pty
+    import struct
+    import termios
+
+import httpx
+import websockets
+
+MAX_OUTPUT_BYTES = 50 * 1024  # 50KB cap on command output
+
+# Sentinel must match the server-side terminal_tool.py so blocked commands
+# are detected as "completed" instead of timing out.
+_SENTINEL = "__CMD_DONE__"
+
+
+def _clean_cmd_for_display(data: str) -> str:
+    """Strip the sentinel wrapper from a PTY input command for display."""
+    cmd = data.strip()
+    # Remove sentinel suffix: '; echo __CMD_DONE__$?__CMD_DONE__' (Unix)
+    # or '& echo __CMD_DONE__%ERRORLEVEL%__CMD_DONE__' (Windows)
+    idx = cmd.find(_SENTINEL)
+    if idx > 0:
+        cmd = cmd[:idx]
+        # Remove the separator before the sentinel ('; echo ' or '& echo ')
+        for sep in ("; echo ", "& echo "):
+            if cmd.endswith(sep):
+                cmd = cmd[:-len(sep)]
+                break
+            # Handle without space
+            if cmd.endswith(sep.rstrip()):
+                cmd = cmd[:-len(sep.rstrip())]
+                break
+    return cmd.strip()
+
+
+SERVER_PRESETS = {
+    "app": {"ws": "wss://api.meetkiwi.ai", "http": "https://api.meetkiwi.ai"},
+    "dev": {"ws": "wss://dev.api.myautobots.com", "http": "https://dev.api.myautobots.com"},
+    "local": {"ws": "ws://localhost:8000", "http": "http://localhost:8000"},
+}
+
+# ---------------------------------------------------------------------------
+# ANSI colors
+# ---------------------------------------------------------------------------
+
+RESET = "\033[0m"
+BOLD = "\033[1m"
+
+# Brand color — ANSI cyan (works reliably across all terminals)
+C = "\033[36m"
+CB = "\033[1;36m"
+GREY = "\033[90m"
+RED = "\033[91m"
+YELLOW = "\033[93m"
+GREEN = "\033[92m"
+
+
+# ---------------------------------------------------------------------------
+# Banner — Kiwi logo (from SVG) + KIWI AI text side by side
+# ---------------------------------------------------------------------------
+
+# Logo: 32 wide x 16 tall (half-block, perfectly symmetric from SVG)
+LOGO = [
+    "              ▄██▄              ",
+    "              ████              ",
+    "    ▄██▄▄     ████     ▄▄██▄    ",
+    "    ▀█████▄   █▀▀█   ▄█████▀   ",
+    "     ▀███▀▀█  ████  █▀▀███▀    ",
+    "       ▀█▄  █ ████ █  ▄█▀      ",
+    "          ▀▀▄██████▄▀▀         ",
+    "▄██████▀▀█▄██████████▄█▀▀██████▄",
+    "▀██████▄▄█▀██████████▀█▄▄██████▀",
+    "          ▄▄▀██████▀▄▄         ",
+    "       ▄█▀  █ ████ █  ▀█▄      ",
+    "     ▄███▄▄█  ████  █▄▄███▄    ",
+    "    ▄█████▀   █▄▄█   ▀█████▄   ",
+    "    ▀██▀▀     ████     ▀▀██▀   ",
+    "              ████              ",
+    "              ▀██▀              ",
+]
+
+# Text: 6 tall
+TEXT = [
+    "██╗  ██╗██╗██╗    ██╗██╗     █████╗ ██╗",
+    "██║ ██╔╝██║██║    ██║██║    ██╔══██╗██║",
+    "█████╔╝ ██║██║ █╗ ██║██║    ███████║██║",
+    "██╔═██╗ ██║██║███╗██║██║    ██╔══██║██║",
+    "██║  ██╗██║╚███╔███╔╝██║    ██║  ██║██║",
+    "╚═╝  ╚═╝╚═╝ ╚══╝╚══╝ ╚═╝    ╚═╝  ╚═╝╚═╝",
+]
+
+TAGLINE = "Terminal Agent for AI-Powered Automation"
+
+
+def print_banner():
+    """Print the Kiwi AI startup banner."""
+    print()
+
+    # Side-by-side: logo (16 lines) with text (6 lines) vertically centered
+    gap = "  "
+    text_start = 5
+    logo_w = 33
+
+    for i, logo_line in enumerate(LOGO):
+        padded_logo = logo_line.ljust(logo_w)
+        text_idx = i - text_start
+        if 0 <= text_idx < len(TEXT):
+            row = f"{padded_logo}{gap}{TEXT[text_idx]}"
+        elif text_idx == len(TEXT) + 1:
+            row = f"{padded_logo}{gap}{TAGLINE}"
+        else:
+            row = padded_logo
+
+        print(f"  {CB}{row}{RESET}")
+
+    print(f"  {GREY}{'━' * 72}{RESET}")
+    print()
+
+
+def print_status(icon: str, message: str, color: str = C):
+    """Print a styled status line."""
+    print(f"  {color}{icon}{RESET}  {message}")
+
+
+def print_section(title: str):
+    """Print a section header."""
+    print(f"\n  {CB}{title}{RESET}")
+    print(f"  {GREY}{'─' * 40}{RESET}")
+
+
+def print_cmd_log(request_id: str, message: str, success: bool | None = None):
+    """Print a command execution log line."""
+    tag = f"{GREY}[{request_id[:8]}]{RESET}"
+    if success is True:
+        icon = f"{GREEN}>{RESET}"
+    elif success is False:
+        icon = f"{RED}x{RESET}"
+    else:
+        icon = f"{C}>{RESET}"
+    print(f"  {icon} {tag} {message}")
+
+
+def print_pty_log(session_id: str, message: str, level: str = "info"):
+    """Print a PTY session log line."""
+    tag = f"{GREY}[PTY {session_id[:8]}]{RESET}"
+    if level == "error":
+        icon = f"{RED}x{RESET}"
+    elif level == "success":
+        icon = f"{GREEN}>{RESET}"
+    else:
+        icon = f"{C}~{RESET}"
+    print(f"  {icon} {tag} {message}")
+
+
+# ---------------------------------------------------------------------------
+# Path validator for restricted mode
+# ---------------------------------------------------------------------------
+
+# Regex to detect Windows absolute paths like C:\ or D:/
+_WIN_ABS_PATH = re.compile(r'^[A-Za-z]:[/\\]')
+# Regex to detect UNC paths like \\server\share
+_UNC_PATH = re.compile(r'^\\\\')
+# Windows environment variables like %USERPROFILE%
+_WIN_ENV_VAR = re.compile(r'%([^%]+)%')
+
+# Regex patterns to find paths embedded anywhere in a command string
+_EMBEDDED_UNIX_PATH = re.compile(r'(?:^|[\s\'";=])(/[^\s\'";|&>]+)')
+_EMBEDDED_HOME_PATH = re.compile(r'(?:^|[\s\'";=])(~[^\s\'";|&>]*)')
+_EMBEDDED_WIN_PATH = re.compile(r'(?:^|[\s\'";=])([A-Za-z]:[/\\][^\s\'";|&>]*)')
+_EMBEDDED_UNC_PATH = re.compile(r'(?:^|[\s\'";=])(\\\\[^\s\'";|&>]+)')
+_EMBEDDED_DOTDOT = re.compile(r'(?:^|[\s\'";=])(\.\.[/\\][^\s\'";|&>]*)')
+
+
+def _scan_paths(command: str) -> list[str]:
+    """Scan the raw command string for embedded path patterns."""
+    paths = []
+    patterns = [_EMBEDDED_HOME_PATH, _EMBEDDED_WIN_PATH, _EMBEDDED_UNC_PATH, _EMBEDDED_DOTDOT]
+    # Unix absolute paths don't exist on Windows; /x tokens are command flags
+    if not IS_WINDOWS:
+        patterns.insert(0, _EMBEDDED_UNIX_PATH)
+    for pattern in patterns:
+        for m in pattern.finditer(command):
+            p = m.group(1).rstrip(";|&'\"")
+            if not p:
+                continue
+            # Skip standard system device paths
+            if p in _SAFE_PATHS or any(p.startswith(sp + "/") for sp in _SAFE_PATHS):
+                continue
+            paths.append(p)
+    return paths
+
+
+def _get_process_cwd(pid: int) -> str | None:
+    """Get the current working directory of a process by PID."""
+    try:
+        import psutil
+        return psutil.Process(pid).cwd()
+    except Exception:
+        return None
+
+
+def validate_command(
+    command: str, allowed_dirs: list[str], pid: int | None = None
+) -> tuple[bool, str]:
+    """Check if a command references paths outside the allowed directories.
+
+    Returns (True, "") if allowed, (False, reason) if blocked.
+    """
+    # Use the shell's actual cwd for resolving relative paths (e.g. ..)
+    base_dir = allowed_dirs[0]
+    if pid is not None:
+        real_cwd = _get_process_cwd(pid)
+        if real_cwd:
+            base_dir = real_cwd
+
+    # Token-based extraction
+    try:
+        tokens = shlex.split(command)
+    except ValueError:
+        tokens = command.split()
+
+    candidates = []
+    for token in tokens:
+        path = _extract_path(token)
+        if path is not None:
+            candidates.append((path, False))
+
+    # Regex scan for embedded paths (catches paths inside quoted strings)
+    for path in _scan_paths(command):
+        candidates.append((path, True))
+
+    # Deduplicate
+    seen = set()
+    for path, from_regex in candidates:
+        if path in seen:
+            continue
+        seen.add(path)
+        resolved = _resolve_path(path, base_dir)
+        if resolved is None:
+            continue
+        # Regex-scanned paths may be string literals (e.g. echo "hello /world"),
+        # so only flag them if the path actually exists on disk.
+        if from_regex and not os.path.exists(resolved):
+            continue
+        if not _is_within_allowed(resolved, allowed_dirs):
+            return False, (
+                f"Restricted: '{path}' resolves to '{resolved}' "
+                f"which is outside the allowed directories"
+            )
+
+    return True, ""
+
+
+_SAFE_PATHS = frozenset({
+    "/dev/null", "/dev/stdin", "/dev/stdout", "/dev/stderr",
+    "/dev/zero", "/dev/random", "/dev/urandom", "/dev/tty",
+    "/dev/fd", "/proc/self",
+})
+
+
+def _extract_path(token: str) -> str | None:
+    """Extract a file path from a token, or None if it doesn't look like one."""
+    # Strip trailing shell metacharacters (;, &&, ||, |, &)
+    token = token.rstrip(";|&")
+    if not token:
+        return None
+    # Strip common shell prefixes that precede paths in redirections
+    for prefix in (">", ">>", "<", "2>", "2>>", "&>", "&>>"):
+        if token.startswith(prefix) and len(token) > len(prefix):
+            token = token[len(prefix):]
+            break
+
+    # Whitelist standard system device paths
+    if token in _SAFE_PATHS or any(token.startswith(p + "/") for p in _SAFE_PATHS):
+        return None
+
+    # Home directory references
+    if token.startswith("~"):
+        return token
+
+    # Unix absolute paths (skip on Windows — /x tokens are command flags like /d, /b, /ad)
+    if token.startswith("/"):
+        if IS_WINDOWS:
+            return None
+        return token
+
+    # Windows absolute paths
+    if IS_WINDOWS:
+        if _WIN_ABS_PATH.match(token):
+            return token
+        if _UNC_PATH.match(token):
+            return token
+        # Expand %VAR% and check
+        expanded = _WIN_ENV_VAR.sub(_expand_win_var, token)
+        if expanded != token and (_WIN_ABS_PATH.match(expanded) or expanded.startswith("/")):
+            return expanded
+
+    # Relative paths with .. traversal
+    if ".." in token.split(os.sep) or ".." in token.split("/"):
+        return token
+
+    return None
+
+
+def _expand_win_var(match: re.Match) -> str:
+    """Expand a Windows %VAR% environment variable."""
+    return os.environ.get(match.group(1), match.group(0))
+
+
+def _resolve_path(path: str, base_dir: str) -> str | None:
+    """Resolve a path to its real absolute location."""
+    try:
+        # Expand ~ to home directory
+        path = os.path.expanduser(path)
+
+        # Make relative paths absolute against base_dir
+        if not os.path.isabs(path):
+            path = os.path.join(base_dir, path)
+
+        # Normalize (resolve ..) and resolve symlinks
+        return os.path.realpath(path)
+    except (OSError, ValueError):
+        return None
+
+
+def _is_within_allowed(resolved_path: str, allowed_dirs: list[str]) -> bool:
+    """Check if a resolved path is within any of the allowed directories."""
+    resolved = os.path.normcase(resolved_path)
+    # macOS filesystem is typically case-insensitive
+    if sys.platform == "darwin":
+        resolved = resolved.lower()
+    for allowed in allowed_dirs:
+        allowed_norm = os.path.normcase(allowed)
+        if sys.platform == "darwin":
+            allowed_norm = allowed_norm.lower()
+        if resolved == allowed_norm or resolved.startswith(allowed_norm + os.sep):
+            return True
+    return False
+
+
+# ---------------------------------------------------------------------------
+# Core logic
+# ---------------------------------------------------------------------------
+
+async def login(http_base_url: str, email: str, password: str) -> str:
+    """Authenticate with email/password and return the JWT access token."""
+    url = f"{http_base_url}/v1/auth/token"
+    async with httpx.AsyncClient(timeout=30) as client:
+        resp = await client.post(
+            url,
+            data={"username": email, "password": password},
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+        )
+        if resp.status_code != 200:
+            detail = resp.text
+            try:
+                detail = resp.json().get("detail", resp.text)
+            except Exception:
+                pass
+            raise Exception(f"Login failed (HTTP {resp.status_code}): {detail}")
+
+        body = resp.json()
+        token = body.get("access_token")
+        if not token:
+            raise Exception(f"No access_token in response: {body}")
+        return token
+
+
+async def run_command(
+    command: str,
+    timeout: int = 120,
+    mode: str = "restricted",
+    allowed_dirs: list[str] | None = None,
+) -> dict:
+    """Execute a shell command and return stdout, stderr, exit_code."""
+    # Validate in restricted mode
+    if mode == "restricted" and allowed_dirs:
+        ok, reason = validate_command(command, allowed_dirs)
+        if not ok:
+            return {"stdout": "", "stderr": reason, "exit_code": 1}
+
+    try:
+        cwd = allowed_dirs[0] if allowed_dirs else None
+        proc = await asyncio.create_subprocess_shell(
+            command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            cwd=cwd,
+        )
+        stdout_bytes, stderr_bytes = await asyncio.wait_for(
+            proc.communicate(), timeout=timeout
+        )
+        stdout = stdout_bytes.decode("utf-8", errors="replace")[:MAX_OUTPUT_BYTES]
+        stderr = stderr_bytes.decode("utf-8", errors="replace")[:MAX_OUTPUT_BYTES]
+        return {"stdout": stdout, "stderr": stderr, "exit_code": proc.returncode}
+    except asyncio.TimeoutError:
+        proc.kill()
+        await proc.wait()
+        return {"stdout": "", "stderr": f"Command timed out after {timeout}s", "exit_code": -1}
+    except Exception as e:
+        return {"stdout": "", "stderr": str(e), "exit_code": -1}
+
+
+if not IS_WINDOWS:
+
+    class PTYProcess:
+        """Manages a single interactive PTY session (Unix/macOS only)."""
+
+        def __init__(self, session_id: str, command: str, cols: int = 120, rows: int = 40):
+            self.session_id = session_id
+            self.command = command
+            self.cols = cols
+            self.rows = rows
+            self.master_fd = None
+            self.slave_fd = None
+            self.proc = None
+            self._read_task = None
+
+        async def start(self, ws, cwd: str | None = None):
+            """Spawn the PTY process and begin reading output."""
+            self.master_fd, self.slave_fd = pty.openpty()
+
+            winsize = struct.pack("HHHH", self.rows, self.cols, 0, 0)
+            fcntl.ioctl(self.slave_fd, termios.TIOCSWINSZ, winsize)
+
+            self.proc = await asyncio.create_subprocess_exec(
+                "/bin/sh", "-c", self.command,
+                stdin=self.slave_fd,
+                stdout=self.slave_fd,
+                stderr=self.slave_fd,
+                preexec_fn=os.setsid,
+                cwd=cwd,
+            )
+
+            os.close(self.slave_fd)
+            self.slave_fd = None
+
+            flags = fcntl.fcntl(self.master_fd, fcntl.F_GETFL)
+            fcntl.fcntl(self.master_fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
+
+            self._read_task = asyncio.create_task(self._read_loop(ws))
+
+        async def _read_loop(self, ws):
+            """Read output from master fd and send over WebSocket."""
+            loop = asyncio.get_event_loop()
+            try:
+                while True:
+                    try:
+                        data = await loop.run_in_executor(None, self._blocking_read)
+                        if data is None:
+                            break
+                        # Echo output to stdout for observability
+                        for line in data.splitlines():
+                            print_cmd_log(self.session_id, line)
+                        await ws.send(json.dumps({
+                            "type": "pty_output",
+                            "session_id": self.session_id,
+                            "data": data,
+                        }))
+                    except OSError:
+                        break
+                    except Exception as e:
+                        print_pty_log(self.session_id, f"Read error: {e}", "error")
+                        break
+
+                if self.proc:
+                    try:
+                        exit_code = await asyncio.wait_for(self.proc.wait(), timeout=5.0)
+                    except asyncio.TimeoutError:
+                        self.proc.kill()
+                        exit_code = -1
+                else:
+                    exit_code = -1
+
+                await ws.send(json.dumps({
+                    "type": "pty_exited",
+                    "session_id": self.session_id,
+                    "exit_code": exit_code,
+                }))
+            except Exception as e:
+                try:
+                    await ws.send(json.dumps({
+                        "type": "pty_error",
+                        "session_id": self.session_id,
+                        "error": str(e),
+                    }))
+                except Exception:
+                    pass
+
+        def _blocking_read(self) -> str | None:
+            """Blocking read from master fd. Returns None on EOF."""
+            import select
+            while True:
+                ready, _, _ = select.select([self.master_fd], [], [], 0.5)
+                if ready:
+                    try:
+                        data = os.read(self.master_fd, 4096)
+                        if not data:
+                            return None
+                        return data.decode("utf-8", errors="replace")
+                    except OSError:
+                        return None
+                if self.proc and self.proc.returncode is not None:
+                    try:
+                        data = os.read(self.master_fd, 4096)
+                        if data:
+                            return data.decode("utf-8", errors="replace")
+                    except OSError:
+                        pass
+                    return None
+
+        async def send_input(self, data: str) -> None:
+            """Write data to the PTY master fd."""
+            if self.master_fd is not None:
+                os.write(self.master_fd, data.encode("utf-8"))
+
+        async def close(self) -> None:
+            """Terminate the process and clean up fds."""
+            if self._read_task and not self._read_task.done():
+                self._read_task.cancel()
+                try:
+                    await self._read_task
+                except asyncio.CancelledError:
+                    pass
+
+            if self.proc and self.proc.returncode is None:
+                try:
+                    self.proc.terminate()
+                    await asyncio.wait_for(self.proc.wait(), timeout=3.0)
+                except (asyncio.TimeoutError, ProcessLookupError):
+                    try:
+                        self.proc.kill()
+                    except ProcessLookupError:
+                        pass
+
+            if self.master_fd is not None:
+                try:
+                    os.close(self.master_fd)
+                except OSError:
+                    pass
+                self.master_fd = None
+
+            if self.slave_fd is not None:
+                try:
+                    os.close(self.slave_fd)
+                except OSError:
+                    pass
+                self.slave_fd = None
+
+
+class PipeProcess:
+    """Manages a persistent shell session using subprocess pipes (Windows-compatible)."""
+
+    def __init__(self, session_id: str, command: str, cols: int = 120, rows: int = 40):
+        self.session_id = session_id
+        self.command = command
+        self.proc = None
+        self._read_task = None
+
+    async def start(self, ws, cwd: str | None = None):
+        """Spawn the shell process with piped stdin/stdout."""
+        if IS_WINDOWS:
+            shell = os.environ.get("COMSPEC", "cmd.exe")
+        else:
+            shell = self.command
+
+        self.proc = await asyncio.create_subprocess_exec(
+            shell,
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+            cwd=cwd,
+        )
+        self._read_task = asyncio.create_task(self._read_loop(ws))
+
+    async def _read_loop(self, ws):
+        """Read output from stdout and send over WebSocket."""
+        try:
+            while True:
+                data = await self.proc.stdout.read(4096)
+                if not data:
+                    break
+                decoded = data.decode("utf-8", errors="replace")
+                # Echo output to stdout for observability
+                for line in decoded.splitlines():
+                    print_cmd_log(self.session_id, line)
+                await ws.send(json.dumps({
+                    "type": "pty_output",
+                    "session_id": self.session_id,
+                    "data": decoded,
+                }))
+
+            if self.proc:
+                try:
+                    exit_code = await asyncio.wait_for(self.proc.wait(), timeout=5.0)
+                except asyncio.TimeoutError:
+                    self.proc.kill()
+                    exit_code = -1
+            else:
+                exit_code = -1
+
+            await ws.send(json.dumps({
+                "type": "pty_exited",
+                "session_id": self.session_id,
+                "exit_code": exit_code,
+            }))
+        except Exception as e:
+            try:
+                await ws.send(json.dumps({
+                    "type": "pty_error",
+                    "session_id": self.session_id,
+                    "error": str(e),
+                }))
+            except Exception:
+                pass
+
+    async def send_input(self, data: str) -> None:
+        """Write data to the process stdin."""
+        if self.proc and self.proc.stdin:
+            self.proc.stdin.write(data.encode("utf-8"))
+            await self.proc.stdin.drain()
+
+    async def close(self) -> None:
+        """Terminate the process and clean up."""
+        if self._read_task and not self._read_task.done():
+            self._read_task.cancel()
+            try:
+                await self._read_task
+            except asyncio.CancelledError:
+                pass
+
+        if self.proc and self.proc.returncode is None:
+            try:
+                self.proc.terminate()
+                await asyncio.wait_for(self.proc.wait(), timeout=3.0)
+            except (asyncio.TimeoutError, ProcessLookupError):
+                try:
+                    self.proc.kill()
+                except ProcessLookupError:
+                    pass
+
+
+async def connect(
+    ws_url: str,
+    token: str,
+    mode: str = "restricted",
+    allowed_dirs: list[str] | None = None,
+):
+    """Connect to the server WebSocket and process commands."""
+    ws_endpoint = f"{ws_url}/v1/terminal/ws/connect"
+    print_status("~", f"Connecting to {BOLD}{ws_endpoint}{RESET}", C)
+
+    try:
+        async with websockets.connect(ws_endpoint) as ws:
+            await ws.send(json.dumps({"type": "auth", "token": token}))
+            auth_resp = json.loads(await ws.recv())
+
+            if auth_resp.get("type") == "error":
+                print_status("x", f"Auth failed: {auth_resp.get('message')}", RED)
+                return
+
+            if auth_resp.get("type") == "auth_ok":
+                user_id = auth_resp.get('user_id', 'unknown')
+                print_status(">", f"Connected as {BOLD}{user_id}{RESET}", GREEN)
+            else:
+                print_status("x", f"Unexpected response: {auth_resp}", RED)
+                return
+
+            print()
+            if mode == "restricted" and allowed_dirs:
+                mode_label = f"{GREEN}restricted{RESET}"
+                print(f"  {C}Ready{RESET} {GREY}| Mode: {mode_label} {GREY}| Ctrl+C to disconnect{RESET}")
+                for d in allowed_dirs:
+                    print(f"         {GREY}Allowed: {d}{RESET}")
+            else:
+                mode_label = f"{YELLOW}unrestricted{RESET}"
+                print(f"  {C}Ready{RESET} {GREY}| Mode: {mode_label} {GREY}| Ctrl+C to disconnect{RESET}")
+            print(f"  {GREY}{'─' * 50}{RESET}")
+            print()
+
+            active_pty_sessions: dict[str, PTYProcess] = {}
+
+            try:
+                async for message in ws:
+                    msg = json.loads(message)
+                    msg_type = msg.get("type")
+
+                    if msg_type == "command":
+                        request_id = msg.get("request_id", "")
+                        command = msg.get("command", "")
+                        print_cmd_log(request_id, f"$ {BOLD}{command}{RESET}")
+
+                        result = await run_command(
+                            command,
+                            mode=mode,
+                            allowed_dirs=allowed_dirs,
+                        )
+                        exit_code = result["exit_code"]
+                        success = exit_code == 0
+                        status_text = (
+                            f"{GREEN}OK{RESET}" if success
+                            else f"{RED}FAILED (exit {exit_code}){RESET}"
+                        )
+                        print_cmd_log(request_id, status_text, success)
+
+                        await ws.send(json.dumps({
+                            "type": "result",
+                            "request_id": request_id,
+                            **result,
+                        }))
+
+                    elif msg_type == "pty_start":
+                        session_id = msg.get("session_id", "")
+                        command = msg.get("command", "bash")
+                        cols = msg.get("cols", 120)
+                        rows = msg.get("rows", 40)
+                        print_pty_log(session_id, f"Starting: {BOLD}{command}{RESET}")
+
+                        try:
+                            if IS_WINDOWS:
+                                pty_proc = PipeProcess(session_id, command, cols, rows)
+                            else:
+                                pty_proc = PTYProcess(session_id, command, cols, rows)
+                            cwd = allowed_dirs[0] if allowed_dirs else None
+                            await pty_proc.start(ws, cwd=cwd)
+                            active_pty_sessions[session_id] = pty_proc
+                            started_msg = {
+                                "type": "pty_started",
+                                "session_id": session_id,
+                                "platform": sys.platform,
+                                "mode": mode,
+                            }
+                            if mode == "restricted" and allowed_dirs:
+                                started_msg["allowed_dirs"] = allowed_dirs
+                            await ws.send(json.dumps(started_msg))
+                            print_pty_log(session_id, "Started", "success")
+                        except Exception as e:
+                            print_pty_log(session_id, f"Failed: {e}", "error")
+                            await ws.send(json.dumps({
+                                "type": "pty_error",
+                                "session_id": session_id,
+                                "error": str(e),
+                            }))
+
+                    elif msg_type == "pty_input":
+                        session_id = msg.get("session_id", "")
+                        data = msg.get("data", "")
+                        display_cmd = _clean_cmd_for_display(data)
+                        if display_cmd:
+                            print_pty_log(session_id, f"$ {BOLD}{display_cmd}{RESET}")
+                        pty_proc = active_pty_sessions.get(session_id)
+                        if pty_proc:
+                            # Validate input in restricted mode
+                            if mode == "restricted" and allowed_dirs:
+                                pty_pid = pty_proc.proc.pid if pty_proc.proc else None
+                                ok, reason = validate_command(data, allowed_dirs, pid=pty_pid)
+                                if not ok:
+                                    # Include sentinel so server detects completion
+                                    err_msg = (
+                                        f"\r\n{reason}\r\n"
+                                        f"{_SENTINEL}1{_SENTINEL}\r\n"
+                                    )
+                                    await ws.send(json.dumps({
+                                        "type": "pty_output",
+                                        "session_id": session_id,
+                                        "data": err_msg,
+                                    }))
+                                    print_pty_log(session_id, f"Blocked: {reason}", "error")
+                                    continue
+                            try:
+                                await pty_proc.send_input(data)
+                            except Exception as e:
+                                print_pty_log(session_id, f"Input error: {e}", "error")
+                        else:
+                            print_pty_log(session_id, "Input for unknown session", "error")
+
+                    elif msg_type == "pty_close":
+                        session_id = msg.get("session_id", "")
+                        pty_proc = active_pty_sessions.pop(session_id, None)
+                        if pty_proc:
+                            print_pty_log(session_id, "Closing")
+                            await pty_proc.close()
+                        else:
+                            print_pty_log(session_id, "Close for unknown session", "error")
+
+                    elif msg_type == "ping":
+                        await ws.send(json.dumps({"type": "pong"}))
+
+                    else:
+                        print_status("?", f"Unknown message: {msg_type}", YELLOW)
+            finally:
+                for sid, pty_proc in active_pty_sessions.items():
+                    try:
+                        await pty_proc.close()
+                    except Exception:
+                        pass
+
+    except websockets.exceptions.ConnectionClosed as e:
+        print()
+        print_status("!", f"Connection closed: {e}", YELLOW)
+    except ConnectionRefusedError:
+        print()
+        print_status("x", f"Could not connect to {ws_endpoint}. Is the server running?", RED)
+    except Exception as e:
+        print()
+        print_status("x", f"Error: {e}", RED)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Kiwi AI CLI Agent — execute terminal commands for LLM agents"
+    )
+    subparsers = parser.add_subparsers(dest="command")
+
+    connect_parser = subparsers.add_parser("connect", help="Connect to the server")
+    connect_parser.add_argument(
+        "--server",
+        default="app",
+        help=(
+            "Server to connect to. Use a preset name (app, dev, local) "
+            "or a full URL (e.g. https://custom.server.com). "
+            "Presets: app=api.meetkiwi.ai, dev=dev.api.myautobots.com, "
+            "local=localhost:8000 (default: app)"
+        ),
+    )
+    connect_parser.add_argument(
+        "--email",
+        default=None,
+        help="Email address (will prompt if not provided)",
+    )
+    connect_parser.add_argument(
+        "--token",
+        default=None,
+        help="Access token to skip login (used by autobots-tui for shared auth)",
+    )
+    connect_parser.add_argument(
+        "--scope",
+        choices=["restricted", "full"],
+        default="restricted",
+        help="Execution scope. restricted (default) limits commands to the current directory.",
+    )
+    connect_parser.add_argument(
+        "--allow",
+        action="append",
+        dest="allow_dirs",
+        default=[],
+        metavar="PATH",
+        help="Additional allowed directory (can be specified multiple times).",
+    )
+
+    args = parser.parse_args()
+
+    if args.command == "connect":
+        print_banner()
+
+        preset = SERVER_PRESETS.get(args.server)
+        if preset:
+            ws_url = preset["ws"]
+            http_url = preset["http"]
+        else:
+            url = args.server.rstrip("/")
+            if url.startswith("wss://"):
+                ws_url = url
+                http_url = url.replace("wss://", "https://", 1)
+            elif url.startswith("ws://"):
+                ws_url = url
+                http_url = url.replace("ws://", "http://", 1)
+            elif url.startswith("https://"):
+                http_url = url
+                ws_url = url.replace("https://", "wss://", 1)
+            elif url.startswith("http://"):
+                http_url = url
+                ws_url = url.replace("http://", "ws://", 1)
+            else:
+                http_url = f"https://{url}"
+                ws_url = f"wss://{url}"
+
+        server_label = args.server if args.server in SERVER_PRESETS else "custom"
+        print_status(">", f"Server: {BOLD}{server_label}{RESET} ({http_url})", C)
+
+        # Build allowed_dirs list
+        mode = args.scope
+        allowed_dirs = None
+        if mode == "restricted":
+            allowed_dirs = [os.path.realpath(os.getcwd())]
+            for d in args.allow_dirs:
+                resolved = os.path.realpath(d)
+                if not os.path.isdir(resolved):
+                    print_status("x", f"Directory not found: {d}", RED)
+                    sys.exit(1)
+                if resolved not in allowed_dirs:
+                    allowed_dirs.append(resolved)
+
+            print_status(">", f"Mode: {BOLD}restricted{RESET}", C)
+            for d in allowed_dirs:
+                print_status(" ", f"  Allowed: {d}", GREY)
+        else:
+            print_status(">", f"Mode: {BOLD}unrestricted{RESET}", YELLOW)
+
+        print()
+
+        loop = asyncio.new_event_loop()
+
+        provided_token = args.token or os.environ.get("KIWI_AUTH_TOKEN")
+        if provided_token:
+            # Use provided token — skip interactive login
+            token = provided_token
+            print_section("Authentication")
+            print_status(">", "Using provided access token", GREEN)
+        else:
+            print_section("Authentication")
+            email = args.email
+            if email:
+                print_status(">", f"Email: {email}", GREY)
+            else:
+                email = input(f"  {C}>{RESET}  Email: ")
+            password = getpass.getpass(f"  {C}>{RESET}  Password: ")
+            print()
+
+            print_status("~", f"Authenticating with {BOLD}{http_url}{RESET}...", C)
+            try:
+                token = loop.run_until_complete(login(http_url, email, password))
+                print_status(">", "Login successful!", GREEN)
+            except Exception as e:
+                print_status("x", f"Login failed: {e}", RED)
+                loop.close()
+                sys.exit(1)
+
+        print_section("Connection")
+
+        _shutting_down = False
+
+        def shutdown():
+            nonlocal _shutting_down
+            if _shutting_down:
+                os._exit(0)
+            _shutting_down = True
+            print()
+            print_status("!", "Disconnecting...", YELLOW)
+            for task in asyncio.all_tasks(loop):
+                task.cancel()
+
+        signal.signal(signal.SIGINT, lambda *_: shutdown())
+
+        try:
+            loop.run_until_complete(connect(ws_url, token, mode=mode, allowed_dirs=allowed_dirs))
+        except asyncio.CancelledError:
+            pass
+        finally:
+            loop.close()
+            print()
+            print(f"  {GREY}{'─' * 50}{RESET}")
+            print_status(">", "Disconnected. Goodbye!", C)
+            print()
+    else:
+        print_banner()
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()