Atomizer/atomizer-dashboard/backend/api/routes/terminal.py

"""
Terminal WebSocket for Claude Code CLI

Provides a PTY-based terminal that runs Claude Code in the dashboard.
Uses pywinpty on Windows for proper interactive terminal support.
"""

from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from typing import Optional
import asyncio
import subprocess
import sys
import os
import signal
import json

router = APIRouter()

# Store active terminal sessions
_terminal_sessions: dict = {}

# Check if winpty is available (for Windows)
try:
    from winpty import PtyProcess
    HAS_WINPTY = True
except ImportError:
    HAS_WINPTY = False


class TerminalSession:
    """Manages a Claude Code terminal session."""

    def __init__(self, session_id: str, working_dir: str):
        self.session_id = session_id
        self.working_dir = working_dir
        self.process = None
        self.websocket: Optional[WebSocket] = None
        self._read_task: Optional[asyncio.Task] = None
        self._running = False
        self._use_winpty = sys.platform == "win32" and HAS_WINPTY

    async def start(self, websocket: WebSocket):
        """Start the Claude Code process."""
        self.websocket = websocket
        self._running = True

        # Determine the claude command
        claude_cmd = "claude"

        try:
            if self._use_winpty:
                # Use winpty for proper PTY on Windows
                # Spawn claude directly - winpty handles the interactive terminal
                import shutil
                claude_path = shutil.which("claude") or "claude"

                # Ensure HOME and USERPROFILE are properly set for Claude CLI auth
                env = {**os.environ}
                env["FORCE_COLOR"] = "1"
                env["TERM"] = "xterm-256color"
                # Claude CLI looks for credentials in HOME/.claude or USERPROFILE/.claude
                # Ensure these are set correctly
                if "USERPROFILE" in env and "HOME" not in env:
                    env["HOME"] = env["USERPROFILE"]

                self.process = PtyProcess.spawn(
                    claude_path,
                    cwd=self.working_dir,
                    env=env
                )
            elif sys.platform == "win32":
                # Fallback: Windows without winpty - use subprocess
                # Run claude with --dangerously-skip-permissions for non-interactive mode
                self.process = subprocess.Popen(
                    ["cmd.exe", "/k", claude_cmd],
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    cwd=self.working_dir,
                    bufsize=0,
                    creationflags=subprocess.CREATE_NEW_PROCESS_GROUP,
                    env={**os.environ, "FORCE_COLOR": "1", "TERM": "xterm-256color"}
                )
            else:
                # On Unix, use pty
                import pty
                master_fd, slave_fd = pty.openpty()
                self.process = subprocess.Popen(
                    [claude_cmd],
                    stdin=slave_fd,
                    stdout=slave_fd,
                    stderr=slave_fd,
                    cwd=self.working_dir,
                    env={**os.environ, "TERM": "xterm-256color"}
                )
                os.close(slave_fd)
                self._master_fd = master_fd

            # Start reading output
            self._read_task = asyncio.create_task(self._read_output())

            await self.websocket.send_json({
                "type": "started",
                "message": f"Claude Code started in {self.working_dir}"
            })

        except FileNotFoundError:
            await self.websocket.send_json({
                "type": "error",
                "message": "Claude Code CLI not found. Please install it with: npm install -g @anthropic-ai/claude-code"
            })
            self._running = False
        except Exception as e:
            await self.websocket.send_json({
                "type": "error",
                "message": f"Failed to start Claude Code: {str(e)}"
            })
            self._running = False

    async def _read_output(self):
        """Read output from the process and send to WebSocket."""
        try:
            while self._running:
                if self._use_winpty:
                    # Read from winpty
                    if self.process and self.process.isalive():
                        loop = asyncio.get_event_loop()
                        try:
                            data = await loop.run_in_executor(
                                None,
                                lambda: self.process.read(4096)
                            )
                            if data:
                                await self.websocket.send_json({
                                    "type": "output",
                                    "data": data
                                })
                        except Exception:
                            break
                    else:
                        break
                elif sys.platform == "win32":
                    # Windows subprocess pipe mode
                    if self.process and self.process.poll() is None:
                        if self.process.stdout:
                            loop = asyncio.get_event_loop()
                            try:
                                # Use non-blocking read with a timeout
                                import msvcrt
                                import ctypes

                                # Read available data
                                data = await loop.run_in_executor(
                                    None,
                                    lambda: self.process.stdout.read(1)
                                )
                                if data:
                                    # Read more if available
                                    more_data = b""
                                    try:
                                        # Try to read more without blocking
                                        while True:
                                            extra = self.process.stdout.read(1)
                                            if extra:
                                                more_data += extra
                                            else:
                                                break
                                    except:
                                        pass

                                    full_data = data + more_data
                                    await self.websocket.send_json({
                                        "type": "output",
                                        "data": full_data.decode("utf-8", errors="replace")
                                    })
                            except Exception:
                                break
                    else:
                        break
                else:
                    # Unix PTY
                    loop = asyncio.get_event_loop()
                    try:
                        data = await loop.run_in_executor(
                            None,
                            lambda: os.read(self._master_fd, 4096)
                        )
                        if data:
                            await self.websocket.send_json({
                                "type": "output",
                                "data": data.decode("utf-8", errors="replace")
                            })
                    except OSError:
                        break

                await asyncio.sleep(0.01)

            # Process ended
            if self.websocket:
                exit_code = -1
                if self._use_winpty:
                    exit_code = self.process.exitstatus if self.process else -1
                elif self.process:
                    exit_code = self.process.poll() if self.process.poll() is not None else -1

                await self.websocket.send_json({
                    "type": "exit",
                    "code": exit_code
                })

        except Exception as e:
            if self.websocket:
                try:
                    await self.websocket.send_json({
                        "type": "error",
                        "message": str(e)
                    })
                except:
                    pass

    async def write(self, data: str):
        """Write input to the process."""
        if not self.process or not self._running:
            return

        try:
            if self._use_winpty:
                self.process.write(data)
            elif sys.platform == "win32":
                if self.process.stdin:
                    self.process.stdin.write(data.encode())
                    self.process.stdin.flush()
            else:
                os.write(self._master_fd, data.encode())
        except Exception as e:
            if self.websocket:
                await self.websocket.send_json({
                    "type": "error",
                    "message": f"Write error: {str(e)}"
                })

    async def resize(self, cols: int, rows: int):
        """Resize the terminal."""
        if self._use_winpty and self.process:
            try:
                self.process.setwinsize(rows, cols)
            except:
                pass
        elif sys.platform != "win32" and hasattr(self, '_master_fd'):
            import struct
            import fcntl
            import termios
            winsize = struct.pack("HHHH", rows, cols, 0, 0)
            fcntl.ioctl(self._master_fd, termios.TIOCSWINSZ, winsize)

    async def stop(self):
        """Stop the terminal session."""
        self._running = False

        if self._read_task:
            self._read_task.cancel()
            try:
                await self._read_task
            except asyncio.CancelledError:
                pass

        if self.process:
            try:
                if self._use_winpty:
                    self.process.terminate()
                elif sys.platform == "win32":
                    self.process.terminate()
                else:
                    os.kill(self.process.pid, signal.SIGTERM)
                    self.process.wait(timeout=2)
            except:
                try:
                    if hasattr(self.process, 'kill'):
                        self.process.kill()
                except:
                    pass

        if sys.platform != "win32" and hasattr(self, '_master_fd'):
            try:
                os.close(self._master_fd)
            except:
                pass


