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, ExtensionContext } from '@earendil-works/pi-coding-agent';
import { Type } from 'typebox';
import {
  existsSync,
  readFileSync,
  writeFileSync,
  unlinkSync,
  mkdirSync,
  renameSync,
} from 'node:fs';
import { join, basename } from 'node:path';
import { homedir } from 'node:os';
import { execSync, spawnSync } from 'node:child_process';
import { registerLeaseLifecycleHooks, type LeaseLifecyclePiApi } from './lease-lifecycle.js';

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

const MOSAIC_HOME = process.env['MOSAIC_HOME'] ?? join(homedir(), '.config', 'mosaic');
const MUTATOR_GATE = join(MOSAIC_HOME, 'tools', 'lease-broker', 'mutator-gate.py');
const LEASE_REVOKER = join(MOSAIC_HOME, 'tools', 'lease-broker', 'revoke-lease.py');
const RECOVERY_COMMAND = join(MOSAIC_HOME, 'tools', 'lease-broker', 'recover-context.py');
const RECEIPT_OBSERVER_CLIENT = join(
  MOSAIC_HOME,
  'tools',
  'lease-broker',
  'receipt-observer-client.py',
);
const RECOVERY_TOOL = 'mosaic_context_recover';

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

// ---------------------------------------------------------------------------
// Native heartbeat (fleet R14/R15)
// ---------------------------------------------------------------------------
// When this agent runs under the Mosaic fleet (MOSAIC_AGENT_NAME set), the
// extension writes its OWN heartbeat in the same .hb contract `fleet ps` reads
// (ts/pid/status[/model]) and touches a `.hb.native` precedence marker so the
// shell sidecar defers. Native HB knows the real turn state (busy/ok), so it is
// more accurate than the pane-PID-only sidecar fallback.
const HB_AGENT_NAME = process.env['MOSAIC_AGENT_NAME'] ?? '';
const HB_RUN_DIR = process.env['MOSAIC_HEARTBEAT_RUN_DIR'] ?? join(MOSAIC_HOME, 'fleet', 'run');
const HB_INTERVAL_MS = (() => {
  const s = Number.parseInt(process.env['MOSAIC_HEARTBEAT_INTERVAL'] ?? '', 10);
  return Number.isFinite(s) && s > 0 ? s * 1000 : 15_000;
})();

function nativeHbEnabled(): boolean {
  return HB_AGENT_NAME.length > 0;
}

function readModelId(ctx: ExtensionContext): string | null {
  const m = ctx.model as unknown as { id?: string; name?: string } | undefined;
  return m?.id ?? m?.name ?? null;
}

function writeNativeHeartbeat(status: 'ok' | 'busy', model: string | null): void {
  if (!nativeHbEnabled()) return;
  try {
    mkdirSync(HB_RUN_DIR, { recursive: true });
    const hb = join(HB_RUN_DIR, `${HB_AGENT_NAME}.hb`);
    const lines = [`ts=${nowIso()}`, `pid=${process.pid}`, `status=${status}`];
    if (model) lines.push(`model=${model}`);
    const tmp = `${hb}.tmp.${process.pid}`;
    writeFileSync(tmp, lines.join('\n') + '\n');
    renameSync(tmp, hb); // atomic replace — fleet ps never reads a partial file
    // Precedence marker: tells the shell sidecar that native HB is authoritative.
    writeFileSync(join(HB_RUN_DIR, `${HB_AGENT_NAME}.hb.native`), nowIso() + '\n');
  } catch {
    // Best-effort: never let heartbeat I/O disrupt the Pi session.
  }
}

function clearNativeMarker(): void {
  if (!nativeHbEnabled()) return;
  try {
    const m = join(HB_RUN_DIR, `${HB_AGENT_NAME}.hb.native`);
    if (existsSync(m)) unlinkSync(m); // native stopping — let the sidecar take over
  } catch {
    /* ignore */
  }
}

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');
}

