stack/packages/mosaic/framework/runtime/pi/mosaic-extension.ts

/**
 * mosaic-extension.ts — Pi Extension for Mosaic Framework
 *
 * Integrates the Mosaic agent framework into Pi sessions launched via `mosaic pi`.
 * Handles:
 *   1. Session start — run repo hooks, detect active mission, display status
 *   2. Session end — run repo hooks, clean up session lock
 *   3. Mission context — inject active mission state into conversation
 *   4. Memory routing — remind agent to use ~/.config/mosaic/memory/
 */

import type { ExtensionAPI } from '@mariozechner/pi-coding-agent';
import { existsSync, readFileSync, writeFileSync, unlinkSync, mkdirSync } from 'node:fs';
import { join, basename } from 'node:path';
import { homedir } from 'node:os';
import { execSync, spawnSync } from 'node:child_process';

// ---------------------------------------------------------------------------
// Config
// ---------------------------------------------------------------------------

const MOSAIC_HOME = process.env['MOSAIC_HOME'] ?? join(homedir(), '.config', 'mosaic');

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

function safeRead(filePath: string): string | null {
  try {
    return readFileSync(filePath, 'utf-8');
  } catch {
    return null;
  }
}

function safeJsonRead(filePath: string): Record<string, unknown> | null {
  const raw = safeRead(filePath);
  if (!raw) return null;
  try {
    return JSON.parse(raw) as Record<string, unknown>;
  } catch {
    return null;
  }
}

function nowIso(): string {
  return new Date().toISOString().replace(/\.\d{3}Z$/, 'Z');
}

// ---------------------------------------------------------------------------
// Mission detection
// ---------------------------------------------------------------------------

interface ActiveMission {
  name: string;
  id: string;
  status: string;
  milestonesTotal: number;
  milestonesCompleted: number;
}

function detectMission(cwd: string): ActiveMission | null {
  const missionFile = join(cwd, '.mosaic', 'orchestrator', 'mission.json');
  const data = safeJsonRead(missionFile);
  if (!data) return null;

  const status = String(data.status ?? 'inactive');
  if (status !== 'active' && status !== 'paused') return null;

  const milestones = Array.isArray(data.milestones) ? data.milestones : [];
  const completed = milestones.filter(
    (m: unknown) =>
      typeof m === 'object' && m !== null && (m as Record<string, unknown>).status === 'completed',
  ).length;

  return {
    name: String(data.name ?? 'unnamed'),
    id: String(data.mission_id ?? ''),
    status,
    milestonesTotal: milestones.length,
    milestonesCompleted: completed,
  };
}

// ---------------------------------------------------------------------------
// Session lock management
// ---------------------------------------------------------------------------

function sessionLockPath(cwd: string): string {
  return join(cwd, '.mosaic', 'orchestrator', 'session.lock');
}

function writeSessionLock(cwd: string): void {
  const lockDir = join(cwd, '.mosaic', 'orchestrator');
  if (!existsSync(lockDir)) return; // Only write lock if orchestrator dir exists

  const lock = {
    session_id: `pi-${new Date().toISOString().replace(/[:.]/g, '-')}-${process.pid}`,
    runtime: 'pi',
    pid: process.pid,
    started_at: nowIso(),
    project_path: cwd,
    milestone_id: '',
  };

  try {
    writeFileSync(sessionLockPath(cwd), JSON.stringify(lock, null, 2) + '\n', 'utf-8');
  } catch {
    // Non-fatal — orchestrator dir may not be writable
  }
}

function cleanSessionLock(cwd: string): void {
  try {
    const lockFile = sessionLockPath(cwd);
    if (existsSync(lockFile)) {
      unlinkSync(lockFile);
    }
  } catch {
    // Non-fatal
  }
}

// ---------------------------------------------------------------------------
// Repo hooks
// ---------------------------------------------------------------------------

function runRepoHook(cwd: string, hookName: string): void {
  const script = join(cwd, 'scripts', 'agent', `${hookName}.sh`);
  if (!existsSync(script)) return;

  try {
    spawnSync('bash', [script], {
      cwd,
      stdio: 'pipe',
      timeout: 30_000,
      env: { ...process.env, MOSAIC_RUNTIME: 'pi' },
    });
  } catch {
    // Non-fatal
  }
}