@router.websocket("/claude")
async def claude_terminal(websocket: WebSocket, working_dir: str = None, study_id: str = None):
    """
    WebSocket endpoint for Claude Code terminal.

    Query params:
        working_dir: Directory to start Claude Code in (defaults to Atomizer root)
        study_id: Optional study ID to set context for Claude

    Client -> Server messages:
        {"type": "input", "data": "user input text"}
        {"type": "resize", "cols": 80, "rows": 24}
        {"type": "stop"}

    Server -> Client messages:
        {"type": "started", "message": "..."}
        {"type": "output", "data": "terminal output"}
        {"type": "exit", "code": 0}
        {"type": "error", "message": "..."}
    """
    await websocket.accept()

    # Default to Atomizer root directory
    if not working_dir:
        working_dir = str(os.path.dirname(os.path.dirname(os.path.dirname(
            os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        ))))

    # Create session
    session_id = f"claude-{id(websocket)}"
    session = TerminalSession(session_id, working_dir)
    _terminal_sessions[session_id] = session

    try:
        # Start Claude Code
        await session.start(websocket)

        # Note: Claude is started in Atomizer root directory so it has access to:
        # - CLAUDE.md (system instructions)
        # - .claude/skills/ (skill definitions)
        # The study_id is available for the user to reference in their prompts

        # Handle incoming messages
        while session._running:
            try:
                message = await websocket.receive_json()

                if message.get("type") == "input":
                    await session.write(message.get("data", ""))
                elif message.get("type") == "resize":
                    await session.resize(
                        message.get("cols", 80),
                        message.get("rows", 24)
                    )
                elif message.get("type") == "stop":
                    break

            except WebSocketDisconnect:
                break
            except Exception as e:
                await websocket.send_json({
                    "type": "error",
                    "message": str(e)
                })

    finally:
        await session.stop()
        _terminal_sessions.pop(session_id, None)


@router.get("/status")
async def terminal_status():
    """Check if Claude Code CLI is available."""
    import shutil

    claude_path = shutil.which("claude")

    return {
        "available": claude_path is not None,
        "path": claude_path,
        "winpty_available": HAS_WINPTY,
        "message": "Claude Code CLI is available" if claude_path else "Claude Code CLI not found. Install with: npm install -g @anthropic-ai/claude-code"
    }