function runPiLeaseRevoker(args: string[]): boolean {
  const result = spawnSync('python3', [LEASE_REVOKER, ...args], {
    encoding: 'utf8',
    timeout: 2_000,
    env: process.env,
  });
  return result.status === 0;
}

function checkPiMutatorGate(toolName: string): { block: true; reason: string } | undefined {
  const result = spawnSync('python3', [MUTATOR_GATE, '--runtime', 'pi'], {
    input: `${JSON.stringify({ tool_name: toolName })}\n`,
    encoding: 'utf8',
    timeout: 2_000,
    env: process.env,
  });
  if (result.status === 0) return undefined;
  const detail = String(result.stderr ?? '')
    .trim()
    .split('\n')[0];
  return {
    block: true,
    reason: detail || 'BLOCKED: Mosaic mutator gate is unavailable or the lease is UNVERIFIED.',
  };
}

function checkPiRecoveryGate(): { block: true; reason: string } | undefined {
  return checkPiMutatorGate(RECOVERY_TOOL);
}

function assistantMessageText(message: unknown): string | undefined {
  if (typeof message !== 'object' || message === null) return undefined;
  const value = message as { role?: unknown; content?: unknown };
  if (value.role !== 'assistant') return undefined;
  if (typeof value.content === 'string') return value.content;
  if (!Array.isArray(value.content)) return undefined;
  const text: string[] = [];
  for (const part of value.content) {
    if (typeof part !== 'object' || part === null) return undefined;
    const typed = part as { type?: unknown; text?: unknown };
    if (typed.type !== 'text' || typeof typed.text !== 'string') return undefined;
    text.push(typed.text);
  }
  return text.join('');
}

function recordPiMessageEnd(message: unknown): void {
  const latestAssistantMessage = assistantMessageText(message);
  if (latestAssistantMessage === undefined) return;
  // This sends finalized Pi message_end content only to the daemon-owned
  // authenticated observer transport, never to the public broker request API.
  spawnSync('python3', [RECEIPT_OBSERVER_CLIENT, '--runtime', 'pi'], {
    input: `${JSON.stringify({ latest_assistant_message: latestAssistantMessage })}\n`,
    encoding: 'utf8',
    timeout: 2_000,
    env: process.env,
  });
}

function runPiRecoveryCommand(params: {
  phase: 'begin' | 'complete';
  construction?: string;
  compactionEpoch?: number;
  requestEpoch?: number;
}): { content: Array<{ type: 'text'; text: string }> } {
  const args = [RECOVERY_COMMAND, params.phase];
  if (params.phase === 'begin') {
    if (
      typeof params.construction !== 'string' ||
      !Number.isInteger(params.compactionEpoch) ||
      !Number.isInteger(params.requestEpoch) ||
      params.compactionEpoch < 0 ||
      params.requestEpoch < 0
    ) {
      return {
        content: [
          { type: 'text', text: 'Recovery begin requires construction and non-negative epochs.' },
        ],
      };
    }
    args.push(
      '--construction',
      params.construction,
      '--compaction-epoch',
      String(params.compactionEpoch),
      '--request-epoch',
      String(params.requestEpoch),
    );
  }
  const result = spawnSync('python3', args, {
    encoding: 'utf8',
    timeout: 3_000,
    env: process.env,
  });
  const output = result.status === 0 ? String(result.stdout ?? '') : String(result.stderr ?? '');
  return { content: [{ type: 'text', text: output || 'Constrained recovery refused.' }] };
}

