amc/amc_server/mixins/spawn.py

import json
import os
import re
import subprocess
import time
import uuid

from amc_server.auth import validate_auth_token
from amc_server.config import SESSIONS_DIR
from amc_server.spawn_config import (
    PENDING_SPAWNS_DIR, PENDING_SPAWN_TTL,
    PROJECTS_DIR,
    _spawn_lock, _spawn_timestamps, SPAWN_COOLDOWN_SEC,
)
from amc_server.zellij import ZELLIJ_BIN, ZELLIJ_SESSION
from amc_server.logging_utils import LOGGER


def _write_pending_spawn(spawn_id, project_path, agent_type):
    """Write a pending spawn record for later correlation by discovery.

    This enables Codex session correlation since env vars don't propagate
    through Zellij's pane spawn mechanism.
    """
    PENDING_SPAWNS_DIR.mkdir(parents=True, exist_ok=True)
    pending_file = PENDING_SPAWNS_DIR / f'{spawn_id}.json'
    data = {
        'spawn_id': spawn_id,
        'project_path': str(project_path),
        'agent_type': agent_type,
        'timestamp': time.time(),
    }
    try:
        pending_file.write_text(json.dumps(data))
    except OSError:
        LOGGER.warning('Failed to write pending spawn file for %s', spawn_id)


def _cleanup_stale_pending_spawns():
    """Remove pending spawn files older than PENDING_SPAWN_TTL."""
    if not PENDING_SPAWNS_DIR.exists():
        return
    now = time.time()
    try:
        for f in PENDING_SPAWNS_DIR.glob('*.json'):
            try:
                if now - f.stat().st_mtime > PENDING_SPAWN_TTL:
                    f.unlink()
            except OSError:
                continue
    except OSError:
        pass

# Agent commands (AC-8, AC-9: full autonomous permissions)
AGENT_COMMANDS = {
    'claude': ['claude', '--dangerously-skip-permissions'],
    'codex': ['codex', '--dangerously-bypass-approvals-and-sandbox'],
}

# Module-level cache for projects list (AC-33)
_projects_cache: list[str] = []

# Characters unsafe for Zellij pane/tab names: control chars, quotes, backticks
_UNSAFE_PANE_CHARS = re.compile(r'[\x00-\x1f\x7f"\'`]')


def _sanitize_pane_name(name):
    """Sanitize a string for use as a Zellij pane name.

    Replaces control characters and quotes with underscores, collapses runs
    of whitespace into a single space, and truncates to 64 chars.
    """
    name = _UNSAFE_PANE_CHARS.sub('_', name)
    name = re.sub(r'\s+', ' ', name).strip()
    return name[:64] if name else 'unnamed'


def load_projects_cache():
    """Scan ~/projects/ and cache the list. Called on server start."""
    global _projects_cache
    try:
        projects = []
        for entry in PROJECTS_DIR.iterdir():
            if entry.is_dir() and not entry.name.startswith('.'):
                projects.append(entry.name)
        projects.sort()
        _projects_cache = projects
    except OSError:
        _projects_cache = []


