claude-cli-proxy/src/main.ts

import { loadConfig } from './config.js';
import { WsClient } from './connection/ws-client.js';
import { PtyManager } from './pty/manager.js';
import type {
  WsMessage,
  SyncStatePayload,
  CreatePtyPayload,
  KillPtyPayload,
  ResizePayload,
  ListDirectoriesPayload,
} from './connection/protocol.js';
import {
  DIR_PTY_OUTPUT,
  DIR_USER_INPUT,
  encodeBinaryFrame,
  decodeBinaryFrame,
} from './connection/protocol.js';
import { readdir, stat } from 'fs/promises';
import { resolve } from 'path';
import { homedir } from 'os';
import { existsSync, readdirSync, statSync, readFileSync } from 'fs';

class ClaudeCliProxy {
  private ws: WsClient;
  private pty: PtyManager;
  private config: ReturnType<typeof loadConfig>;

  constructor() {
    this.config = loadConfig();

    this.pty = new PtyManager();

    this.ws = new WsClient({
      url: this.config.serverUrl,
      token: this.config.proxyToken,
      onMessage: (msg) => this.handleMessage(msg),
      onBinary: (data) => this.handleBinaryFrame(data),
      onConnected: () => this.onConnected(),
      onDisconnected: () => this.onDisconnected(),
    });
  }

  async start(): Promise<void> {
    const mode = this.config.isDev ? 'DEV' : 'PROD';
    console.log(`[proxy] Starting claude-cli-proxy (${mode}) → ${this.config.serverUrl}`);

    this.ws.connect();

    process.on('SIGINT', () => this.shutdown());
    process.on('SIGTERM', () => this.shutdown());

    console.log('[proxy] Running. Press Ctrl+C to stop.');
  }

  private onConnected(): void {
    console.log('[proxy] Connected to server, waiting for sync...');
  }

  private onDisconnected(): void {
    console.log('[proxy] Disconnected from server, PTY sessions continue running');
  }

  // === Binary frame handling ===

  private handleBinaryFrame(data: Buffer): void {
    const { direction, chatId, data: payload } = decodeBinaryFrame(data);
    if (direction === DIR_USER_INPUT) {
      this.pty.writeToPty(chatId, payload.toString('utf-8'));
    }
  }

  // === JSON message handling ===

  private async handleMessage(msg: WsMessage): Promise<void> {
    console.log(`[proxy] ← ${msg.type}`);
    try {
      switch (msg.type) {
        case 'sync_state':
          await this.handleSyncState(msg.payload as SyncStatePayload);
          break;
        case 'create_pty':
          this.handleCreatePty(msg.payload as CreatePtyPayload);
          break;
        case 'kill_pty':
          this.handleKillPty(msg.payload as KillPtyPayload);
          break;
        case 'resize':
          this.handleResize(msg.payload as ResizePayload);
          break;
        case 'request_replay' as any: {
          const { chatId } = msg.payload as { chatId: string };
          const replay = this.pty.getReplayBuffer(chatId);
          if (replay) {
            this.ws.sendBinary(encodeBinaryFrame(DIR_PTY_OUTPUT, chatId, replay));
          }
          break;
        }
        case 'list_directories':
          await this.handleListDirectories(msg.payload as ListDirectoriesPayload, msg.requestId);
          break;
        case 'list_sessions':
          await this.handleListSessions(msg.payload as { workDir: string }, msg.requestId);
          break;
        default:
          console.warn(`[proxy] Unknown message type: ${msg.type}`);
      }
    } catch (err: any) {
      console.error(`[proxy] Error handling ${msg.type}:`, err);
      this.ws.send({ type: 'error', payload: { message: err.message } });
    }
  }

  private async handleSyncState(payload: SyncStatePayload): Promise<void> {
    console.log(`[proxy] Syncing state: ${payload.chats.length} chats`);

    const activePtys = new Set(this.pty.listPtys());
    const chatIds = new Set(payload.chats.map((c) => c.id));

    for (const chat of payload.chats) {
      if (activePtys.has(chat.id)) {
        // PTY exists — send replay + ready
        const replay = this.pty.getReplayBuffer(chat.id);
        if (replay) {
          this.ws.sendBinary(encodeBinaryFrame(DIR_PTY_OUTPUT, chat.id, replay));
        }
        this.ws.send({ type: 'pty_ready', payload: { chatId: chat.id } });
      } else if (chat.workDir) {
        // PTY not running — create it
        this.handleCreatePty({
          chatId: chat.id,
          dir: chat.workDir,
          resumeSessionId: chat.sessionId ?? undefined,
          cols: chat.cols ?? 120,
          rows: chat.rows ?? 40,
        });
      }
    }

    // Kill PTYs that no longer have chats
    for (const chatId of activePtys) {
      if (!chatIds.has(chatId)) {
        this.pty.killPty(chatId);
        console.log(`[proxy] Killed orphan PTY: ${chatId}`);
      }
    }
  }

