agentic-chat-connect/src/connectLocalController.ts

import os from 'node:os';
import path from 'node:path';
import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from '@earendil-works/pi-coding-agent';
import type { LocalAgentServerMessage } from '@pi-platform/shared';
import {
  clearReconnectCredential,
  loadReconnectCredential,
  saveReconnectCredential,
} from './credentialStore.js';
import { LocalAgentWsClient } from './wsClient.js';

export interface ConnectLocalControllerDeps {
  pi: ExtensionAPI;
  credentialBaseDir?: string;
  wsUrl?: string;
  createWsClient?: (options: ConstructorParameters<typeof LocalAgentWsClient>[0]) => LocalAgentWsClient;
}

export class ConnectLocalController {
  private readonly pi: ExtensionAPI;
  private readonly credentialBaseDir: string;
  private readonly wsUrl: string;
  private readonly createWsClient: ConnectLocalControllerDeps['createWsClient'];
  private client: LocalAgentWsClient | undefined;
  private projectName: string | undefined;
  private connectionEpoch: number | undefined;
  private remoteTaskReserved = false;
  private agentBusy = false;
  private pairWaiters: Array<{
    resolve: (value: LocalAgentServerMessage & { type: 'pair.ok' }) => void;
    reject: (error: Error) => void;
  }> = [];

  constructor(deps: ConnectLocalControllerDeps) {
    this.pi = deps.pi;
    this.credentialBaseDir = deps.credentialBaseDir ?? path.join(os.homedir(), '.pi-platform', 'connect-local');
    this.wsUrl = deps.wsUrl ?? process.env.PI_PLATFORM_LOCAL_AGENT_WS_URL ?? 'ws://127.0.0.1:3000/api/local-agent/ws';
    this.createWsClient = deps.createWsClient ?? ((options) => new LocalAgentWsClient(options));
  }

  registerCommands(): void {
    this.pi.registerCommand('connect', {
      description: 'Pair this interactive Pi session with the cloud platform',
      handler: async (args, ctx) => {
        await this.handleConnect(args, ctx);
      },
    });
    this.pi.registerCommand('disconnect', {
      description: 'Disconnect the local agent bridge',
      handler: async (_args, ctx) => {
        await this.handleDisconnect(ctx);
      },
    });
    this.pi.registerCommand('connected', {
      description: 'Show local agent bridge connection status',
      handler: async (_args, ctx) => {
        this.notifyStatus(ctx);
      },
    });
  }

  bindLifecycle(): void {
    this.pi.on('agent_start', async (_event, ctx) => {
      this.agentBusy = true;
      this.sendPresence('busy', ctx);
    });
    this.pi.on('agent_settled', async (_event, ctx) => {
      this.agentBusy = false;
      this.remoteTaskReserved = false;
      this.sendPresence('idle', ctx);
    });
    this.pi.on('session_shutdown', async () => {
      this.shutdown();
    });
  }

  scheduleAutoReconnect(): void {
    queueMicrotask(() => {
      void this.tryAutoReconnect();
    });
  }

  async handleConnect(args: string, ctx: ExtensionCommandContext): Promise<void> {
    const parts = args.trim().split(/\s+/).filter(Boolean);
    const token = parts[0];
    if (!token) {
      ctx.ui.notify('Usage: /connect <pairing-token> [project-name]', 'warning');
      return;
    }
    const projectName = parts[1] ?? path.basename(ctx.cwd);
    try {
      await this.connectWithPairingToken(token, projectName);
      ctx.ui.notify(`Connected local project "${projectName}"`, 'info');
    } catch (error) {
      const message = error instanceof Error ? error.message : String(error);
      ctx.ui.notify(`Connect failed: ${message}`, 'error');
    }
  }

  async handleDisconnect(ctx: ExtensionCommandContext): Promise<void> {
    this.shutdown();
    await clearReconnectCredential(this.credentialBaseDir);
    ctx.ui.notify('Disconnected from cloud platform', 'info');
  }

