Codewalkers/apps/server/agent/cleanup-manager.ts

/**
 * CleanupManager — Worktree, branch, and log cleanup for agents.
 *
 * Extracted from MultiProviderAgentManager. Handles all filesystem
 * and git cleanup operations, plus orphan detection and reconciliation.
 */

import { promisify } from 'node:util';
import { execFile } from 'node:child_process';
import { readFile, readdir, rm, cp, mkdir } from 'node:fs/promises';
import { existsSync, readdirSync } from 'node:fs';
import { join } from 'node:path';
import type { AgentRepository } from '../db/repositories/agent-repository.js';
import type { ProjectRepository } from '../db/repositories/project-repository.js';
import type { EventBus, AgentCrashedEvent } from '../events/index.js';
import { createModuleLogger } from '../logger/index.js';
import { SimpleGitWorktreeManager } from '../git/manager.js';
import { getProjectCloneDir } from '../git/project-clones.js';
import { getStreamParser } from './providers/parsers/index.js';
import { FileTailer } from './file-tailer.js';
import { getProvider } from './providers/registry.js';
import type { StreamEvent } from './providers/parsers/index.js';
import type { SignalManager } from './lifecycle/signal-manager.js';
import { isPidAlive } from './process-manager.js';

const log = createModuleLogger('cleanup-manager');
const execFileAsync = promisify(execFile);

export class CleanupManager {
  constructor(
    private workspaceRoot: string,
    private repository: AgentRepository,
    private projectRepository: ProjectRepository,
    private eventBus?: EventBus,
    private debug: boolean = false,
    private signalManager?: SignalManager,
  ) {}

  /**
   * Resolve the agent's working directory path.
   */
  private getAgentWorkdir(alias: string): string {
    return join(this.workspaceRoot, 'agent-workdirs', alias);
  }

  /**
   * Resolve the actual working directory for an agent, probing for the
   * workspace/ subdirectory used by standalone agents.
   */
  private resolveAgentCwd(worktreeId: string): string {
    const base = this.getAgentWorkdir(worktreeId);

    // Fast path: .cw/output exists at the base level
    if (existsSync(join(base, '.cw', 'output'))) {
      return base;
    }

    // Standalone agents use a workspace/ subdirectory
    const workspaceSub = join(base, 'workspace');
    if (existsSync(join(workspaceSub, '.cw'))) {
      return workspaceSub;
    }

    // Initiative-based agents may have written .cw/ inside a project
    // subdirectory (e.g. agent-workdirs/<name>/codewalk-district/.cw/).
    // Probe immediate children for a .cw/output directory.
    try {
      const entries = readdirSync(base, { withFileTypes: true });
      for (const entry of entries) {
        if (entry.isDirectory() && entry.name !== '.cw') {
          const projectSub = join(base, entry.name);
          if (existsSync(join(projectSub, '.cw', 'output'))) {
            return projectSub;
          }
        }
      }
    } catch {
      // base dir may not exist
    }

    return base;
  }

  /**
   * Remove git worktrees for an agent.
   * Handles both initiative-linked (multi-project) and standalone agents.
   */
  async removeAgentWorktrees(alias: string, initiativeId: string | null): Promise<void> {
    const agentWorkdir = this.getAgentWorkdir(alias);

    try {
      await readdir(agentWorkdir);
    } catch {
      return;
    }

    if (initiativeId) {
      const projects = await this.projectRepository.findProjectsByInitiativeId(initiativeId);
      for (const project of projects) {
        try {
          const clonePath = join(this.workspaceRoot, getProjectCloneDir(project.name, project.id));
          const wm = new SimpleGitWorktreeManager(clonePath, undefined, agentWorkdir);
          await wm.remove(project.name);
        } catch (err) {
          log.warn({ alias, project: project.name, err: err instanceof Error ? err.message : String(err) }, 'failed to remove project worktree');
        }
      }
    } else {
      try {
        const wm = new SimpleGitWorktreeManager(this.workspaceRoot, undefined, agentWorkdir);
        await wm.remove('workspace');
      } catch (err) {
        log.warn({ alias, err: err instanceof Error ? err.message : String(err) }, 'failed to remove standalone worktree');
      }
    }

    await rm(agentWorkdir, { recursive: true, force: true });
    await this.pruneWorktrees(initiativeId);
    // Also prune project clone repos (catches errand worktree refs)
    await this.pruneProjectRepos();
  }

