@bytespell/amux 0.0.10 → 0.0.12

package/src/session.ts

@@ -0,0 +1,719 @@
+import { spawn, type ChildProcess } from 'child_process';
+import { EventEmitter } from 'events';
+import os from 'os';
+import path from 'path';
+import { Writable, Readable } from 'stream';
+
+import * as acp from '@agentclientprotocol/sdk';
+
+import { AmuxClient } from './client.js';
+import { TerminalManager } from './terminal.js';
+import { StateManager } from './state.js';
+import { parseMessageToContentBlocks } from './message-parser.js';
+import type {
+  AgentSessionConfig,
+  AgentConfig,
+  SessionRestoreInfo,
+  SessionMetadata,
+} from './types.js';
+import { AGENTS as DEFAULT_AGENTS } from './types.js';
+
+const INIT_TIMEOUT_MS = 90000;
+
+/**
+ * Events emitted by AgentSession
+ */
+export interface AgentSessionEvents {
+  ready: {
+    cwd: string;
+    sessionId: string | null;
+    sessionRestore?: SessionRestoreInfo;
+    capabilities: acp.AgentCapabilities | null;
+    agent: { type: string; name: string };
+    availableAgents: Array<{ id: string; name: string }>;
+  };
+  update: acp.SessionUpdate;
+  turn_start: Record<string, never>;
+  turn_end: Record<string, never>;
+  permission_request: {
+    requestId: string;
+    toolCall?: acp.ToolCall;
+    options: acp.PermissionOption[];
+  };
+  prompt_complete: { stopReason?: string };
+  session_created: { sessionId: string | null };
+  session_switched: { sessionId: string };
+  history_replay: {
+    previousSessionId: string;
+    events: unknown[];
+    eventCount: number;
+  };
+  error: { message: string };
+  agent_exit: { code: number | null; signal: string | null };
+  connecting: Record<string, never>;
+}
+
+/**
+ * Helper to wrap promises with timeout
+ */
+function withTimeout<T>(promise: Promise<T>, ms: number, operation: string): Promise<T> {
+  return Promise.race([
+    promise,
+    new Promise<T>((_, reject) =>
+      setTimeout(() => reject(new Error(`${operation} timed out after ${ms}ms`)), ms)
+    ),
+  ]);
+}
+
+/**
+ * AgentSession - Core ACP session management with EventEmitter interface.
+ *
+ * Emits events for all session activity. Transport-agnostic - wire up
+ * to WebSocket, HTTP, IPC, or anything else.
+ *
+ * @example
+ * ```typescript
+ * const session = new AgentSession({
+ *   instanceId: 'my-instance',
+ *   basePath: __dirname,
+ * });
+ *
+ * session.on('ready', (data) => console.log('Ready:', data));
+ * session.on('update', (data) => console.log('Update:', data));
+ *
+ * await session.spawnAgent();
+ * await session.prompt('Hello!');
+ * ```
+ */
+export class AgentSession extends EventEmitter {
+  // Session state
+  private _cwd: string;
+  private _sessionId: string | null;
+  private _agentType: string;
+  private _systemContext?: string;
+
+  // Agent process state
+  private agentProcess: ChildProcess | null = null;
+  private acpConnection: acp.ClientSideConnection | null = null;
+  private _agentCapabilities: acp.AgentCapabilities | null = null;
+
+  // Components
+  private terminalManager: TerminalManager;
+  private amuxClient: AmuxClient | null = null;
+  private stateManager: StateManager;
+
+  // Config
+  private instanceId: string;
+  private basePath: string;
+  private fixedCwd?: string;
+  private agents: Record<string, AgentConfig>;
+
+  constructor(config: AgentSessionConfig) {
+    super();
+
+    this.instanceId = config.instanceId;
+    this.basePath = config.basePath;
+    this.fixedCwd = config.fixedCwd;
+    this._systemContext = config.systemContext;
+    this.agents = DEFAULT_AGENTS;
+
+    // Create state manager
+    this.stateManager = new StateManager(config.stateDir);
+
+    // Load persisted state
+    const persistedState = this.stateManager.loadState(this.instanceId);
+    this._cwd = config.fixedCwd ?? persistedState.cwd ?? os.homedir();
+    this._sessionId = persistedState.sessionId ?? null;
+    this._agentType = config.agentType ?? persistedState.agentType ?? 'claude-code';
+
+    // Create terminal manager with cwd getter
+    this.terminalManager = new TerminalManager(() => this._cwd);
+  }
+
+  // Type-safe event emitter methods
+  override emit<K extends keyof AgentSessionEvents>(
+    event: K,
+    data: AgentSessionEvents[K]
+  ): boolean {
+    return super.emit(event, data);
+  }
+
+  override on<K extends keyof AgentSessionEvents>(
+    event: K,
+    listener: (data: AgentSessionEvents[K]) => void
+  ): this {
+    return super.on(event, listener);
+  }
+
+  override once<K extends keyof AgentSessionEvents>(
+    event: K,
+    listener: (data: AgentSessionEvents[K]) => void
+  ): this {
+    return super.once(event, listener);
+  }
+
+  /**
+   * Emit update and store to history
+   */
+  private emitUpdate(update: acp.SessionUpdate): void {
+    this.emit('update', update);
+    this.stateManager.storeEvent(this._sessionId, { type: 'session_update', update });
+  }
+
+  /**
+   * Emit turn marker and store to history
+   */
+  private emitTurnMarker(event: 'turn_start' | 'turn_end'): void {
+    this.emit(event, {});
+    this.stateManager.storeEvent(this._sessionId, { type: event });
+  }
+
+  // Getters for read-only access
+  get cwd(): string { return this._cwd; }
+  get sessionId(): string | null { return this._sessionId; }
+  get agentType(): string { return this._agentType; }
+  get agentCapabilities(): acp.AgentCapabilities | null { return this._agentCapabilities; }
+  get isConnected(): boolean { return this.acpConnection !== null; }
+  get systemContext(): string | undefined { return this._systemContext; }
+
+  /**
+   * Get current agent info
+   */
+  getAgentInfo(): { type: string; name: string } {
+    return {
+      type: this._agentType,
+      name: this.agents[this._agentType]?.name ?? 'Unknown',
+    };
+  }
+
+  /**
+   * Get available agents list
+   */
+  getAvailableAgents(): Array<{ id: string; name: string }> {
+    return Object.entries(this.agents).map(([id, a]) => ({ id, name: a.name }));
+  }
+
+  /**
+   * Save current state to disk
+   */
+  private saveState(): void {
+    this.stateManager.saveState(this.instanceId, {
+      cwd: this._cwd,
+      sessionId: this._sessionId,
+      agentType: this._agentType,
+    });
+  }
+
+  /**
+   * Name the session after the first user message
+   */
+  private maybeNameSession(userMessage: string): void {
+    if (!this._sessionId) return;
+
+    const session = this.stateManager.getSession(this._sessionId);
+    if (!session) return;
+
+    const isDefaultName = session.name.startsWith('Session ') && session.name.length <= 16;
+    if (!isDefaultName) return;
+
+    let displayName = userMessage.trim();
+    if (displayName.length > 50) {
+      displayName = displayName.slice(0, 50);
+      const lastSpace = displayName.lastIndexOf(' ');
+      if (lastSpace > 30) {
+        displayName = displayName.slice(0, lastSpace);
+      }
+      displayName += '...';
+    }
+
+    this.stateManager.updateSessionMetadata(this._sessionId, { name: displayName });
+    console.log(`[amux] Named session: "${displayName}"`);
+  }
+
+  /**
+   * Spawn and initialize the ACP agent
+   */
+  async spawnAgent(): Promise<void> {
+    if (this.agentProcess) return;
+
+    const agent = this.agents[this._agentType];
+    if (!agent) {
+      console.error('[amux] Unknown agent type:', this._agentType);
+      this.emit('error', { message: `Unknown agent type: ${this._agentType}` });
+      return;
+    }
+
+    console.log(`[amux] Spawning ${agent.name} agent in cwd:`, this._cwd);
+    this.emit('connecting', {});
+
+    const agentBin = path.join(this.basePath, 'node_modules', '.bin', agent.bin);
+    const newProcess = spawn(agentBin, [], {
+      stdio: ['pipe', 'pipe', 'pipe'],
+      env: { ...process.env },
+      cwd: this._cwd,
+    });
+    this.agentProcess = newProcess;
+
+    newProcess.stderr?.on('data', (data: Buffer) => {
+      console.error('[amux] Agent stderr:', data.toString());
+    });
+
+    newProcess.on('error', (err) => {
+      console.error('[amux] Agent process error:', err);
+      if (this.agentProcess === newProcess) {
+        this.emit('error', { message: `Failed to start agent: ${err.message}` });
+        this.agentProcess = null;
+        this.acpConnection = null;
+      }
+    });
+
+    newProcess.on('exit', (code, signal) => {
+      console.log(`[amux] Agent exited: code=${code}, signal=${signal}`);
+      if (this.agentProcess === newProcess) {
+        this.agentProcess = null;
+        this.acpConnection = null;
+        this.emit('agent_exit', { code, signal });
+      }
+    });
+
+    // Create ACP streams
+    const input = Writable.toWeb(newProcess.stdin!);
+    const output = Readable.toWeb(newProcess.stdout!);
+
+    // Create the ACP client connection
+    this.amuxClient = new AmuxClient(this.terminalManager, (event) => {
+      // Route client events through our emitter
+      if (event.type === 'session_update' && event.update) {
+        this.emitUpdate( event.update as acp.SessionUpdate);
+      } else if (event.type === 'permission_request') {
+        this.emit('permission_request', {
+          requestId: event.requestId as string,
+          toolCall: event.toolCall as acp.ToolCall | undefined,
+          options: event.options as acp.PermissionOption[],
+        });
+      }
+    });
+    const stream = acp.ndJsonStream(input, output);
+    this.acpConnection = new acp.ClientSideConnection(() => this.amuxClient!, stream);
+
+    try {
+      const initResult = await withTimeout(
+        this.acpConnection.initialize({
+          protocolVersion: acp.PROTOCOL_VERSION,
+          clientCapabilities: {
+            fs: {
+              readTextFile: true,
+              writeTextFile: true,
+            },
+            terminal: true,
+          },
+        }),
+        INIT_TIMEOUT_MS,
+        'Agent initialization'
+      );
+
+      console.log(`[amux] Connected to agent (protocol v${initResult.protocolVersion})`);
+      this._agentCapabilities = initResult.agentCapabilities ?? null;
+
+      const sessionRestore = await this.initializeSession(initResult.agentCapabilities);
+
+      this.emit('ready', {
+        cwd: this._cwd,
+        sessionId: this._sessionId,
+        sessionRestore,
+        capabilities: this._agentCapabilities,
+        agent: this.getAgentInfo(),
+        availableAgents: this.getAvailableAgents(),
+      });
+
+    } catch (err) {
+      console.error('[amux] Failed to initialize agent:', err);
+      this.emit('error', { message: `Failed to initialize agent: ${(err as Error).message}` });
+      newProcess.kill();
+      if (this.agentProcess === newProcess) {
+        this.agentProcess = null;
+        this.acpConnection = null;
+      }
+    }
+  }
+
+  /**
+   * Initialize or restore session
+   */
+  private async initializeSession(
+    capabilities: acp.AgentCapabilities | undefined
+  ): Promise<SessionRestoreInfo | undefined> {
+    if (!this.acpConnection) return undefined;
+
+    const canLoadSession = capabilities?.loadSession === true;
+    const hasExistingSession = Boolean(this._sessionId);
+    const previousSessionId = this._sessionId;
+
+    let sessionRestore: SessionRestoreInfo | undefined = undefined;
+
+    if (this._sessionId && canLoadSession) {
+      let loadSucceeded = false;
+      try {
+        await withTimeout(
+          this.acpConnection.loadSession({
+            sessionId: this._sessionId,
+            cwd: this._cwd,
+            mcpServers: [],
+          }),
+          INIT_TIMEOUT_MS,
+          'Session load'
+        );
+
+        await new Promise(resolve => setTimeout(resolve, 100));
+        sessionRestore = { restored: true };
+        loadSucceeded = true;
+
+        if (this._sessionId) {
+          this.stateManager.registerSession(this._sessionId, this._cwd, this._agentType);
+        }
+      } catch (loadErr) {
+        console.log(`[amux] Session load failed: ${(loadErr as Error).message}`);
+      }
+
+      if (!loadSucceeded) {
+        const sessionResult = await this.createNewSession();
+        sessionRestore = await this.replayHistoryFromPreviousSession(previousSessionId);
+        this.broadcastModeInfo(sessionResult);
+      }
+    } else {
+      const sessionResult = await this.createNewSession();
+
+      if (hasExistingSession && previousSessionId) {
+        sessionRestore = await this.replayHistoryFromPreviousSession(previousSessionId);
+      }
+      this.broadcastModeInfo(sessionResult);
+    }
+
+    if (this._systemContext && sessionRestore?.restored !== true) {
+      await this.injectSystemContext();
+    }
+
+    return sessionRestore;
+  }
+
+  /**
+   * Inject system context as the first message
+   */
+  private async injectSystemContext(): Promise<void> {
+    if (!this._systemContext || !this.acpConnection || !this._sessionId) return;
+
+    console.log('[amux] Injecting system context');
+
+    const contextPrompt = `<system-context>\n${this._systemContext}\n</system-context>\n\nI understand. I'm ready to help based on the context provided above.`;
+
+    try {
+      await this.acpConnection.prompt({
+        sessionId: this._sessionId,
+        prompt: [{ type: 'text', text: contextPrompt }],
+      });
+      console.log('[amux] System context injected successfully');
+    } catch (err) {
+      console.error('[amux] Failed to inject system context:', err);
+    }
+  }
+
+  /**
+   * Replay history events from a previous session
+   */
+  private async replayHistoryFromPreviousSession(
+    previousSessionId: string | null
+  ): Promise<SessionRestoreInfo> {
+    if (!previousSessionId) {
+      return { restored: false, reason: 'No previous session to restore.' };
+    }
+
+    const history = this.stateManager.loadHistory(previousSessionId);
+    if (history.length === 0) {
+      return { restored: false, reason: 'No history found for previous session.' };
+    }
+
+    this.emit('history_replay', {
+      previousSessionId,
+      events: history,
+      eventCount: history.length,
+    });
+
+    return {
+      restored: true,
+      reason: `Restored ${history.length} events from previous session.`,
+    };
+  }
+
+  /**
+   * Create a new session
+   */
+  private async createNewSession(): Promise<acp.NewSessionResponse> {
+    if (!this.acpConnection) {
+      throw new Error('No ACP connection');
+    }
+
+    const sessionResult = await withTimeout(
+      this.acpConnection.newSession({ cwd: this._cwd, mcpServers: [] }),
+      INIT_TIMEOUT_MS,
+      'New session creation'
+    );
+    this._sessionId = sessionResult.sessionId;
+    this.saveState();
+
+    if (this._sessionId) {
+      this.stateManager.registerSession(this._sessionId, this._cwd, this._agentType);
+    }
+
+    return sessionResult;
+  }
+
+  /**
+   * Broadcast mode info if available
+   */
+  private broadcastModeInfo(sessionResult: acp.NewSessionResponse): void {
+    if (sessionResult?.modes) {
+      this.emit('update', {
+        sessionUpdate: 'current_mode_update',
+        currentModeId: sessionResult.modes.currentModeId,
+      });
+    }
+  }
+
+  /**
+   * Send a prompt to the agent
+   */
+  async prompt(message: string): Promise<{ stopReason?: string }> {
+    if (!this.acpConnection || !this._sessionId) {
+      throw new Error('Agent not ready');
+    }
+
+    console.log('[amux] Sending prompt:', message.slice(0, 100));
+
+    this.emitTurnMarker('turn_start');
+    this.stateManager.incrementSessionMessageCount(this._sessionId);
+    this.maybeNameSession(message);
+
+    // Emit and store user message
+    this.emitUpdate( {
+      sessionUpdate: 'user_message_chunk',
+      content: { type: 'text', text: message },
+    });
+
+    try {
+      // Parse @mentions into resource_link blocks
+      const content = parseMessageToContentBlocks(message, this._cwd);
+      const result = await this.acpConnection.prompt({
+        sessionId: this._sessionId,
+        prompt: content,
+      });
+      this.emit('prompt_complete', { stopReason: result.stopReason });
+      return { stopReason: result.stopReason };
+    } catch (err) {
+      console.error('[amux] Prompt error:', err);
+      this.emit('error', { message: `Prompt failed: ${(err as Error).message}` });
+      throw err;
+    } finally {
+      this.emitTurnMarker('turn_end');
+    }
+  }
+
+  /**
+   * Cancel current operation
+   */
+  async cancel(): Promise<void> {
+    if (this.acpConnection && this._sessionId) {
+      await this.acpConnection.cancel({ sessionId: this._sessionId });
+    }
+  }
+
+  /**
+   * Handle permission response from UI
+   */
+  respondToPermission(requestId: string, optionId: string): void {
+    this.amuxClient?.handlePermissionResponse(requestId, optionId);
+  }
+
+  /**
+   * List available sessions
+   */
+  listSessions(): SessionMetadata[] {
+    const sessions = this.stateManager.listSessions(this._cwd);
+    return sessions.map(s => ({
+      ...s,
+      isCurrent: s.id === this._sessionId,
+    }));
+  }
+
+  /**
+   * Load history for current session
+   */
+  loadHistory(): unknown[] {
+    return this.stateManager.loadHistory(this._sessionId);
+  }
+
+  /**
+   * Switch to a different session
+   */
+  async switchSession(sessionId: string): Promise<void> {
+    if (!this.acpConnection) {
+      throw new Error('Agent not ready');
+    }
+
+    const canLoadSession = this._agentCapabilities?.loadSession === true;
+    const canResume = this._agentCapabilities?.sessionCapabilities?.resume !== undefined;
+
+    if (!canLoadSession && !canResume) {
+      throw new Error('Agent does not support session switching');
+    }
+
+    try {
+      if (canLoadSession) {
+        await withTimeout(
+          this.acpConnection.loadSession({
+            sessionId,
+            cwd: this._cwd,
+            mcpServers: [],
+          }),
+          INIT_TIMEOUT_MS,
+          'Session load'
+        );
+      } else {
+        await withTimeout(
+          this.acpConnection.unstable_resumeSession({
+            sessionId,
+            cwd: this._cwd,
+            mcpServers: [],
+          }),
+          INIT_TIMEOUT_MS,
+          'Session resume'
+        );
+      }
+
+      this._sessionId = sessionId;
+      this.saveState();
+
+      this.emit('session_switched', { sessionId: this._sessionId });
+
+      const history = this.stateManager.loadHistory(sessionId);
+      if (history.length > 0) {
+        this.emit('history_replay', {
+          previousSessionId: sessionId,
+          events: history,
+          eventCount: history.length,
+        });
+      }
+
+    } catch (err) {
+      console.error(`[amux] Failed to switch session:`, err);
+      throw err;
+    }
+  }
+
+  /**
+   * Change working directory (requires agent restart)
+   */
+  async changeCwd(newPath: string): Promise<void> {
+    if (this.fixedCwd) {
+      throw new Error('Working directory is fixed for this session');
+    }
+
+    this._cwd = newPath;
+    this._sessionId = null;
+    this.saveState();
+
+    await this.killAgent();
+    await this.spawnAgent();
+  }
+
+  /**
+   * Create a new session
+   */
+  async newSession(): Promise<void> {
+    if (!this.acpConnection) {
+      throw new Error('No ACP connection');
+    }
+
+    const sessionResult = await this.acpConnection.newSession({
+      cwd: this._cwd,
+      mcpServers: [],
+    });
+    this._sessionId = sessionResult.sessionId;
+    this.saveState();
+
+    if (this._sessionId) {
+      this.stateManager.registerSession(this._sessionId, this._cwd, this._agentType);
+    }
+
+    this.emit('session_created', { sessionId: this._sessionId });
+    this.broadcastModeInfo(sessionResult);
+
+    if (this._systemContext) {
+      await this.injectSystemContext();
+    }
+  }
+
+  /**
+   * Set session mode
+   */
+  async setMode(modeId: string): Promise<void> {
+    if (!this.acpConnection || !this._sessionId) {
+      throw new Error('Agent not ready');
+    }
+
+    await this.acpConnection.setSessionMode({
+      sessionId: this._sessionId,
+      modeId,
+    });
+  }
+
+  /**
+   * Set session model
+   */
+  async setModel(modelId: string): Promise<void> {
+    if (!this.acpConnection || !this._sessionId) {
+      throw new Error('Agent not ready');
+    }
+
+    await this.acpConnection.unstable_setSessionModel({
+      sessionId: this._sessionId,
+      modelId,
+    });
+  }
+
+  /**
+   * Change agent type (requires agent restart)
+   */
+  async changeAgent(agentType: string): Promise<void> {
+    if (!this.agents[agentType]) {
+      throw new Error(`Unknown agent type: ${agentType}`);
+    }
+
+    if (agentType !== this._agentType) {
+      this._agentType = agentType;
+      this._sessionId = null;
+      this.saveState();
+
+      await this.killAgent();
+      await this.spawnAgent();
+    }
+  }
+
+  /**
+   * Kill the agent process
+   */
+  async killAgent(): Promise<void> {
+    if (this.agentProcess) {
+      this.agentProcess.kill();
+      this.agentProcess = null;
+      this.acpConnection = null;
+    }
+  }
+
+  /**
+   * Shutdown (cleanup all resources)
+   */
+  shutdown(): void {
+    this.terminalManager.killAll();
+    this.agentProcess?.kill();
+  }
+}