  notifyStatus(ctx: ExtensionCommandContext): void {
    if (!this.client?.isOpen || !this.projectName) {
      ctx.ui.notify('Local agent: offline', 'info');
      return;
    }
    ctx.ui.notify(`Local agent: online (${this.projectName})`, 'info');
  }

  async connectWithPairingToken(token: string, projectName: string): Promise<void> {
    await this.ensureClient();
    const pairOk = await new Promise<LocalAgentServerMessage & { type: 'pair.ok' }>((resolve, reject) => {
      const timeout = setTimeout(() => reject(new Error('Pairing timed out')), 30_000);
      this.pairWaiters.push({
        resolve: (value) => {
          clearTimeout(timeout);
          resolve(value);
        },
        reject: (error) => {
          clearTimeout(timeout);
          reject(error);
        },
      });
      try {
        this.client!.send({ type: 'pair', token, projectName });
      } catch (error) {
        clearTimeout(timeout);
        reject(error);
      }
    });

    this.projectName = pairOk.projectName;
    this.connectionEpoch = pairOk.connectionEpoch;
    await saveReconnectCredential(this.credentialBaseDir, {
      projectName: pairOk.projectName,
      credential: pairOk.reconnectCredential,
      wsUrl: this.wsUrl,
      updatedAt: new Date().toISOString(),
    });
  }

  async tryAutoReconnect(): Promise<void> {
    const stored = await loadReconnectCredential(this.credentialBaseDir);
    if (!stored?.credential || !stored.projectName) {
      return;
    }
    try {
      await this.ensureClient(stored.wsUrl || this.wsUrl);
      this.client!.send({
        type: 'reconnect',
        credential: stored.credential,
        projectName: stored.projectName,
        connectionEpoch: this.connectionEpoch,
      });
      this.projectName = stored.projectName;
    } catch {
      /* extension must not block Pi; reconnect can be retried manually */
    }
  }

  shutdown(): void {
    this.client?.close();
    this.client = undefined;
    this.projectName = undefined;
    this.connectionEpoch = undefined;
    this.remoteTaskReserved = false;
  }

  private async ensureClient(wsUrl = this.wsUrl): Promise<void> {
    if (this.client?.isOpen) {
      return;
    }
    this.client = this.createWsClient!({
      wsUrl,
      onServerMessage: (message) => this.handleServerMessage(message),
      onClose: () => {
        this.projectName = undefined;
      },
    });
    await this.client.connect();
  }

  private handleServerMessage(message: LocalAgentServerMessage): void {
    if (message.type === 'pair.ok') {
      const waiter = this.pairWaiters.shift();
      waiter?.resolve(message);
      return;
    }
    if (message.type === 'error') {
      const waiter = this.pairWaiters.shift();
      waiter?.reject(new Error(message.message));
      return;
    }
    if (message.type === 'task.dispatch') {
      void this.handleTaskDispatch(message);
    }
  }

  private handleTaskDispatch(message: Extract<LocalAgentServerMessage, { type: 'task.dispatch' }>): void {
    if (!this.client?.isOpen) {
      return;
    }
    if (this.remoteTaskReserved || this.agentBusy) {
      this.client.send({
        type: 'task.result',
        taskId: message.taskId,
        status: 'rejected',
        reason: 'agent_busy',
      });
      return;
    }

    this.remoteTaskReserved = true;
    this.client.send({ type: 'presence', status: 'busy' });
    this.pi.sendUserMessage(message.prompt);
    this.client.send({
      type: 'task.result',
      taskId: message.taskId,
      status: 'accepted',
    });
  }

  private sendPresence(status: 'idle' | 'busy', ctx: ExtensionContext): void {
    if (!this.client?.isOpen) {
      return;
    }
    if (status === 'idle' && this.remoteTaskReserved) {
      return;
    }
    if (!ctx.isIdle() && status === 'idle') {
      return;
    }
    this.client.send({ type: 'presence', status });
  }
}