  /**
   * Delete agent/<alias> branches from all relevant repos.
   */
  async removeAgentBranches(alias: string, initiativeId: string | null): Promise<void> {
    const branchName = `agent/${alias}`;
    const repoPaths: string[] = [];

    if (initiativeId) {
      const projects = await this.projectRepository.findProjectsByInitiativeId(initiativeId);
      for (const project of projects) {
        repoPaths.push(join(this.workspaceRoot, getProjectCloneDir(project.name, project.id)));
      }
    } else {
      repoPaths.push(this.workspaceRoot);
    }

    for (const repoPath of repoPaths) {
      try {
        await execFileAsync('git', ['branch', '-D', branchName], { cwd: repoPath });
      } catch {
        // Branch may not exist
      }
    }
  }

  /**
   * Remove log directory for an agent.
   */
  async removeAgentLogs(agentName: string): Promise<void> {
    const logDir = join(this.workspaceRoot, '.cw', 'agent-logs', agentName);
    await rm(logDir, { recursive: true, force: true });
  }

  /**
   * Run git worktree prune on all relevant repos.
   */
  async pruneWorktrees(initiativeId: string | null): Promise<void> {
    const repoPaths: string[] = [];

    if (initiativeId) {
      const projects = await this.projectRepository.findProjectsByInitiativeId(initiativeId);
      for (const project of projects) {
        repoPaths.push(join(this.workspaceRoot, getProjectCloneDir(project.name, project.id)));
      }
    } else {
      repoPaths.push(this.workspaceRoot);
    }

    for (const repoPath of repoPaths) {
      try {
        await execFileAsync('git', ['worktree', 'prune'], { cwd: repoPath });
      } catch (err) {
        log.warn({ repoPath, err: err instanceof Error ? err.message : String(err) }, 'failed to prune worktrees');
      }
    }
  }

  /**
   * Run git worktree prune on all project clone repos.
   * Catches errand worktree refs that aren't covered by initiative-based pruning.
   */
  private async pruneProjectRepos(): Promise<void> {
    const reposPath = join(this.workspaceRoot, 'repos');
    try {
      const repoDirs = await readdir(reposPath);
      for (const repoDir of repoDirs) {
        try {
          await execFileAsync('git', ['worktree', 'prune'], { cwd: join(reposPath, repoDir) });
        } catch { /* ignore */ }
      }
    } catch { /* no repos dir */ }
  }

  /**
   * Clean up orphaned agent workdirs (directories with no matching DB agent).
   */
  async cleanupOrphanedWorkdirs(): Promise<void> {
    const workdirsPath = join(this.workspaceRoot, 'agent-workdirs');
    let entries: string[];
    try {
      entries = await readdir(workdirsPath);
    } catch {
      return;
    }

    const agents = await this.repository.findAll();
    const knownAliases = new Set(agents.map(a => a.name));

    for (const entry of entries) {
      if (!knownAliases.has(entry)) {
        log.info({ orphan: entry }, 'removing orphaned agent workdir');
        try {
          await rm(join(workdirsPath, entry), { recursive: true, force: true });
        } catch (err) {
          log.warn({ orphan: entry, err: err instanceof Error ? err.message : String(err) }, 'failed to remove orphaned workdir');
        }
      }
    }

    try {
      await execFileAsync('git', ['worktree', 'prune'], { cwd: this.workspaceRoot });
    } catch { /* ignore */ }

    const reposPath = join(this.workspaceRoot, 'repos');
    try {
      const repoDirs = await readdir(reposPath);
      for (const repoDir of repoDirs) {
        try {
          await execFileAsync('git', ['worktree', 'prune'], { cwd: join(reposPath, repoDir) });
        } catch { /* ignore */ }
      }
    } catch { /* no repos dir */ }
  }