// ---------------------------------------------------------------------------
// 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();
  let hbStatus: 'ok' | 'busy' = 'ok';
  let hbModel: string | null = null;
  let hbTimer: ReturnType<typeof setInterval> | null = null;

  // ── Compaction observers and same-PID generation rollover ─────────────
  registerLeaseLifecycleHooks(pi as unknown as LeaseLifecyclePiApi, runPiLeaseRevoker);

  // ── Whole mutator-class authorization gate ────────────────────────────
  // Every Pi tool, including unknown/custom tools, reaches the broker-backed
  // class gate before execution. Broker/script failure blocks fail-closed.
  pi.on('tool_call', async (event) => checkPiMutatorGate(event.toolName));

  // Pi records only a finalized assistant entry at message_end. It never uses
  // after_provider_response, which occurs before stream consumption.
  pi.on('message_end', async (event) => {
    recordPiMessageEnd((event as unknown as { message?: unknown }).message);
  });

  // The recovery custom tool is the only Pi invocation that maps to the
  // broker's exempt RECOVERY_TOOL identity. It is not a Bash exception.
  pi.registerTool({
    name: RECOVERY_TOOL,
    label: 'Mosaic Context Recovery',
    description:
      'Run the constrained broker-backed context recovery flow. This is the sole ungated mutator.',
    parameters: Type.Object({
      phase: Type.Union([Type.Literal('begin'), Type.Literal('complete')]),
      construction: Type.Optional(Type.String()),
      compactionEpoch: Type.Optional(Type.Integer({ minimum: 0 })),
      requestEpoch: Type.Optional(Type.Integer({ minimum: 0 })),
    }),
    async execute(_toolCallId, params) {
      const blocked = checkPiRecoveryGate();
      if (blocked !== undefined) return { content: [{ type: 'text', text: blocked.reason }] };
      return runPiRecoveryCommand(params);
    },
  });

  // ── 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');
    }

    // Native heartbeat: write immediately, then on an interval. Idle = 'ok';
    // turn_start/turn_end flip the status so `fleet ps` reflects real activity.
    if (nativeHbEnabled()) {
      hbModel = readModelId(ctx);
      writeNativeHeartbeat('ok', hbModel);
      hbTimer = setInterval(() => writeNativeHeartbeat(hbStatus, hbModel), HB_INTERVAL_MS);
      if (typeof hbTimer.unref === 'function') hbTimer.unref();
    }
  });

  // ── Turn lifecycle → accurate busy/ok heartbeat ───────────────────────
  pi.on('turn_start', async (_event, ctx) => {
    hbStatus = 'busy';
    hbModel = readModelId(ctx) ?? hbModel;
    writeNativeHeartbeat('busy', hbModel);
  });
  pi.on('turn_end', async (_event, ctx) => {
    hbStatus = 'ok';
    hbModel = readModelId(ctx) ?? hbModel;
    writeNativeHeartbeat('ok', hbModel);
  });

  // ── Session Shutdown ──────────────────────────────────────────────────
  // (The pi API event is 'session_shutdown'; the prior 'session_end' handler
  // never fired — fixed here so repo hooks + lock cleanup actually run.)
  pi.on('session_shutdown', async (_event, _ctx) => {
    if (hbTimer) {
      clearInterval(hbTimer);
      hbTimer = null;
    }
    clearNativeMarker();

    // 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');
      }
    },
  });

  // ── Register mosaic_mission_status tool (model-callable) ──────────────
  // R14 "proper tool usage": give the agent a first-class tool to load its
  // active Mosaic mission, milestone progress, task counts, and latest
  // scratchpad — so it self-orients on in-flight work before planning,
  // instead of shelling out or guessing. Mirrors the /mosaic-status command
  // but returns the summary as tool output the LLM can read.
  pi.registerTool({
    name: 'mosaic_mission_status',
    label: 'Mosaic Mission Status',
    description:
      'Return the active Mosaic mission, milestone progress, task counts, and latest scratchpad for the current project. Returns a note when no mission is active.',
    promptSnippet: 'Read the active Mosaic mission + task state for the current project',
    promptGuidelines: [
      'Use mosaic_mission_status at the start of a session or task to load the active mission, milestone progress, and open tasks before planning work.',
    ],
    parameters: Type.Object({}),
    async execute(_toolCallId, _params, _signal, _onUpdate, _ctx) {
      const mission = detectMission(sessionCwd);
      const text = mission
        ? buildMissionSummary(sessionCwd, mission)
        : 'No active Mosaic mission in this project.';
      return {
        content: [{ type: 'text', text }],
        details: mission ? { ...mission } : { active: false },
      };
    },
  });
}