  private handleCreatePty(payload: CreatePtyPayload): void {
    this.pty.createPty(
      payload.chatId,
      payload.dir,
      payload.cols,
      payload.rows,
      // onData — PTY output → send to backend
      (data) => {
        this.ws.sendBinary(encodeBinaryFrame(DIR_PTY_OUTPUT, payload.chatId, data));
      },
      // onExit — PTY exited
      (exitCode) => {
        console.log(`[proxy] PTY ${payload.chatId} exited with code ${exitCode}`);
        this.ws.send({ type: 'pty_closed', payload: { chatId: payload.chatId, exitCode } });
      },
      payload.resumeSessionId,
    );

    this.ws.send({ type: 'pty_ready', payload: { chatId: payload.chatId } });
  }

  private handleKillPty(payload: KillPtyPayload): void {
    this.pty.killPty(payload.chatId);
    this.ws.send({ type: 'pty_closed', payload: { chatId: payload.chatId, exitCode: 0 } });
  }

  private handleResize(payload: ResizePayload): void {
    this.pty.resizePty(payload.chatId, payload.cols, payload.rows);
  }

  // === Directory listing ===

  private async handleListDirectories(payload: ListDirectoriesPayload, requestId?: string): Promise<void> {
    const dirPath = payload.path.replace(/^~/, process.env.HOME || '');
    const entries: Array<{ name: string; isDirectory: boolean }> = [];

    try {
      const items = await readdir(dirPath);
      for (const item of items) {
        if (item.startsWith('.')) continue;
        try {
          const itemStat = await stat(resolve(dirPath, item));
          entries.push({ name: item, isDirectory: itemStat.isDirectory() });
        } catch {
          // Skip inaccessible items
        }
      }
    } catch {
      // Directory doesn't exist or can't be read
    }

    this.ws.send({
      type: 'directories',
      requestId,
      payload: { requestId: requestId ?? '', entries },
    });
  }

  // === Session listing ===

  private async handleListSessions(payload: { workDir: string }, requestId?: string): Promise<void> {
    const expandedDir = (payload.workDir || '').replace(/^~/, homedir());
    const encoded = expandedDir.replace(/[/_]/g, '-');
    const projectDir = resolve(homedir(), '.claude', 'projects', encoded);

    const sessions: Array<{ id: string; prompt: string; modified: string; messages: number }> = [];

    if (existsSync(projectDir)) {
      const files = readdirSync(projectDir).filter((f) => f.endsWith('.jsonl'));

      for (const f of files) {
        try {
          const fullPath = resolve(projectDir, f);
          const s = statSync(fullPath);
          const data = readFileSync(fullPath, 'utf-8');
          const lines = data.split('\n').filter(Boolean);

          let firstPrompt = '';
          let messageCount = 0;
          for (const line of lines) {
            try {
              const j = JSON.parse(line);
              if (j.type === 'user' || j.type === 'assistant') messageCount++;
              if (!firstPrompt && j.type === 'user' && j.message?.content) {
                const c = j.message.content;
                firstPrompt = typeof c === 'string'
                  ? c.substring(0, 80)
                  : (Array.isArray(c) ? (c.find((b: any) => b.type === 'text')?.text?.substring(0, 80) ?? '') : '');
              }
            } catch {}
          }

          if (firstPrompt) {
            sessions.push({
              id: f.replace('.jsonl', ''),
              prompt: firstPrompt,
              modified: new Date(s.mtimeMs).toISOString(),
              messages: messageCount,
            });
          }
        } catch {}
      }

      sessions.sort((a, b) => new Date(b.modified).getTime() - new Date(a.modified).getTime());
    }

    this.ws.send({
      type: 'sessions_list',
      requestId,
      payload: { sessions: sessions.slice(0, 15) },
    });
  }

  private shutdown(): void {
    console.log('\n[proxy] Shutting down...');
    this.pty.killAll();
    this.ws.disconnect();
    process.exit(0);
  }
}

const proxy = new ClaudeCliProxy();
proxy.start().catch((err) => {
  console.error('[proxy] Fatal error:', err);
  process.exit(1);
});