// ---------------------------------------------------------------------------
// Build mission summary for notifications
// ---------------------------------------------------------------------------

function buildMissionSummary(cwd: string, mission: ActiveMission): string {
  const lines: string[] = [
    `Mission: ${mission.name}`,
    `Status: ${mission.status} | Milestones: ${mission.milestonesCompleted}/${mission.milestonesTotal}`,
  ];

  // Task counts
  const tasksFile = join(cwd, 'docs', 'TASKS.md');
  const tasksContent = safeRead(tasksFile);
  if (tasksContent) {
    const tableRows = tasksContent
      .split('\n')
      .filter((l) => l.startsWith('|') && !l.includes('---'));
    const total = Math.max(0, tableRows.length - 1); // minus header
    const done = (tasksContent.match(/\|\s*done\s*\|/gi) ?? []).length;
    lines.push(`Tasks: ${done} done / ${total} total`);
  }

  // Latest scratchpad
  try {
    const spDir = join(cwd, 'docs', 'scratchpads');
    if (existsSync(spDir)) {
      const files = execSync(`ls -t "${spDir}"/*.md 2>/dev/null | head -1`, {
        encoding: 'utf-8',
        timeout: 5000,
      }).trim();
      if (files) lines.push(`Scratchpad: ${basename(files)}`);
    }
  } catch {
    // Non-fatal
  }

  lines.push('', 'Read ORCHESTRATOR-PROTOCOL.md + TASKS.md before proceeding.');
  return lines.join('\n');
}

// ---------------------------------------------------------------------------
// Extension registration
// ---------------------------------------------------------------------------

export default function register(pi: ExtensionAPI) {
  let sessionCwd = process.cwd();

  // ── Session Start ─────────────────────────────────────────────────────
  pi.on('session_start', async (_event, ctx) => {
    sessionCwd = process.cwd();

    // Run repo session-start hook
    runRepoHook(sessionCwd, 'session-start');

    // Detect active mission
    const mission = detectMission(sessionCwd);

    if (mission) {
      // Write session lock for orchestrator awareness
      writeSessionLock(sessionCwd);

      const summary = buildMissionSummary(sessionCwd, mission);
      ctx.ui.notify(`🎯 Active Mosaic Mission\n${summary}`, 'info');
    } else {
      ctx.ui.notify('Mosaic framework loaded', 'info');
    }
  });

  // ── Session End ───────────────────────────────────────────────────────
  pi.on('session_end', async (_event, _ctx) => {
    // Run repo session-end hook
    runRepoHook(sessionCwd, 'session-end');

    // Clean up session lock
    cleanSessionLock(sessionCwd);
  });

  // ── Register /mosaic-status command ───────────────────────────────────
  pi.registerCommand('mosaic-status', {
    description: 'Show Mosaic mission status for the current project',
    handler: async (_args, ctx) => {
      const mission = detectMission(sessionCwd);
      if (!mission) {
        ctx.ui.notify('No active Mosaic mission in this project.', 'info');
        return;
      }
      const summary = buildMissionSummary(sessionCwd, mission);
      ctx.ui.notify(`🎯 Mission Status\n${summary}`, 'info');
    },
  });

  // ── Register /mosaic-memory command ───────────────────────────────────
  pi.registerCommand('mosaic-memory', {
    description: 'Show Mosaic memory directory path and contents',
    handler: async (_args, ctx) => {
      const memDir = join(MOSAIC_HOME, 'memory');
      if (!existsSync(memDir)) {
        ctx.ui.notify(`Memory directory: ${memDir} (empty)`, 'info');
        return;
      }
      try {
        const files = execSync(`ls -la "${memDir}" 2>/dev/null`, {
          encoding: 'utf-8',
          timeout: 5000,
        }).trim();
        ctx.ui.notify(`Memory directory: ${memDir}\n${files}`, 'info');
      } catch {
        ctx.ui.notify(`Memory directory: ${memDir}`, 'info');
      }
    },
  });
}