  /**
   * Clean up orphaned agent log directories (directories with no matching DB agent).
   */
  async cleanupOrphanedLogs(): Promise<void> {
    const logsPath = join(this.workspaceRoot, '.cw', 'agent-logs');
    let entries: string[];
    try {
      entries = await readdir(logsPath);
    } catch {
      return;
    }

    const agents = await this.repository.findAll();
    const knownNames = new Set(agents.map(a => a.name));

    for (const entry of entries) {
      if (!knownNames.has(entry)) {
        log.info({ orphan: entry }, 'removing orphaned agent log dir');
        try {
          await rm(join(logsPath, entry), { recursive: true, force: true });
        } catch (err) {
          log.warn({ orphan: entry, err: err instanceof Error ? err.message : String(err) }, 'failed to remove orphaned log dir');
        }
      }
    }
  }

  /**
   * Get the relative subdirectory names of dirty worktrees for an agent.
   * Returns an empty array if all worktrees are clean or the workdir doesn't exist.
   */
  async getDirtyWorktreePaths(alias: string, initiativeId: string | null): Promise<{ name: string; absPath: string }[]> {
    const agentWorkdir = this.getAgentWorkdir(alias);

    try {
      await readdir(agentWorkdir);
    } catch {
      return [];
    }

    const worktreePaths: { absPath: string; name: string }[] = [];

    if (initiativeId) {
      const projects = await this.projectRepository.findProjectsByInitiativeId(initiativeId);
      for (const project of projects) {
        worktreePaths.push({ absPath: join(agentWorkdir, project.name), name: project.name });
      }
    } else {
      worktreePaths.push({ absPath: join(agentWorkdir, 'workspace'), name: 'workspace' });
    }

    const dirty: { name: string; absPath: string }[] = [];
    for (const { absPath, name } of worktreePaths) {
      try {
        const { stdout } = await execFileAsync('git', ['status', '--porcelain'], { cwd: absPath });
        if (stdout.trim().length > 0) dirty.push({ name, absPath });
      } catch {
        dirty.push({ name, absPath });
      }
    }
    return dirty;
  }

  /**
   * Check if all project worktrees for an agent are clean (no uncommitted/untracked files).
   */
  async isWorkdirClean(alias: string, initiativeId: string | null): Promise<boolean> {
    const dirty = await this.getDirtyWorktreePaths(alias, initiativeId);
    if (dirty.length > 0) {
      log.info({ alias, dirtyWorktrees: dirty }, 'workdir has uncommitted changes');
    }
    return dirty.length === 0;
  }

  /**
   * Archive agent workdir and logs to .cw/debug/ before removal.
   */
  async archiveForDebug(alias: string, agentId: string): Promise<void> {
    const agentWorkdir = this.getAgentWorkdir(alias);
    const debugWorkdir = join(this.workspaceRoot, '.cw', 'debug', 'workdirs', alias);
    const logDir = join(this.workspaceRoot, '.cw', 'agent-logs', alias);
    const debugLogDir = join(this.workspaceRoot, '.cw', 'debug', 'agent-logs', alias);

    try {
      if (existsSync(agentWorkdir)) {
        await mkdir(join(this.workspaceRoot, '.cw', 'debug', 'workdirs'), { recursive: true });
        await cp(agentWorkdir, debugWorkdir, { recursive: true });
        log.debug({ alias, debugWorkdir }, 'archived workdir for debug');
      }
    } catch (err) {
      log.warn({ alias, err: err instanceof Error ? err.message : String(err) }, 'failed to archive workdir for debug');
    }

    try {
      if (existsSync(logDir)) {
        await mkdir(join(this.workspaceRoot, '.cw', 'debug', 'agent-logs'), { recursive: true });
        await cp(logDir, debugLogDir, { recursive: true });
        log.debug({ agentId, debugLogDir }, 'archived logs for debug');
      }
    } catch (err) {
      log.warn({ agentId, err: err instanceof Error ? err.message : String(err) }, 'failed to archive logs for debug');
    }
  }