class SpawnMixin:

    def _handle_spawn(self):
        """POST /api/spawn handler."""
        # Verify auth token (AC-38)
        auth_header = self.headers.get('Authorization', '')
        if not validate_auth_token(auth_header):
            self._send_json(401, {'ok': False, 'error': 'Unauthorized', 'code': 'UNAUTHORIZED'})
            return

        # Parse JSON body
        try:
            content_length = int(self.headers.get('Content-Length', 0))
            body = json.loads(self.rfile.read(content_length))
            if not isinstance(body, dict):
                self._json_error(400, 'Invalid JSON body')
                return
        except (json.JSONDecodeError, ValueError):
            self._json_error(400, 'Invalid JSON body')
            return

        project = body.get('project', '')
        agent_type = body.get('agent_type', '')

        # Validate params (returns resolved_path to avoid TOCTOU)
        validation = self._validate_spawn_params(project, agent_type)
        if 'error' in validation:
            self._send_json(400, {
                'ok': False,
                'error': validation['error'],
                'code': validation['code'],
            })
            return

        resolved_path = validation['resolved_path']
        spawn_id = str(uuid.uuid4())

        # Acquire _spawn_lock with 15s timeout
        acquired = _spawn_lock.acquire(timeout=15)
        if not acquired:
            self._send_json(503, {
                'ok': False,
                'error': 'Server busy - another spawn in progress',
                'code': 'SERVER_BUSY',
            })
            return

        try:
            # Check rate limit inside lock
            # Use None sentinel to distinguish "never spawned" from "spawned at time 0"
            # (time.monotonic() can be close to 0 on fresh process start)
            now = time.monotonic()
            last_spawn = _spawn_timestamps.get(project)
            if last_spawn is not None and now - last_spawn < SPAWN_COOLDOWN_SEC:
                remaining = SPAWN_COOLDOWN_SEC - (now - last_spawn)
                self._send_json(429, {
                    'ok': False,
                    'error': f'Rate limited - wait {remaining:.0f}s before spawning in {project}',
                    'code': 'RATE_LIMITED',
                })
                return

            # Execute spawn
            result = self._spawn_agent_in_project_tab(
                project, resolved_path, agent_type, spawn_id,
            )

            # Update timestamp only on success
            if result.get('ok'):
                _spawn_timestamps[project] = time.monotonic()

            status_code = 200 if result.get('ok') else 500
            result['spawn_id'] = spawn_id
            self._send_json(status_code, result)
        finally:
            _spawn_lock.release()

    def _validate_spawn_params(self, project, agent_type):
        """Validate spawn parameters. Returns resolved_path or error dict."""
        if not project or not isinstance(project, str):
            return {'error': 'Project name is required', 'code': 'MISSING_PROJECT'}

        # Reject whitespace-only names
        if not project.strip():
            return {'error': 'Project name is required', 'code': 'MISSING_PROJECT'}

        # Reject null bytes and control characters (U+0000-U+001F, U+007F)
        if '\x00' in project or re.search(r'[\x00-\x1f\x7f]', project):
            return {'error': 'Invalid project name', 'code': 'INVALID_PROJECT'}

        # Reject path traversal characters (/, \, ..)
        if '/' in project or '\\' in project or '..' in project:
            return {'error': 'Invalid project name', 'code': 'INVALID_PROJECT'}

        # Resolve symlinks and verify under PROJECTS_DIR
        candidate = PROJECTS_DIR / project
        try:
            resolved = candidate.resolve()
        except OSError:
            return {'error': f'Project not found: {project}', 'code': 'PROJECT_NOT_FOUND'}

        # Symlink escape check
        try:
            resolved.relative_to(PROJECTS_DIR.resolve())
        except ValueError:
            return {'error': 'Invalid project name', 'code': 'INVALID_PROJECT'}

        if not resolved.is_dir():
            return {'error': f'Project not found: {project}', 'code': 'PROJECT_NOT_FOUND'}

        if agent_type not in AGENT_COMMANDS:
            return {
                'error': f'Invalid agent type: {agent_type}. Must be one of: {", ".join(sorted(AGENT_COMMANDS))}',
                'code': 'INVALID_AGENT_TYPE',
            }

        return {'resolved_path': resolved}

    def _check_zellij_session_exists(self):
        """Check if the target Zellij session exists."""
        try:
            result = subprocess.run(
                [ZELLIJ_BIN, 'list-sessions'],
                capture_output=True,
                text=True,
                timeout=5,
            )
            if result.returncode != 0:
                return False
            # Strip ANSI escape codes (Zellij outputs colored text)
            ansi_pattern = re.compile(r'\x1b\[[0-9;]*m')
            output = ansi_pattern.sub('', result.stdout)
            # Parse line-by-line to avoid substring false positives
            for line in output.splitlines():
                # Zellij outputs "session_name [Created ...]" or just "session_name"
                session_name = line.strip().split()[0] if line.strip() else ''
                if session_name == ZELLIJ_SESSION:
                    return True
            return False
        except FileNotFoundError:
            return False
        except subprocess.TimeoutExpired:
            return False
        except OSError:
            return False

    def _wait_for_session_file(self, spawn_id, timeout=10.0):
        """Poll for a session file matching spawn_id."""
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            try:
                for f in SESSIONS_DIR.glob('*.json'):
                    try:
                        data = json.loads(f.read_text())
                        if isinstance(data, dict) and data.get('spawn_id') == spawn_id:
                            return True
                    except (json.JSONDecodeError, OSError):
                        continue
            except OSError:
                pass
            time.sleep(0.25)
        return False

    def _spawn_agent_in_project_tab(self, project, project_path, agent_type, spawn_id):
        """Spawn an agent in a project-named Zellij tab."""
        # Clean up stale pending spawns opportunistically
        _cleanup_stale_pending_spawns()

        # For Codex, write pending spawn record before launching.
        # Zellij doesn't propagate env vars to pane commands, so discovery
        # will match the session to this record by CWD + timestamp.
        # (Claude doesn't need this - amc-hook writes spawn_id directly)
        if agent_type == 'codex':
            _write_pending_spawn(spawn_id, project_path, agent_type)

        # Check session exists
        if not self._check_zellij_session_exists():
            return {
                'ok': False,
                'error': f'Zellij session "{ZELLIJ_SESSION}" not found',
                'code': 'SESSION_NOT_FOUND',
            }

        # Create/switch to project tab
        try:
            result = subprocess.run(
                [
                    ZELLIJ_BIN, '--session', ZELLIJ_SESSION,
                    'action', 'go-to-tab-name', '--create', project,
                ],
                capture_output=True,
                text=True,
                timeout=5,
            )
            if result.returncode != 0:
                return {
                    'ok': False,
                    'error': f'Failed to create tab: {result.stderr.strip() or "unknown error"}',
                    'code': 'TAB_ERROR',
                }
        except FileNotFoundError:
            return {'ok': False, 'error': f'Zellij not found at {ZELLIJ_BIN}', 'code': 'ZELLIJ_NOT_FOUND'}
        except subprocess.TimeoutExpired:
            return {'ok': False, 'error': 'Zellij tab creation timed out', 'code': 'TIMEOUT'}
        except OSError as e:
            return {'ok': False, 'error': str(e), 'code': 'SPAWN_ERROR'}

        # Build agent command
        agent_cmd = AGENT_COMMANDS[agent_type]
        pane_name = _sanitize_pane_name(f'{agent_type}-{project}')

        # Spawn pane with agent command
        env = os.environ.copy()
        env['AMC_SPAWN_ID'] = spawn_id

        try:
            result = subprocess.run(
                [
                    ZELLIJ_BIN, '--session', ZELLIJ_SESSION,
                    'action', 'new-pane',
                    '--name', pane_name,
                    '--cwd', str(project_path),
                    '--',
                ] + agent_cmd,
                env=env,
                capture_output=True,
                text=True,
                timeout=5,
            )
            if result.returncode != 0:
                return {
                    'ok': False,
                    'error': f'Failed to spawn pane: {result.stderr.strip() or "unknown error"}',
                    'code': 'SPAWN_ERROR',
                }
        except FileNotFoundError:
            return {'ok': False, 'error': f'Zellij not found at {ZELLIJ_BIN}', 'code': 'ZELLIJ_NOT_FOUND'}
        except subprocess.TimeoutExpired:
            return {'ok': False, 'error': 'Pane spawn timed out', 'code': 'TIMEOUT'}
        except OSError as e:
            return {'ok': False, 'error': str(e), 'code': 'SPAWN_ERROR'}

        # Wait for session file to appear
        found = self._wait_for_session_file(spawn_id)
        if not found:
            LOGGER.warning(
                'Session file not found for spawn_id=%s after timeout (agent may still be starting)',
                spawn_id,
            )

        return {'ok': True, 'session_file_found': found}

    def _handle_projects(self):
        """GET /api/projects - return cached projects list."""
        self._send_json(200, {'ok': True, 'projects': list(_projects_cache)})

    def _handle_projects_refresh(self):
        """POST /api/projects/refresh - refresh and return projects list."""
        load_projects_cache()
        self._send_json(200, {'ok': True, 'projects': list(_projects_cache)})

    def _handle_health(self):
        """GET /api/health - check server and Zellij status."""
        zellij_ok = self._check_zellij_session_exists()
        self._send_json(200, {
            'ok': True,
            'zellij_session': ZELLIJ_SESSION,
            'zellij_available': zellij_ok,
        })