  /**
   * Auto-cleanup agent workdir after successful completion.
   * Removes worktrees, branches, and logs. Preserves DB record.
   */
  async autoCleanupAfterCompletion(
    agentId: string,
    alias: string,
    initiativeId: string | null,
  ): Promise<{ clean: boolean; removed: boolean }> {
    const agentWorkdir = this.getAgentWorkdir(alias);

    // Idempotent: if workdir is already gone, nothing to do
    if (!existsSync(agentWorkdir)) {
      return { clean: true, removed: true };
    }

    const clean = await this.isWorkdirClean(alias, initiativeId);
    if (!clean) {
      return { clean: false, removed: false };
    }

    if (this.debug) {
      await this.archiveForDebug(alias, agentId);
    }

    let worktreeRemoved = true;
    try {
      await this.removeAgentWorktrees(alias, initiativeId);
    } catch (err) {
      log.warn({ agentId, alias, err: err instanceof Error ? err.message : String(err) }, 'auto-cleanup: failed to remove worktrees');
      worktreeRemoved = false;
    }

    try {
      await this.removeAgentBranches(alias, initiativeId);
    } catch (err) {
      log.warn({ agentId, alias, err: err instanceof Error ? err.message : String(err) }, 'auto-cleanup: failed to remove branches');
    }

    try {
      await this.removeAgentLogs(alias);
    } catch (err) {
      log.warn({ agentId, err: err instanceof Error ? err.message : String(err) }, 'auto-cleanup: failed to remove logs');
    }

    log.info({ agentId, alias }, 'auto-cleanup: workdir and logs removed');
    return { clean: true, removed: worktreeRemoved };
  }

  /**
   * Reconcile agent state after server restart.
   * Checks all agents in 'running' status:
   * - If PID is still alive: create FileTailer to resume streaming
   * - If PID is dead but output file exists: process the output
   * - Otherwise: mark as crashed
   *
   * @param activeAgents - Shared map from manager to register live agents
   * @param onStreamEvent - Callback for stream events from tailer
   * @param onAgentOutput - Callback to process raw agent output
   * @param pollForCompletion - Callback to start polling for completion
   */
  async reconcileAfterRestart(
    activeAgents: Map<string, {
      agentId: string;
      pid: number;
      tailer: FileTailer;
      outputFilePath: string;
      agentCwd?: string;
    }>,
    onStreamEvent: (agentId: string, event: StreamEvent) => void,
    onAgentOutput: (agentId: string, rawOutput: string, provider: NonNullable<ReturnType<typeof getProvider>>) => Promise<void>,
    pollForCompletion: (agentId: string, pid: number) => void,
    onRawContent?: (agentId: string, agentName: string, content: string) => void,
  ): Promise<void> {
    const runningAgents = await this.repository.findByStatus('running');
    log.info({ runningCount: runningAgents.length }, 'reconciling agents after restart');

    for (const agent of runningAgents) {
      const alive = agent.pid ? isPidAlive(agent.pid) : false;
      log.info({ agentId: agent.id, pid: agent.pid, alive }, 'reconcile: checking agent');

      if (alive && agent.outputFilePath) {
        log.debug({ agentId: agent.id, pid: agent.pid }, 'reconcile: resuming streaming for alive agent');

        const parser = getStreamParser(agent.provider);
        const tailer = new FileTailer({
          filePath: agent.outputFilePath,
          agentId: agent.id,
          parser,
          onEvent: (event) => onStreamEvent(agent.id, event),
          startFromBeginning: false,
          onRawContent: onRawContent
            ? (content) => onRawContent(agent.id, agent.name, content)
            : undefined,
        });

        tailer.start().catch((err) => {
          log.warn({ agentId: agent.id, err: err instanceof Error ? err.message : String(err) }, 'failed to start tailer during reconcile');
        });

        const pid = agent.pid!;

        // Resolve actual agent cwd — standalone agents run in workspace/ subdir
        const resolvedCwd = this.resolveAgentCwd(agent.worktreeId);

        activeAgents.set(agent.id, {
          agentId: agent.id,
          pid,
          tailer,
          outputFilePath: agent.outputFilePath,
          agentCwd: resolvedCwd,
        });

        pollForCompletion(agent.id, pid);
      } else if (agent.outputFilePath) {
        // CRITICAL FIX: Check for signal.json completion FIRST before parsing raw output
        // Resolve actual agent cwd — standalone agents run in workspace/ subdir
        const agentWorkdir = this.resolveAgentCwd(agent.worktreeId);
        const hasValidSignal = this.signalManager ? await this.signalManager.readSignal(agentWorkdir) : null;

        if (hasValidSignal) {
          log.debug({ agentId: agent.id }, 'found valid signal.json, processing as completion');
          try {
            const signalFile = join(agentWorkdir, '.cw/output/signal.json');
            const signalContent = await readFile(signalFile, 'utf-8');
            const provider = getProvider(agent.provider);
            if (provider) {
              await onAgentOutput(agent.id, signalContent, provider);
              continue;
            }
          } catch (err) {
            log.error({
              agentId: agent.id,
              err: err instanceof Error ? err.message : String(err)
            }, 'reconcile: failed to process signal.json');
            // Fall through to raw output processing
          }
        }

        try {
          const rawOutput = await readFile(agent.outputFilePath, 'utf-8');
          if (rawOutput.trim()) {
            const provider = getProvider(agent.provider);
            if (provider) {
              // Check if agent actually completed successfully before processing
              const hasCompletionResult = this.checkForCompletionResult(rawOutput);
              if (hasCompletionResult) {
                log.info({ agentId: agent.id }, 'reconcile: processing completed agent output');
                try {
                  await onAgentOutput(agent.id, rawOutput, provider);
                  continue;
                } catch (err) {
                  log.error({
                    agentId: agent.id,
                    err: err instanceof Error ? err.message : String(err)
                  }, 'reconcile: failed to process completed agent output');
                  // Mark as crashed since processing failed
                  await this.repository.update(agent.id, { status: 'crashed' });
                  this.emitCrashed(agent, `Failed to process output: ${err instanceof Error ? err.message : String(err)}`);
                  continue;
                }
              }
            }
          }
        } catch (readErr) {
          log.warn({
            agentId: agent.id,
            err: readErr instanceof Error ? readErr.message : String(readErr)
          }, 'reconcile: failed to read output file');
        }
        log.warn({ agentId: agent.id }, 'reconcile: marking agent crashed (no valid output)');
        await this.repository.update(agent.id, { status: 'crashed' });
        this.emitCrashed(agent, 'Server restarted, agent output not found or invalid');
      } else {
        log.warn({ agentId: agent.id }, 'reconcile: marking agent crashed');
        await this.repository.update(agent.id, { status: 'crashed' });
        this.emitCrashed(agent, 'Server restarted while agent was running');
      }
    }

    try {
      await this.cleanupOrphanedWorkdirs();
    } catch (err) {
      log.warn({ err: err instanceof Error ? err.message : String(err) }, 'orphaned workdir cleanup failed');
    }
    try {
      await this.cleanupOrphanedLogs();
    } catch (err) {
      log.warn({ err: err instanceof Error ? err.message : String(err) }, 'orphaned log cleanup failed');
    }
  }

  /**
   * Check if the agent output contains a completion result line.
   * This indicates the agent finished successfully, even if processing fails.
   */
  private checkForCompletionResult(rawOutput: string): boolean {
    try {
      const lines = rawOutput.trim().split('\n');
      for (const line of lines) {
        try {
          const parsed = JSON.parse(line);
          // Look for Claude CLI result events with success status
          if (parsed.type === 'result' && parsed.subtype === 'success') {
            return true;
          }
          // Look for other providers' completion indicators
          if (parsed.status === 'done' || parsed.status === 'questions') {
            return true;
          }
        } catch { /* skip non-JSON lines */ }
      }
    } catch { /* invalid output format */ }
    return false;
  }

  /**
   * Emit a crashed event for an agent.
   */
  private emitCrashed(agent: { id: string; name: string; taskId: string | null }, error: string): void {
    if (this.eventBus) {
      const event: AgentCrashedEvent = {
        type: 'agent:crashed',
        timestamp: new Date(),
        payload: {
          agentId: agent.id,
          name: agent.name,
          taskId: agent.taskId ?? '',
          error,
        },
      };
      this.eventBus.emit(event);
    }
  }


}